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;
        private float baitDropTime;
        private Vector3 baitPosition;

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

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

        public void Tick(float deltaTime)
        {
            elapsed += deltaTime;
        }

        public void DropBait(Vector3 position)
        {
            baitPosition = position;
            baitDropTime = elapsed;
            if (config != null)
            {
                config.Scenario = AvoidanceScenarioKind.BaitScramble;
            }
        }

        public Vector3 GetPreferredVelocity(BenchmarkAgent agent, float deltaTime)
        {
            if (!targets.TryGetValue(agent, out var target))
            {
                target = CreateTarget(agent);
                targets[agent] = target;
            }

            if (config.Scenario == AvoidanceScenarioKind.BaitScramble)
            {
                return GetBaitScrambleVelocity(agent);
            }

            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.BaitScramble:
                    return baitPosition;
                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);
        }

        private Vector3 GetBaitScrambleVelocity(BenchmarkAgent agent)
        {
            var position = agent.transform.position;
            var flatOffset = new Vector3(position.x - baitPosition.x, 0f, position.z - baitPosition.z);
            if (flatOffset.sqrMagnitude < 0.0001f)
            {
                flatOffset = Quaternion.Euler(0f, agent.GroupId * 137.5f, 0f) * Vector3.forward;
            }

            var outward = flatOffset.normalized;
            var tangentSign = agent.GroupId % 2 == 0 ? 1f : -1f;
            var tangent = new Vector3(-outward.z, 0f, outward.x) * tangentSign;
            var distance = flatOffset.magnitude;
            var age = elapsed - baitDropTime;

            Vector3 desired;
            if (age < 1.2f)
            {
                desired = outward * agent.MaxSpeed * 1.15f + tangent * agent.MaxSpeed * 0.25f;
            }
            else
            {
                var ringRadius = Mathf.Lerp(5.5f, 2.2f, Mathf.Clamp01((age - 1.2f) / 7.5f));
                var radialError = distance - ringRadius;
                var radial = -outward * Mathf.Clamp(radialError * 0.55f, -0.75f, 1f);
                var swirl = tangent * Mathf.Lerp(0.35f, 0.95f, Mathf.Clamp01((age - 1.2f) / 5f));
                var baitPull = -outward * Mathf.Lerp(0.25f, 0.55f, Mathf.Clamp01((age - 1.2f) / 4f));
                desired = (radial + swirl + baitPull).normalized * agent.MaxSpeed;
            }

            if (config.DimensionMode == AvoidanceDimensionMode.XYZ3D)
            {
                desired.y = Mathf.Clamp(baitPosition.y - position.y, -1f, 1f) * agent.MaxSpeed * 0.35f;
            }
            else
            {
                desired.y = 0f;
            }

            return desired;
        }
    }
}
