using AutoRig.Voxel;

namespace AutoRig.Solve.Organic;

// s&box compat: the engine defines Vector2/Vector3 in the GLOBAL namespace, which
// shadows using-directive imports - alias explicitly to System.Numerics.
using Vector3 = System.Numerics.Vector3;

/// <summary>One node of a curve skeleton.</summary>
public sealed class SkeletonNode
{
    public required Vector3 Position { get; init; }

    /// <summary>Local half-thickness of the volume (boundary distance × cell size).</summary>
    public required float Thickness { get; init; }

    public List<int> Neighbors { get; } = new();
}

/// <summary>A curve skeleton: nodes with edges, tips (degree 1) and the core node.</summary>
public sealed class SkeletonGraph
{
    public List<SkeletonNode> Nodes { get; } = new();

    /// <summary>Indices of degree-1 nodes (excluding an isolated core).</summary>
    public List<int> Tips { get; } = new();

    /// <summary>Index of the core (deepest) node.</summary>
    public int Core { get; set; }
}

/// <summary>
/// Extracts a curve skeleton from a solid voxel grid: from each extremity tip, walk
/// down the core-geodesic field along the medial ridge (prefer deep voxels); paths
/// merge where they touch earlier paths (junctions); voxel paths are resampled into
/// graph nodes.
/// </summary>
public static class CurveSkeleton
{
    const int ResampleInterval = 2;

    public static SkeletonGraph Extract( VoxelGrid grid )
    {
        ArgumentNullException.ThrowIfNull( grid );

        var field = DistanceField.FromGrid( grid );
        var core = Extremities.Core( grid, field );
        var coreGeodesic = DistanceField.GeodesicFrom( grid, core );
        var tips = Extremities.Find( grid, field );

        var graph = new SkeletonGraph();

        // Core node always exists.
        graph.Core = 0;
        graph.Nodes.Add( new SkeletonNode
        {
            Position = grid.CenterOf( core.X, core.Y, core.Z ),
            Thickness = field.At( core.X, core.Y, core.Z ) * grid.CellSize,
        } );

        if ( tips.Count == 0 )
            return graph;

        // Voxel → (path id, raw index) claims.
        var claims = new Dictionary<(int X, int Y, int Z), (int Path, int Index)>();
        var pathNodesOf = new List<Dictionary<int, int>>(); // per path: raw index → graph node

        for ( var pathId = 0; pathId < tips.Count; pathId++ )
        {
            var raw = new List<(int X, int Y, int Z)>();
            var current = tips[pathId];
            var terminalNode = -1; // node the finished path connects into

            while ( true )
            {
                if ( current == core )
                {
                    terminalNode = graph.Core;
                    break;
                }
                if ( claims.TryGetValue( current, out var claim ) )
                {
                    var owner = pathNodesOf[claim.Path];
                    terminalNode = owner.Count == 0
                        ? graph.Core
                        : owner[NearestSampledIndex( claim.Index, owner )];
                    break;
                }

                raw.Add( current );

                // Step: among 26-neighbors with strictly smaller core distance, prefer
                // the deepest (medial ridge); ties by smaller core distance then order.
                var currentDistance = coreGeodesic[grid.Index( current.X, current.Y, current.Z )];
                (int X, int Y, int Z) best = default;
                var bestDepth = -1;
                var bestDistance = int.MaxValue;
                var found = false;
                for ( var dz = -1; dz <= 1; dz++ )
                    for ( var dy = -1; dy <= 1; dy++ )
                        for ( var dx = -1; dx <= 1; dx++ )
                        {
                            if ( dx == 0 && dy == 0 && dz == 0 )
                                continue;
                            var nx = current.X + dx; var ny = current.Y + dy; var nz = current.Z + dz;
                            if ( nx < 0 || nx >= grid.SizeX || ny < 0 || ny >= grid.SizeY
                                || nz < 0 || nz >= grid.SizeZ || !grid.IsSolid( nx, ny, nz ) )
                                continue;
                            var d = coreGeodesic[grid.Index( nx, ny, nz )];
                            if ( d == int.MaxValue || d >= currentDistance )
                                continue;
                            var depth = field.At( nx, ny, nz );
                            if ( depth > bestDepth || (depth == bestDepth && d < bestDistance) )
                            {
                                best = (nx, ny, nz);
                                bestDepth = depth;
                                bestDistance = d;
                                found = true;
                            }
                        }
                if ( !found )
                {
                    terminalNode = graph.Core; // isolated pocket: connect to core directly
                    break;
                }
                current = best;
            }

            // Claim raw voxels and resample into nodes.
            var nodeOf = new Dictionary<int, int>();
            var previousNode = -1;
            for ( var i = 0; i < raw.Count; i++ )
            {
                claims[raw[i]] = (pathId, i);
                if ( i % ResampleInterval != 0 )
                    continue;
                var voxel = raw[i];
                var node = graph.Nodes.Count;
                graph.Nodes.Add( new SkeletonNode
                {
                    Position = grid.CenterOf( voxel.X, voxel.Y, voxel.Z ),
                    Thickness = field.At( voxel.X, voxel.Y, voxel.Z ) * grid.CellSize,
                } );
                nodeOf[i] = node;
                if ( previousNode >= 0 )
                    Connect( graph, previousNode, node );
                previousNode = node;
            }
            // Map every raw index to its nearest sampled node for junction lookups.
            pathNodesOf.Add( nodeOf );

            // Connect the path's inner end to its terminal (core or junction node).
            if ( previousNode >= 0 && terminalNode >= 0 )
                Connect( graph, previousNode, terminalNode );
            else if ( previousNode < 0 && terminalNode >= 0 && raw.Count == 0 )
            {
                // Tip immediately terminated (tip == core or claimed): nothing to add.
            }
        }

        // Prune stubby branches: a genuine limb is long relative to the thickness of
        // the junction it hangs off; the corner of a fat box-end produces a branch
        // that dives into the deep body within a few steps. Iterate to stability
        // (removing a stub can turn its junction into a stub's base again).
        PruneStubs( graph, grid.CellSize );

        // Tips = degree-1 nodes that are not the core.
        for ( var i = 0; i < graph.Nodes.Count; i++ )
            if ( graph.Nodes[i].Neighbors.Count == 1 && i != graph.Core )
                graph.Tips.Add( i );

        return graph;
    }

