using System.Collections.Generic;
using UnityEngine;

namespace FishROV.AvoidanceBenchmark
{
    public sealed class AvoidanceScenarioController
    {
        private readonly Dictionary<BenchmarkAgent, Vector3> targets = new Dictionary<BenchmarkAgent, Vector3>();
        private AvoidanceBenchmarkConfig config;
        private float elapsed;

        public void Reset(AvoidanceBenchmarkConfig benchmarkConfig, IReadOnlyList<BenchmarkAgent> agents)
        {
            config = benchmarkConfig.Clone();
            elapsed = 0f;
            targets.Clear();

            foreach (var agent in agents)
            {
                targets[agent] = CreateTarget(agent);
            }
        }

        public Vector3 GetPreferredVelocity(BenchmarkAgent agent, float deltaTime)
        {
            elapsed += deltaTime;

            if (!targets.TryGetValue(agent, out var target))
            {
                target = CreateTarget(agent);
                targets[agent] = target;
            }

            if ((target - agent.transform.position).sqrMagnitude < 4f)
            {
                target = CreateTarget(agent);
                targets[agent] = target;
            }

            var desired = target - agent.transform.position;
            if (config.DimensionMode != AvoidanceDimensionMode.XYZ3D)
            {
                desired.y = 0f;
            }

            if (desired.sqrMagnitude < 0.0001f)
            {
                return Vector3.zero;
            }

            return desired.normalized * agent.MaxSpeed;
        }

        private Vector3 CreateTarget(BenchmarkAgent agent)
        {
            var radius = config.ArenaRadius * 0.85f;

            switch (config.Scenario)
            {
                case AvoidanceScenarioKind.CrossFlow:
                    return agent.GroupId % 2 == 0
                        ? new Vector3(radius, HeightFor(agent), Random.Range(-radius, radius))
                        : new Vector3(-radius, HeightFor(agent), Random.Range(-radius, radius));
                case AvoidanceScenarioKind.NarrowPassage:
                    return new Vector3(agent.GroupId % 2 == 0 ? radius : -radius, HeightFor(agent), Random.Range(-3f, 3f));
                case AvoidanceScenarioKind.DenseCircleStress:
                    var from = agent.transform.position;
                    return new Vector3(-from.x, HeightFor(agent), -from.z);
                case AvoidanceScenarioKind.FreeSwim:
                default:
                    var circle = Random.insideUnitCircle.normalized * Random.Range(radius * 0.25f, radius);
                    return new Vector3(circle.x, HeightFor(agent), circle.y);
            }
        }

        private float HeightFor(BenchmarkAgent agent)
        {
            if (config.DimensionMode == AvoidanceDimensionMode.XZ2D)
            {
                return 0f;
            }

            if (config.DimensionMode == AvoidanceDimensionMode.XZ25D)
            {
                return (agent.GroupId % 5 - 2) * 0.22f;
            }

            return Random.Range(-4f, 4f);
        }
    }
}
