using AutoRig.Analyze;
using AutoRig.Mesh;
using AutoRig.Rig;

namespace AutoRig.Solve;

// 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>
/// Rigs a machine: one bone per rigid part, hierarchy from the contact graph (root =
/// largest part, BFS outward), pivots at contact centers, hinge axes from flat contact
/// interfaces or part elongation, rigid binding.
/// </summary>
public static class MechanicalSolver
{
    public static RigResult Rig( AnalysisResult analysis )
    {
        ArgumentNullException.ThrowIfNull( analysis );
        var parts = analysis.Parts;
        if ( parts.Count == 0 )
            return FloorSolver.Rig( analysis );

        // Contact lookup per pair.
        var contactsByPair = new Dictionary<(int A, int B), List<PartContact>>();
        foreach ( var contact in analysis.Contacts )
        {
            var key = (contact.PartA, contact.PartB);
            if ( !contactsByPair.TryGetValue( key, out var list ) )
                contactsByPair[key] = list = new List<PartContact>();
            list.Add( contact );
        }
        var adjacency = new List<int>[parts.Count];
        for ( var i = 0; i < parts.Count; i++ ) adjacency[i] = new List<int>();
        foreach ( var (a, b) in contactsByPair.Keys )
        {
            adjacency[a].Add( b );
            adjacency[b].Add( a );
        }

        // BFS from the largest part (parts are sorted by area; index 0 is largest).
        var order = new List<int> { 0 };
        var parentPart = new int[parts.Count];
        Array.Fill( parentPart, -1 );
        var visited = new bool[parts.Count];
        visited[0] = true;
        for ( var head = 0; head < order.Count; head++ )
        {
            var current = order[head];
            foreach ( var next in adjacency[current].OrderBy( i => i ) )
            {
                if ( visited[next] )
                    continue;
                visited[next] = true;
                parentPart[next] = current;
                order.Add( next );
            }
        }
        // Disconnected parts attach directly to the root, largest first.
        for ( var i = 0; i < parts.Count; i++ )
        {
            if ( visited[i] )
                continue;
            visited[i] = true;
            parentPart[i] = 0;
            order.Add( i );
        }

        // Emit joints in BFS order (parents precede children by construction).
        var jointIndexOfPart = new int[parts.Count];
        var skeleton = new RigSkeleton();
        var usedNames = new HashSet<string>( StringComparer.Ordinal );
        foreach ( var partIndex in order )
        {
            var part = parts[partIndex];
            var parent = parentPart[partIndex];

            Vector3 position;
            var hinge = Vector3.Zero;
            if ( parent < 0 )
            {
                position = part.Bounds.Center;
            }
            else
            {
                var key = (Math.Min( parent, partIndex ), Math.Max( parent, partIndex ));
                if ( contactsByPair.TryGetValue( key, out var pairContacts ) )
                {
                    position = Vector3.Zero;
                    foreach ( var c in pairContacts ) position += c.Center;
                    position /= pairContacts.Count;
                    hinge = InferHinge( analysis.Mesh, part, pairContacts );
                }
                else
                {
                    position = part.Bounds.Center; // disconnected: pivot at own center
                }
            }

            var name = NameUtil.Sanitize( part.Name );
            if ( !usedNames.Add( name ) )
            {
                var n = 2;
                while ( !usedNames.Add( $"{name}_{n}" ) ) n++;
                name = $"{name}_{n}";
            }

            jointIndexOfPart[partIndex] = skeleton.Joints.Count;
            skeleton.Joints.Add( new RigJoint
            {
                Name = name,
                Parent = parent < 0 ? -1 : jointIndexOfPart[parent],
                Position = position,
                HingeAxis = hinge,
            } );
        }

        // Rigid binding: vertex → its part's joint.
        var vertexBone = new int[analysis.Mesh.Positions.Length];
        foreach ( var part in parts )
        {
            var joint = jointIndexOfPart[part.Index];
            foreach ( var tri in part.TriangleIndices )
                for ( var k = 0; k < 3; k++ )
                    vertexBone[analysis.Mesh.Triangles[tri * 3 + k]] = joint;
        }

        var result = new RigResult
        {
            Skeleton = skeleton,
            Weights = SkinWeights.Rigid( analysis.Mesh.Positions.Length, vertexBone ),
            SolverName = "mechanical",
            Degraded = false,
            Explanation = $"{parts.Count} rigid part(s); root '{skeleton.Joints[skeleton.Root].Name}'.",
        };
        result.Skeleton.Validate();
        result.Weights.Validate( analysis.Mesh, result.Skeleton );
        return result;
    }

    /// <summary>
    /// Hinge for a part articulating at its parent contact: the contact plane normal
    /// when the interface is flat; else the part's own elongation axis when clearly
    /// cylindrical; else free (Zero).
    /// </summary>
    static Vector3 InferHinge( RigMesh mesh, MeshPart part, List<PartContact> contacts )
    {
        // Flat interface: contact cloud planar → plane normal is the rotation axis.
        var main = contacts[0];
        foreach ( var c in contacts )
            if ( c.SampleCount > main.SampleCount )
                main = c;

        if ( main.Extent > 0f && IsPlanar( main ) )
            return main.PlaneNormal;

        // Elongated (shaft-like) part: rotate about its own long axis.
        var points = new List<Vector3>( part.TriangleIndices.Length );
        var stride = Math.Max( 1, part.TriangleIndices.Length / 2000 );
        for ( var i = 0; i < part.TriangleIndices.Length; i += stride )
            points.Add( mesh.Positions[mesh.Triangles[part.TriangleIndices[i] * 3]] );
        if ( points.Count >= 3 )
        {
            var center = Vector3.Zero;
            foreach ( var p in points ) center += p;
            center /= points.Count;
            var axes = PartContacts.PrincipalAxes( points, center );
            var (v0, v1) = (Variance( points, center, axes[0] ), Variance( points, center, axes[1] ));
            if ( v1 > 1e-12f && v0 / v1 > 4f )
                return axes[0];
        }
        return Vector3.Zero;
    }

    /// <summary>Planar when the contact cloud's normal-direction spread is tiny vs its extent.</summary>
    static bool IsPlanar( PartContact contact )
        // With only aggregate data (center/extent/axes), accept flat interfaces via the
        // plane normal being well-defined: Extent > 0 and the cloud not degenerate.
        // The PlaneNormal is the smallest-variance axis by construction; a spherical
        // cloud would give an arbitrary but still harmless axis, so gate on extent.
        => contact.SampleCount >= 4 && contact.Extent > 1e-6f;

    static float Variance( List<Vector3> points, Vector3 center, Vector3 axis )
    {
        float total = 0;
        foreach ( var p in points )
        {
            var d = Vector3.Dot( p - center, axis );
            total += d * d;
        }
        return total / points.Count;
    }

}