    static void PruneStubs( SkeletonGraph graph, float cellSize )
    {
        var removed = new bool[graph.Nodes.Count];
        bool changed;
        do
        {
            changed = false;
            for ( var tip = 0; tip < graph.Nodes.Count; tip++ )
            {
                if ( removed[tip] || tip == graph.Core || graph.Nodes[tip].Neighbors.Count != 1 )
                    continue;

                // Walk through degree-2 nodes to the branch's junction (degree >= 3 or core).
                var branch = new List<int> { tip };
                var previous = -1;
                var current = tip;
                while ( true )
                {
                    var next = graph.Nodes[current].Neighbors.FirstOrDefault( n => n != previous, -1 );
                    if ( next < 0 )
                        break;
                    if ( graph.Nodes[next].Neighbors.Count >= 3 || next == graph.Core )
                    {
                        // A stub is short relative to its junction AND stays fat along
                        // its length (a box corner never leaves the body volume). A
                        // short but THIN branch (a chunky duck's neck) is a real limb.
                        var lengthSteps = branch.Count * ResampleInterval;
                        var junctionThickness = graph.Nodes[next].Thickness / cellSize;
                        var thicknesses = branch
                            .Select( n => graph.Nodes[n].Thickness / cellSize )
                            .OrderBy( t => t )
                            .ToList();
                        var medianThickness = thicknesses[thicknesses.Count / 2];
                        if ( lengthSteps < 2f * junctionThickness
                            && medianThickness > 0.45f * junctionThickness )
                        {
                            foreach ( var node in branch )
                            {
                                removed[node] = true;
                                foreach ( var neighbor in graph.Nodes[node].Neighbors )
                                    graph.Nodes[neighbor].Neighbors.Remove( node );
                                graph.Nodes[node].Neighbors.Clear();
                            }
                            changed = true;
                        }
                        break;
                    }
                    branch.Add( next );
                    previous = current;
                    current = next;
                }
            }
        } while ( changed );

        // Compact: rebuild node list without removed nodes.
        if ( !removed.Any( r => r ) )
            return;
        var remap = new int[graph.Nodes.Count];
        var kept = new List<SkeletonNode>();
        for ( var i = 0; i < graph.Nodes.Count; i++ )
        {
            if ( removed[i] )
            {
                remap[i] = -1;
                continue;
            }
            remap[i] = kept.Count;
            kept.Add( graph.Nodes[i] );
        }
        foreach ( var node in kept )
        {
            for ( var n = node.Neighbors.Count - 1; n >= 0; n-- )
            {
                var mapped = remap[node.Neighbors[n]];
                if ( mapped < 0 )
                    node.Neighbors.RemoveAt( n );
                else
                    node.Neighbors[n] = mapped;
            }
        }
        graph.Core = remap[graph.Core] >= 0 ? remap[graph.Core] : 0;
        graph.Nodes.Clear();
        graph.Nodes.AddRange( kept );
    }

    static int NearestSampledIndex( int rawIndex, Dictionary<int, int> nodeOf )
    {
        if ( nodeOf.Count == 0 )
            return 0; // degenerate: caller connects to core
        var best = -1;
        var bestDelta = int.MaxValue;
        foreach ( var key in nodeOf.Keys )
        {
            var delta = Math.Abs( key - rawIndex );
            if ( delta < bestDelta )
            {
                bestDelta = delta;
                best = key;
            }
        }
        return best;
    }

    static void Connect( SkeletonGraph graph, int a, int b )
    {
        if ( a == b )
            return;
        if ( !graph.Nodes[a].Neighbors.Contains( b ) )
            graph.Nodes[a].Neighbors.Add( b );
        if ( !graph.Nodes[b].Neighbors.Contains( a ) )
            graph.Nodes[b].Neighbors.Add( a );
    }
}