using System.Collections.Generic;
using UnityEngine;

namespace FishROV.AvoidanceBenchmark
{
    public sealed class AvoidanceCollisionProbe
    {
        private const float RefreshInterval = 0.25f;

        private readonly Dictionary<Vector2Int, List<BenchmarkAgent>> buckets =
            new Dictionary<Vector2Int, List<BenchmarkAgent>>(1024);

        private float refreshTimer;

        public int HorizontalOverlapPairs { get; private set; }
        public int SpatialOverlapPairs { get; private set; }
        public float MinimumClearance { get; private set; }

        public void Update(IReadOnlyList<BenchmarkAgent> agents, AvoidanceDimensionMode dimensionMode, float deltaTime)
        {
            refreshTimer -= deltaTime;
            if (refreshTimer > 0f)
            {
                return;
            }

            refreshTimer = RefreshInterval;
            Recalculate(agents, dimensionMode);
        }

        private void Recalculate(IReadOnlyList<BenchmarkAgent> agents, AvoidanceDimensionMode dimensionMode)
        {
            HorizontalOverlapPairs = 0;
            SpatialOverlapPairs = 0;
            MinimumClearance = float.PositiveInfinity;

            buckets.Clear();
            var cellSize = CalculateCellSize(agents);

            for (var i = 0; i < agents.Count; i++)
            {
                var agent = agents[i];
                if (agent == null)
                {
                    continue;
                }

                var cell = ToCell(agent.transform.position, cellSize);
                if (!buckets.TryGetValue(cell, out var bucket))
                {
                    bucket = new List<BenchmarkAgent>(8);
                    buckets[cell] = bucket;
                }

                bucket.Add(agent);
            }

            for (var i = 0; i < agents.Count; i++)
            {
                var agent = agents[i];
                if (agent == null)
                {
                    continue;
                }

                var center = ToCell(agent.transform.position, cellSize);
                for (var x = -1; x <= 1; x++)
                {
                    for (var y = -1; y <= 1; y++)
                    {
                        var key = new Vector2Int(center.x + x, center.y + y);
                        if (!buckets.TryGetValue(key, out var bucket))
                        {
                            continue;
                        }

                        CountPairs(agent, bucket, dimensionMode);
                    }
                }
            }

            if (float.IsPositiveInfinity(MinimumClearance))
            {
                MinimumClearance = 0f;
            }
        }

        private void CountPairs(BenchmarkAgent agent, List<BenchmarkAgent> bucket, AvoidanceDimensionMode dimensionMode)
        {
            for (var i = 0; i < bucket.Count; i++)
            {
                var other = bucket[i];
                if (other == null || other.GetInstanceID() <= agent.GetInstanceID())
                {
                    continue;
                }

                var combinedRadius = agent.Radius + other.Radius;
                var delta = other.transform.position - agent.transform.position;
                var flatDistance = new Vector2(delta.x, delta.z).magnitude;
                var flatClearance = flatDistance - combinedRadius;
                MinimumClearance = Mathf.Min(MinimumClearance, flatClearance);

                if (flatClearance < 0f)
                {
                    HorizontalOverlapPairs++;
                }

                var spatialDistance = dimensionMode == AvoidanceDimensionMode.XYZ3D
                    ? delta.magnitude
                    : new Vector3(delta.x, delta.y, delta.z).magnitude;
                if (spatialDistance < combinedRadius)
                {
                    SpatialOverlapPairs++;
                }
            }
        }

        private static float CalculateCellSize(IReadOnlyList<BenchmarkAgent> agents)
        {
            var maxRadius = 0.5f;
            for (var i = 0; i < agents.Count; i++)
            {
                var agent = agents[i];
                if (agent != null)
                {
                    maxRadius = Mathf.Max(maxRadius, agent.Radius);
                }
            }

            return Mathf.Max(0.5f, maxRadius * 2.2f);
        }

        private static Vector2Int ToCell(Vector3 position, float cellSize)
        {
            return new Vector2Int(
                Mathf.FloorToInt(position.x / cellSize),
                Mathf.FloorToInt(position.z / cellSize));
        }
    }
}
