using RobotNet10.Common.Models;

namespace RobotNet10.Common;

/// <summary>
/// KD-Tree implementation for efficient 2D spatial search of nodes.
/// Thread-safe for read operations after construction.
/// Time complexity: O(n log n) build, O(log n) search average case.
/// </summary>
public class KDTree
{
    private readonly KDTreeNode? _root;
    private readonly List<KDTreeData> _nodes;

    /// <summary>
    /// Initializes a new KD-Tree from a collection of nodes.
    /// </summary>
    /// <param name="nodes">The nodes to index. Original collection is not modified.</param>
    /// <exception cref="ArgumentNullException">Thrown when nodes is null.</exception>
    public KDTree(IEnumerable<KDTreeData> nodes)
    {
        ArgumentNullException.ThrowIfNull(nodes);

        // Create a copy to avoid mutating input
        _nodes = [.. nodes];

        if (_nodes.Count == 0)
        {
            _root = null;
            return;
        }

        _root = BuildTree(0, _nodes.Count - 1, 0);
    }

    /// <summary>
    /// Gets the total number of nodes in the tree.
    /// </summary>
    public int Count => _nodes.Count;

    /// <summary>
    /// Builds the KD-Tree using index-based recursion to avoid memory allocation overhead.
    /// Time complexity: O(n log n)
    /// </summary>
    private KDTreeNode? BuildTree(int start, int end, int depth)
    {
        if (start > end)
            return null;

        int axis = depth % 2;

        // Use QuickSelect to find median without full sort
        int medianIndex = QuickSelect(start, end, (start + end) / 2, axis);

        return new KDTreeNode(
            node: _nodes[medianIndex],
            axis: axis,
            left: BuildTree(start, medianIndex - 1, depth + 1),
            right: BuildTree(medianIndex + 1, end, depth + 1)
        );
    }

    /// <summary>
    /// QuickSelect algorithm to find the k-th smallest element.
    /// Time complexity: O(n) average, O(n²) worst case.
    /// </summary>
    private int QuickSelect(int left, int right, int k, int axis)
    {
        while (left < right)
        {
            int pivotIndex = Partition(left, right, axis);

            if (pivotIndex == k)
                return k;
            else if (k < pivotIndex)
                right = pivotIndex - 1;
            else
                left = pivotIndex + 1;
        }
        return left;
    }

    /// <summary>
    /// Partitions the array for QuickSelect using median-of-three pivot selection.
    /// This is the CORRECTED version that handles all edge cases properly.
    /// </summary>
    private int Partition(int left, int right, int axis)
    {
        // Handle small subarrays
        if (right - left < 2)
        {
            if (right > left && CompareNodes(_nodes[right], _nodes[left], axis) < 0)
                Swap(left, right);
            return left;
        }

        // Use median-of-three pivot selection for better performance
        int mid = left + (right - left) / 2; // Overflow-safe

        // Sort left, mid, right to get median as pivot
        if (CompareNodes(_nodes[mid], _nodes[left], axis) < 0)
            Swap(left, mid);
        if (CompareNodes(_nodes[right], _nodes[left], axis) < 0)
            Swap(left, right);
        if (CompareNodes(_nodes[right], _nodes[mid], axis) < 0)
            Swap(mid, right);

        // Now: nodes[left] <= nodes[mid] <= nodes[right]
        // Use mid as pivot and hide it at right-1
        KDTreeData pivot = _nodes[mid];
        Swap(mid, right - 1);

        // Partition with pivot at right-1
        int i = left + 1;      // Start after left (which is already <= pivot)
        int j = right - 2;     // Start before pivot position

        while (i <= j)
        {
            // Find element >= pivot from left
            while (i <= j && CompareNodes(_nodes[i], pivot, axis) < 0)
                i++;

            // Find element <= pivot from right
            while (i <= j && CompareNodes(_nodes[j], pivot, axis) > 0)
                j--;

            if (i < j)
            {
                Swap(i, j);
                i++;
                j--;
            }
            else
            {
                break;
            }
        }

        // Put pivot in its final position
        Swap(i, right - 1);
        return i;
    }

    /// <summary>
    /// Compares two nodes along the specified axis.
    /// </summary>
    private static int CompareNodes(KDTreeData a, KDTreeData b, int axis)
    {
        return axis == 0 ? a.X.CompareTo(b.X) : a.Y.CompareTo(b.Y);
    }

    /// <summary>
    /// Swaps two elements in the nodes list using tuple deconstruction.
    /// </summary>
    private void Swap(int i, int j)
    {
        if (i != j)
        {
            (_nodes[j], _nodes[i]) = (_nodes[i], _nodes[j]);
        }
    }

