#nullable enable
using Editor;
using Sandbox;
using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;

namespace BlenderActions;

/// <summary>Represents the result of locating a target vertex.</summary>
public readonly record struct VertexSnapResult(bool Found, Vector3 Vertex);

/// <summary>Stores reusable world-space source vertices for vertex snapping.</summary>
public sealed class VertexSnapSource
{
    /// <summary>Handles new.</summary>
    internal List<Vector3> Vertices { get; } = new(4096);
    /// <summary>Handles new.</summary>
    internal HashSet<ModelRenderer> VisitedRenderers { get; } = new();

    /// <summary>Clears reusable source-vertex collections.</summary>
    internal void Clear()
    {
        Vertices.Clear();
        VisitedRenderers.Clear();
    }
}

/// <summary>Provides cached screen-space vertex snapping for modal transforms.</summary>
public static class VertexSnapService
{
    /// <summary>Defines the maximum world distance used to trace a target renderer.</summary>
    private const float TraceLength = 100000f;
    /// <summary>Defines the logical screen-space radius used to acquire target vertices.</summary>
    private const float SnapRadiusPixels = 16f;
    /// <summary>Defines the projected-vertex spatial hash cell size in pixels.</summary>
    private const float ProjectionCellSize = 32f;
    /// <summary>Defines local-vertex quantization used for model vertex deduplication.</summary>
    private const float Quantization = 10000f;

    /// <summary>Handles new.</summary>
    private static ConditionalWeakTable<Model, CachedVertices> _vertexCache = new();
    /// <summary>Handles new.</summary>
    private static ConditionalWeakTable<ModelRenderer, ProjectedVertexIndex> _targetIndices = new();

    /// <summary>Captures current world-space vertices from selected model renderers.</summary>
    public static void CaptureSourceSnapshot(
        VertexSnapSource destination,
        IReadOnlyCollection<GameObject> selectedObjects)
    {
        destination.Clear();

        foreach(var selectedObject in selectedObjects)
        {
            if(!selectedObject.IsValid())
                continue;

            foreach(var renderer in selectedObject.GetComponentsInChildren<ModelRenderer>(
                includeDisabled: true,
                includeSelf: true))
            {
                if(renderer == null ||
                    !destination.VisitedRenderers.Add(renderer) ||
                    renderer.Model == null ||
                    !renderer.Model.IsValid)
                {
                    continue;
                }

                AppendWorldVertices(renderer, destination.Vertices);
            }
        }
    }

    /// <summary>Finds the nearest target vertex under the pointer on the traced renderer.</summary>
    public static VertexSnapResult FindTargetVertex(
        Scene scene,
        CameraComponent camera,
        SceneViewportWidget viewport,
        IReadOnlyCollection<GameObject> ignoredObjects)
    {
        var mousePosition = SceneViewportWidget.MousePosition;
        var ray = camera.ScreenPixelToRay(mousePosition);
        var trace = scene.Trace
            .Ray(ray, TraceLength)
            .UseRenderMeshes(true, true)
            .UseHitPosition(true);

        foreach(var gameObject in ignoredObjects)
        {
            if(gameObject.IsValid())
                trace = trace.IgnoreGameObjectHierarchy(gameObject);
        }

        var hit = trace.Run();

        if(!hit.Hit || hit.GameObject == null)
            return default;

        var renderer = hit.Component as ModelRenderer ??
            hit.GameObject.GetComponent<ModelRenderer>(true);

        if(renderer == null || renderer.Model == null || !renderer.Model.IsValid)
            return default;

        var threshold = SnapRadiusPixels * MathF.Max(viewport.DpiScale, 1f);
        var index = _targetIndices.GetValue(renderer, _ => new ProjectedVertexIndex());
        index.Update(renderer, camera);

        return index.TryFindNearest(mousePosition, threshold, out var vertex)
            ? new VertexSnapResult(true, vertex)
            : default;
    }

    /// <summary>Finds the source vertex closest on screen after translation.</summary>
    public static bool TryFindClosestTranslatedSource(
        VertexSnapSource source,
        CameraComponent camera,
        Vector3 target,
        Vector3 translation,
        out Vector3 sourceVertex)
    {
        return TryFindClosestSource(
            source,
            camera,
            target,
            SourceTransform.Translate,
            translation,
            Vector3.Zero,
            Rotation.Identity,
            Vector3.One,
            out sourceVertex);
    }

    /// <summary>Finds the source vertex closest on screen after rotation.</summary>
    public static bool TryFindClosestRotatedSource(
        VertexSnapSource source,
        CameraComponent camera,
        Vector3 target,
        Vector3 pivot,
        Rotation rotation,
        out Vector3 sourceVertex)
    {
        return TryFindClosestSource(
            source,
            camera,
            target,
            SourceTransform.Rotate,
            Vector3.Zero,
            pivot,
            rotation,
            Vector3.One,
            out sourceVertex);
    }

