using System.Numerics;
using AutoRig.Mesh;

namespace AutoRig.Analyze;

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


/// <summary>
/// Splits a mesh into connected parts: positions within the weld tolerance are treated
/// as one vertex, and triangles sharing a welded vertex belong to the same part.
/// </summary>
public static class MeshSegmenter
{
    /// <summary>
    /// Vertex → canonical vertex map under positional welding. Tolerance &lt;= 0 selects an
    /// automatic tolerance of 1e-5 × the largest bounds dimension.
    /// </summary>
    public static int[] WeldMap( RigMesh mesh, float tolerance )
    {
        ArgumentNullException.ThrowIfNull( mesh );
        if ( tolerance <= 0f )
        {
            var size = mesh.ComputeBounds().Size;
            tolerance = 1e-5f * MathF.Max( size.X, MathF.Max( size.Y, size.Z ) );
            if ( tolerance <= 0f )
                tolerance = 1e-7f;
        }

        var cell = tolerance;
        var grid = new Dictionary<(long, long, long), List<int>>();
        var map = new int[mesh.Positions.Length];
        var toleranceSquared = tolerance * tolerance;

        for ( var i = 0; i < mesh.Positions.Length; i++ )
        {
            var p = mesh.Positions[i];
            var cx = (long)MathF.Floor( p.X / cell );
            var cy = (long)MathF.Floor( p.Y / cell );
            var cz = (long)MathF.Floor( p.Z / cell );

            var canonical = -1;
            for ( var dx = -1L; dx <= 1 && canonical < 0; dx++ )
                for ( var dy = -1L; dy <= 1 && canonical < 0; dy++ )
                    for ( var dz = -1L; dz <= 1 && canonical < 0; dz++ )
                    {
                        if ( !grid.TryGetValue( (cx + dx, cy + dy, cz + dz), out var bucket ) )
                            continue;
                        foreach ( var candidate in bucket )
                        {
                            if ( Vector3.DistanceSquared( mesh.Positions[candidate], p ) <= toleranceSquared )
                            {
                                canonical = candidate;
                                break;
                            }
                        }
                    }

            if ( canonical < 0 )
            {
                canonical = i;
                if ( !grid.TryGetValue( (cx, cy, cz), out var own ) )
                    grid[(cx, cy, cz)] = own = new List<int>();
                own.Add( i );
            }
            map[i] = canonical;
        }
        return map;
    }

    /// <summary>Segments the mesh into connected parts, sorted by descending surface area.</summary>
    public static IReadOnlyList<MeshPart> Segment( RigMesh mesh, float weldTolerance )
    {
        ArgumentNullException.ThrowIfNull( mesh );
        mesh.Validate();

        var weld = WeldMap( mesh, weldTolerance );

        // Union-find over canonical vertices.
        var parent = new int[mesh.Positions.Length];
        for ( var i = 0; i < parent.Length; i++ ) parent[i] = weld[i];

        int Find( int i )
        {
            while ( parent[i] != i )
            {
                parent[i] = parent[parent[i]];
                i = parent[i];
            }
            return i;
        }
        void Union( int a, int b )
        {
            a = Find( a ); b = Find( b );
            if ( a != b ) parent[b] = a;
        }

        for ( var t = 0; t < mesh.Triangles.Length; t += 3 )
        {
            Union( mesh.Triangles[t], mesh.Triangles[t + 1] );
            Union( mesh.Triangles[t], mesh.Triangles[t + 2] );
        }

        // Group triangles by component root.
        var byRoot = new Dictionary<int, List<int>>();
        for ( var tri = 0; tri < mesh.TriangleCount; tri++ )
        {
            var root = Find( mesh.Triangles[tri * 3] );
            if ( !byRoot.TryGetValue( root, out var list ) )
                byRoot[root] = list = new List<int>();
            list.Add( tri );
        }

        // Assemble parts.
        var assembled = new List<(int[] Triangles, Aabb3 Bounds, float Area, string Tag, int VertexCount)>();
        foreach ( var list in byRoot.Values )
        {
            var vertices = new HashSet<int>();
            var tagVotes = new Dictionary<int, int>();
            float area = 0;
            var min = new Vector3( float.MaxValue );
            var max = new Vector3( float.MinValue );

            foreach ( var tri in list )
            {
                var a = mesh.Positions[mesh.Triangles[tri * 3]];
                var b = mesh.Positions[mesh.Triangles[tri * 3 + 1]];
                var c = mesh.Positions[mesh.Triangles[tri * 3 + 2]];
                area += Vector3.Cross( b - a, c - a ).Length() * 0.5f;
                for ( var k = 0; k < 3; k++ )
                {
                    var v = mesh.Triangles[tri * 3 + k];
                    vertices.Add( v );
                    min = Vector3.Min( min, mesh.Positions[v] );
                    max = Vector3.Max( max, mesh.Positions[v] );
                }
                var tag = mesh.TriangleTags[tri];
                tagVotes[tag] = tagVotes.GetValueOrDefault( tag ) + 1;
            }

            var dominantTag = 0;
            var bestVotes = -1;
            foreach ( var (tag, votes) in tagVotes )
                if ( votes > bestVotes ) { bestVotes = votes; dominantTag = tag; }

            assembled.Add( (list.ToArray(), new Aabb3( min, max ), area,
                mesh.Tags.Length > 0 ? mesh.Tags[dominantTag] : "part", vertices.Count) );
        }

        assembled.Sort( ( x, y ) => y.Area.CompareTo( x.Area ) );

        // Unique names: dominant tag, suffixed on repeats.
        var used = new Dictionary<string, int>( StringComparer.Ordinal );
        var parts = new List<MeshPart>( assembled.Count );
        foreach ( var (triangles, bounds, area, tag, vertexCount) in assembled )
        {
            var name = tag;
            if ( used.TryGetValue( tag, out var seen ) )
                name = $"{tag}_{seen + 1}";
            used[tag] = used.GetValueOrDefault( tag ) + 1;

            parts.Add( new MeshPart
            {
                Index = parts.Count,
                TriangleIndices = triangles,
                Bounds = bounds,
                SurfaceArea = area,
                Name = name,
                VertexCount = vertexCount,
            } );
        }
        return parts;
    }
}