    /// <summary>
    /// Finds the nearest node to the given coordinates within the specified distance limit.
    /// Time complexity: O(log n) average, O(n) worst case.
    /// </summary>
    /// <param name="x">X coordinate</param>
    /// <param name="y">Y coordinate</param>
    /// <param name="limitDistance">Maximum search distance</param>
    /// <returns>The nearest node within limit, or null if none found</returns>
    /// <exception cref="ArgumentOutOfRangeException">Thrown when limitDistance is negative.</exception>
    public KDTreeData? FindNearest(double x, double y, double limitDistance)
    {
        ArgumentOutOfRangeException.ThrowIfNegative(limitDistance);

        if (_root == null)
            return null;

        double limitDistSquared = limitDistance * limitDistance;
        var result = FindNearestRecursive(_root, x, y, null, double.MaxValue, limitDistSquared);

        return result.BestNode;
    }

    /// <summary>
    /// Recursively finds the nearest node using squared distances to avoid sqrt operations.
    /// </summary>
    private static SearchResult FindNearestRecursive(
        KDTreeNode? node,
        double x,
        double y,
        KDTreeData? bestNode,
        double bestDistSquared,
        double limitDistSquared)
    {
        if (node == null)
            return new SearchResult(bestNode, bestDistSquared);

        // Calculate squared distance (avoid sqrt for performance)
        double dx = node.Node.X - x;
        double dy = node.Node.Y - y;
        double distSquared = dx * dx + dy * dy;

        // Update best if this node is closer and within limit
        if (distSquared < bestDistSquared && distSquared <= limitDistSquared)
        {
            bestNode = node.Node;
            bestDistSquared = distSquared;
        }

        // Determine which side to search first
        double delta = node.Axis == 0 ? x - node.Node.X : y - node.Node.Y;
        KDTreeNode? nearSide = delta < 0 ? node.Left : node.Right;
        KDTreeNode? farSide = delta < 0 ? node.Right : node.Left;

        // Search near side
        var result = FindNearestRecursive(nearSide, x, y, bestNode, bestDistSquared, limitDistSquared);
        bestNode = result.BestNode;
        bestDistSquared = result.BestDistSquared;

        // Only search far side if it could contain a closer point
        double deltaSquared = delta * delta;
        if (deltaSquared < bestDistSquared)
        {
            result = FindNearestRecursive(farSide, x, y, bestNode, bestDistSquared, limitDistSquared);
            bestNode = result.BestNode;
            bestDistSquared = result.BestDistSquared;
        }

        return new SearchResult(bestNode, bestDistSquared);
    }

    /// <summary>
    /// Finds all nodes within the specified radius from the given coordinates.
    /// </summary>
    /// <param name="x">X coordinate</param>
    /// <param name="y">Y coordinate</param>
    /// <param name="radius">Search radius</param>
    /// <returns>List of all nodes within the radius</returns>
    public List<KDTreeData> FindInRadius(double x, double y, double radius)
    {
        ArgumentOutOfRangeException.ThrowIfNegative(radius);

        if (_root == null)
            return [];

        var result = new List<KDTreeData>();
        double radiusSquared = radius * radius;
        FindInRadiusRecursive(_root, x, y, radiusSquared, result);
        return result;
    }

    /// <summary>
    /// Recursively finds all nodes within radius.
    /// </summary>
    private static void FindInRadiusRecursive(
        KDTreeNode? node,
        double x,
        double y,
        double radiusSquared,
        List<KDTreeData> result)
    {
        if (node == null)
            return;

        double dx = node.Node.X - x;
        double dy = node.Node.Y - y;
        double distSquared = dx * dx + dy * dy;

        if (distSquared <= radiusSquared)
            result.Add(node.Node);

        double delta = node.Axis == 0 ? x - node.Node.X : y - node.Node.Y;
        double deltaSquared = delta * delta;

        // Search both sides if sphere intersects splitting plane
        if (delta < 0)
        {
            FindInRadiusRecursive(node.Left, x, y, radiusSquared, result);
            if (deltaSquared <= radiusSquared)
                FindInRadiusRecursive(node.Right, x, y, radiusSquared, result);
        }
        else
        {
            FindInRadiusRecursive(node.Right, x, y, radiusSquared, result);
            if (deltaSquared <= radiusSquared)
                FindInRadiusRecursive(node.Left, x, y, radiusSquared, result);
        }
    }

    /// <summary>
    /// Internal struct to return search results without allocations.
    /// </summary>
    private readonly struct SearchResult(KDTreeData? bestNode, double bestDistSquared)
    {
        public readonly KDTreeData? BestNode = bestNode;
        public readonly double BestDistSquared = bestDistSquared;
    }
}