    /// <summary>Finds the source vertex closest on screen after scaling.</summary>
    public static bool TryFindClosestScaledSource(
        VertexSnapSource source,
        CameraComponent camera,
        Vector3 target,
        Vector3 pivot,
        Vector3 multiplier,
        out Vector3 sourceVertex)
    {
        return TryFindClosestSource(
            source,
            camera,
            target,
            SourceTransform.Scale,
            Vector3.Zero,
            pivot,
            Rotation.Identity,
            multiplier,
            out sourceVertex);
    }

    /// <summary>Clears model and projected-vertex caches after hotload.</summary>
    [EditorEvent.Hotload]
    private static void ClearCaches()
    {
        _vertexCache = new ConditionalWeakTable<Model, CachedVertices>();
        _targetIndices = new ConditionalWeakTable<ModelRenderer, ProjectedVertexIndex>();
    }

    /// <summary>Finds the screen-space closest source vertex after a supplied transform.</summary>
    private static bool TryFindClosestSource(
        VertexSnapSource source,
        CameraComponent camera,
        Vector3 target,
        SourceTransform transform,
        Vector3 translation,
        Vector3 pivot,
        Rotation rotation,
        Vector3 multiplier,
        out Vector3 sourceVertex)
    {
        sourceVertex = Vector3.Zero;

        if(source.Vertices.Count == 0)
            return false;

        var targetScreen = camera.PointToScreenPixels(target, out var targetBehind);

        if(targetBehind)
            return false;

        var bestDistance = float.MaxValue;
        var found = false;

        for(var index = 0; index < source.Vertices.Count; index++)
        {
            var original = source.Vertices[index];
            var candidate = transform switch
            {
                SourceTransform.Translate => original + translation,
                SourceTransform.Rotate => pivot + rotation * (original - pivot),
                SourceTransform.Scale => pivot + (original - pivot).MultiplyComponents(multiplier),
                _ => original
            };

            var screen = camera.PointToScreenPixels(candidate, out var isBehind);

            if(isBehind)
                continue;

            var distance = (screen - targetScreen).Length;

            if(distance >= bestDistance)
                continue;

            bestDistance = distance;
            sourceVertex = candidate;
            found = true;
        }

        return found;
    }

    /// <summary>Appends one renderer's transformed model vertices to a reusable destination.</summary>
    private static void AppendWorldVertices(
        ModelRenderer renderer,
        List<Vector3> destination)
    {
        var vertices = GetVertices(renderer.Model);

        for(var index = 0; index < vertices.Length; index++)
            destination.Add(ToWorld(renderer.GameObject, vertices[index]));
    }

    /// <summary>Transforms a local model vertex into world space.</summary>
    private static Vector3 ToWorld(GameObject gameObject, Vector3 localVertex)
    {
        var scaled = localVertex.MultiplyComponents(gameObject.WorldScale);
        return gameObject.WorldPosition + gameObject.WorldRotation * scaled;
    }

    /// <summary>Returns cached deduplicated local-space vertices for a model.</summary>
    private static Vector3[] GetVertices(Model model)
    {
        return _vertexCache.GetValue(model, CreateCache).Vertices;
    }

    /// <summary>Creates a deduplicated local-space vertex cache for a model.</summary>
    private static CachedVertices CreateCache(Model model)
    {
        var unique = new Dictionary<QuantizedVertex, Vector3>();

        foreach(var vertex in model.GetVertices())
        {
            var position = vertex.Position;
            var key = new QuantizedVertex(
                (int)MathF.Round(position.x * Quantization),
                (int)MathF.Round(position.y * Quantization),
                (int)MathF.Round(position.z * Quantization));

            if(!unique.ContainsKey(key))
                unique.Add(key, position);
        }

        var vertices = new Vector3[unique.Count];
        unique.Values.CopyTo(vertices, 0);
        return new CachedVertices(vertices);
    }

    /// <summary>Combines two screen-space cell coordinates into one dictionary key.</summary>
    private static long CellKey(int x, int y)
    {
        return ((long)x << 32) ^ (uint)y;
    }

    /// <summary>Identifies the transform applied while evaluating source vertices.</summary>
    private enum SourceTransform
    {
        Translate,
        Rotate,
        Scale
    }

    /// <summary>Stores deduplicated local-space vertices for one model.</summary>
    private sealed class CachedVertices
    {
        /// <summary>Initializes a new cached vertices instance.</summary>
        public CachedVertices(Vector3[] vertices)
        {
            Vertices = vertices;
        }

        /// <summary>Gets the reusable captured world-space vertex list.</summary>
        public Vector3[] Vertices { get; }
    }

    /// <summary>Indexes one renderer's projected vertices in screen-space cells.</summary>
    private sealed class ProjectedVertexIndex
    {
        /// <summary>Handles new.</summary>
        private readonly Dictionary<long, List<ProjectedVertex>> _cells = new();
        /// <summary>Handles new.</summary>
        private readonly Stack<List<ProjectedVertex>> _bucketPool = new();

        /// <summary>Stores the model represented by the current projected index.</summary>
        private Model? _model;
        /// <summary>Stores the indexed renderer world position.</summary>
        private Vector3 _objectPosition;
        /// <summary>Stores the indexed renderer world rotation.</summary>
        private Rotation _objectRotation;
        /// <summary>Stores the indexed renderer world scale.</summary>
        private Vector3 _objectScale;
        /// <summary>Stores the camera position used to build the index.</summary>
        private Vector3 _cameraPosition;
        /// <summary>Stores the camera rotation used to build the index.</summary>
        private Rotation _cameraRotation;
        /// <summary>Stores the camera render size used to build the index.</summary>
        private Vector2? _cameraSize;
        /// <summary>Stores the perspective field of view used to build the index.</summary>
        private float _fieldOfView;
        /// <summary>Stores the orthographic height used to build the index.</summary>
        private float _orthographicHeight;
        /// <summary>Tracks whether the indexed camera uses orthographic projection.</summary>
        private bool _orthographic;

        /// <summary>Rebuilds the projected vertex index when renderer or camera state changes.</summary>
        public void Update(ModelRenderer renderer, CameraComponent camera)
        {
            var gameObject = renderer.GameObject;
            var cameraObject = camera.GameObject;

            if(ReferenceEquals(_model, renderer.Model) &&
                _objectPosition.Equals(gameObject.WorldPosition) &&
                _objectRotation.Equals(gameObject.WorldRotation) &&
                _objectScale.Equals(gameObject.WorldScale) &&
                _cameraPosition.Equals(cameraObject.WorldPosition) &&
                _cameraRotation.Equals(cameraObject.WorldRotation) &&
                _cameraSize.Equals(camera.CustomSize) &&
                _fieldOfView.Equals(camera.FieldOfView) &&
                _orthographicHeight.Equals(camera.OrthographicHeight) &&
                _orthographic == camera.Orthographic)
            {
                return;
            }

            RecycleCells();
            _model = renderer.Model;
            _objectPosition = gameObject.WorldPosition;
            _objectRotation = gameObject.WorldRotation;
            _objectScale = gameObject.WorldScale;
            _cameraPosition = cameraObject.WorldPosition;
            _cameraRotation = cameraObject.WorldRotation;
            _cameraSize = camera.CustomSize;
            _fieldOfView = camera.FieldOfView;
            _orthographicHeight = camera.OrthographicHeight;
            _orthographic = camera.Orthographic;

            var vertices = GetVertices(renderer.Model);

            for(var index = 0; index < vertices.Length; index++)
            {
                var world = ToWorld(gameObject, vertices[index]);
                var screen = camera.PointToScreenPixels(world, out var isBehind);

                if(isBehind)
                    continue;

                var cellX = (int)MathF.Floor(screen.x / ProjectionCellSize);
                var cellY = (int)MathF.Floor(screen.y / ProjectionCellSize);
                var key = CellKey(cellX, cellY);

                if(!_cells.TryGetValue(key, out var bucket))
                {
                    bucket = _bucketPool.Count > 0
                        ? _bucketPool.Pop()
                        : new List<ProjectedVertex>();
                    _cells.Add(key, bucket);
                }

                bucket.Add(new ProjectedVertex(screen, world));
            }
        }

        /// <summary>Finds the nearest indexed vertex within a screen-space radius.</summary>
        public bool TryFindNearest(
            Vector2 screenPosition,
            float radius,
            out Vector3 worldVertex)
        {
            worldVertex = Vector3.Zero;
            var centerX = (int)MathF.Floor(screenPosition.x / ProjectionCellSize);
            var centerY = (int)MathF.Floor(screenPosition.y / ProjectionCellSize);
            var cellRadius = Math.Max(1, (int)MathF.Ceiling(radius / ProjectionCellSize));
            var bestSquaredDistance = radius * radius;
            var found = false;

            for(var x = centerX - cellRadius; x <= centerX + cellRadius; x++)
            {
                for(var y = centerY - cellRadius; y <= centerY + cellRadius; y++)
                {
                    if(!_cells.TryGetValue(CellKey(x, y), out var bucket))
                        continue;

                    for(var index = 0; index < bucket.Count; index++)
                    {
                        var candidate = bucket[index];
                        var delta = candidate.Screen - screenPosition;
                        var squaredDistance = delta.x * delta.x + delta.y * delta.y;

                        if(squaredDistance > bestSquaredDistance)
                            continue;

                        bestSquaredDistance = squaredDistance;
                        worldVertex = candidate.World;
                        found = true;
                    }
                }
            }

            return found;
        }

        /// <summary>Clears projected cells and returns their lists to the bucket pool.</summary>
        private void RecycleCells()
        {
            foreach(var bucket in _cells.Values)
            {
                bucket.Clear();
                _bucketPool.Push(bucket);
            }

            _cells.Clear();
        }
    }

    /// <summary>Pairs a projected screen position with its world-space vertex.</summary>
    private readonly record struct ProjectedVertex(Vector2 Screen, Vector3 World);
    /// <summary>Provides a quantized key for deduplicating model vertices.</summary>
    private readonly record struct QuantizedVertex(int X, int Y, int Z);
}