using System.Numerics;
using AutoRig.Mesh;

namespace AutoRig.Formats.Fbx;

using Vector2 = System.Numerics.Vector2;
using Vector3 = System.Numerics.Vector3;

/// <summary>
/// Extracts a combined <see cref="RigMesh"/> from an FBX file (binary or ASCII 7.x):
/// reads Geometry objects (Vertices / PolygonVertexIndex / LayerElementNormal /
/// LayerElementUV), connects each geometry to its Model via the OO connection table,
/// bakes the model's evaluated world transform, and tags triangles with the model name.
/// </summary>
public static class FbxMeshImporter
{
    /// <exception cref="FormatException">Malformed file, FBX 6.x, or no triangle geometry.</exception>
    public static RigMesh Import( byte[] data, string sourceName )
        => ImportWithSource( data, sourceName ).Mesh;

    /// <summary>
    /// Import that also reports, per output vertex, which geometry control point it
    /// came from (corner dedup splits control points), plus the parsed scene — the
    /// donor importer maps FBX skin clusters (control-point indexed) through this.
    /// </summary>
    internal static (RigMesh Mesh, List<(long GeometryId, int ControlPoint)> VertexSource,
        FbxNode Root, FbxScene Scene) ImportWithSource( byte[] data, string sourceName )
    {
        var root = FbxTokenizer.Parse( data );
        var scene = FbxScene.Build( root );
        var vertexSource = new List<(long GeometryId, int ControlPoint)>();

        // Geometry id → owning Model, from raw OO connections (child geometry → parent model);
        // plus the material chain Model ← Material ← Texture ← Video (embedded Content).
        var geometryToModel = new Dictionary<long, FbxObject>();
        var modelToMaterial = new Dictionary<long, FbxObject>();
        var materialToTexture = new Dictionary<long, long>();
        var textureToVideo = new Dictionary<long, FbxObject>();
        if ( root.Child( "Connections" ) is { } connections )
        {
            foreach ( var c in connections.ChildrenNamed( "C" ) )
            {
                if ( c.Properties.Count < 3
                    || c.Properties[0] is not string kind || (kind != "OO" && kind != "OP")
                    || c.Properties[1] is not (long or int)
                    || c.Properties[2] is not (long or int) )
                    continue;

                var srcId = c.Prop<long>( 1 );
                var dstId = c.Prop<long>( 2 );
                if ( !scene.ObjectsById.TryGetValue( srcId, out var src )
                    || !scene.ObjectsById.TryGetValue( dstId, out var dst ) )
                    continue;

                if ( kind == "OO" && src.NodeType == "Geometry" && dst.NodeType == "Model" )
                    geometryToModel.TryAdd( srcId, dst );
                else if ( kind == "OO" && src.NodeType == "Material" && dst.NodeType == "Model" )
                    modelToMaterial.TryAdd( dstId, src );
                else if ( src.NodeType == "Texture" && dst.NodeType == "Material" )
                    materialToTexture.TryAdd( dstId, srcId );   // any channel; diffuse first wins
                else if ( kind == "OO" && src.NodeType == "Video" && dst.NodeType == "Texture" )
                    textureToVideo.TryAdd( dstId, src );
            }
        }

        var worldCache = new Dictionary<long, Matrix4x4>();
        Matrix4x4 WorldOf( FbxObject model )
        {
            if ( worldCache.TryGetValue( model.Id, out var cached ) )
                return cached;
            var local = FbxTransform.FromModel( scene, model ).LocalMatrixDefault();
            var world = model.ModelParent is { } parent ? local * WorldOf( parent ) : local;
            worldCache[model.Id] = world;
            return world;
        }

        var outPositions = new List<Vector3>();
        var outNormals = new List<Vector3>();
        var outUvs = new List<Vector2>();
        var outTriangles = new List<int>();
        var outTriangleTags = new List<int>();
        var outTriangleMaterials = new List<int>();
        var tags = new List<string>();
        var tagIndex = new Dictionary<string, int>( StringComparer.Ordinal );
        var anyNormal = false;
        var anyUv = false;

        // FBX material object → RigMaterial index (built on first use).
        var materials = new List<RigMaterial>();
        var materialIndexById = new Dictionary<long, int>();
        int MaterialOf( FbxObject model )
        {
            if ( model is null || !modelToMaterial.TryGetValue( model.Id, out var material ) )
                return -1;
            if ( materialIndexById.TryGetValue( material.Id, out var existing ) )
                return existing;

            var rigMaterial = new RigMaterial
            {
                Name = material.Name is { Length: > 0 } n ? n : $"material{materials.Count}",
            };
            if ( material.Properties.TryGetValue( "DiffuseColor", out var diffuse )
                && diffuse.Values.Count >= 3 )
                rigMaterial.Tint = diffuse.GetVector3();
            if ( materialToTexture.TryGetValue( material.Id, out var textureId ) )
            {
                if ( textureToVideo.TryGetValue( textureId, out var video )
                    && video.Node.Child( "Content" ) is { Properties.Count: > 0 } content
                    && content.Properties[0] is byte[] imageBytes && imageBytes.Length > 0 )
                    rigMaterial.BaseColorImage = imageBytes;
                if ( scene.ObjectsById.TryGetValue( textureId, out var texture ) )
                {
                    var relative = texture.Node.Child( "RelativeFilename" )
                        ?? texture.Node.Child( "FileName" );
                    if ( relative is { Properties.Count: > 0 }
                        && relative.Properties[0] is string relativePath
                        && relativePath.Length > 0 )
                        rigMaterial.SourceImagePath = relativePath;
                }
            }

            materialIndexById[material.Id] = materials.Count;
            materials.Add( rigMaterial );
            return materials.Count - 1;
        }

        foreach ( var obj in scene.ObjectsById.Values )
        {
            if ( obj.NodeType != "Geometry" )
                continue;
            var geometry = obj.Node;
            var verticesNode = geometry.Child( "Vertices" );
            var indexNode = geometry.Child( "PolygonVertexIndex" );
            if ( verticesNode is null || indexNode is null )
                continue;

            var vertices = verticesNode.AsDoubleArray( 0 );
            var polygonVertexIndex = indexNode.AsIntArray( 0 );
            if ( vertices.Length < 9 || polygonVertexIndex.Length < 3 )
                continue;
            var vertexCount = vertices.Length / 3;

            geometryToModel.TryGetValue( obj.Id, out var model );
            var world = model is not null ? WorldOf( model ) : Matrix4x4.Identity;
            Matrix4x4.Invert( world, out var inverse );
            var normalMatrix = Matrix4x4.Transpose( inverse );

            var tagName = model?.Name is { Length: > 0 } n ? n
                : obj.Name is { Length: > 0 } g ? g : $"geometry{obj.Id}";
            if ( !tagIndex.TryGetValue( tagName, out var tag ) )
            {
                tag = tags.Count;
                tags.Add( tagName );
                tagIndex.Add( tagName, tag );
            }

            var normals = ReadNormalLayer( geometry, out var normalMapping, out var normalIndices );
            var uvs = ReadUvLayer( geometry, out var uvMapping, out var uvIndices );

            // Corner dedup: position index + quantized normal/uv values.
            var corner = new Dictionary<(int V, long N, long U), int>();

            // Walk polygons: PolygonVertexIndex terminates each polygon with ~index.
            var polygon = new List<(int OutIndex, int PositionIndex)>();
            var cornerOrdinal = 0;
            for ( var i = 0; i < polygonVertexIndex.Length; i++ )
            {
                var raw = polygonVertexIndex[i];
                var positionIndex = raw < 0 ? ~raw : raw;
                if ( positionIndex < 0 || positionIndex >= vertexCount )
                    throw new FormatException(
                        $"FBX: geometry '{obj.Name}' polygon index {positionIndex} out of range (have {vertexCount} vertices)." );

                var normal = SampleLayer3( normals, normalMapping, normalIndices, positionIndex, cornerOrdinal );
                var uv = SampleLayer2( uvs, uvMapping, uvIndices, positionIndex, cornerOrdinal );
                if ( normal is not null ) anyNormal = true;
                if ( uv is not null ) anyUv = true;

                var key = (positionIndex, Quantize3( normal ), Quantize2( uv ));
                if ( !corner.TryGetValue( key, out var outIndex ) )
                {
                    outIndex = outPositions.Count;
                    vertexSource.Add( (obj.Id, positionIndex) );
                    var p = new Vector3(
                        (float)vertices[positionIndex * 3],
                        (float)vertices[positionIndex * 3 + 1],
                        (float)vertices[positionIndex * 3 + 2] );
                    outPositions.Add( Vector3.Transform( p, world ) );

                    if ( normal is { } nrm )
                    {
                        var transformed = Vector3.TransformNormal( nrm, normalMatrix );
                        var length = transformed.Length();
                        outNormals.Add( length > 1e-8f ? transformed / length : Vector3.UnitZ );
                    }
                    else
                    {
                        outNormals.Add( Vector3.UnitZ );
                    }
                    outUvs.Add( uv ?? Vector2.Zero );
                    corner.Add( key, outIndex );
                }

                polygon.Add( (outIndex, positionIndex) );
                cornerOrdinal++;

                if ( raw < 0 ) // polygon terminator: fan-triangulate
                {
                    var triangleMaterial = MaterialOf( model );
                    for ( var t = 2; t < polygon.Count; t++ )
                    {
                        outTriangles.Add( polygon[0].OutIndex );
                        outTriangles.Add( polygon[t - 1].OutIndex );
                        outTriangles.Add( polygon[t].OutIndex );
                        outTriangleTags.Add( tag );
                        outTriangleMaterials.Add( triangleMaterial );
                    }
                    polygon.Clear();
                }
            }
        }

        if ( outTriangles.Count == 0 )
            throw new FormatException( $"FBX: '{sourceName}' contains no triangle geometry." );

        var result = new RigMesh
        {
            SourceName = sourceName,
            Positions = [ .. outPositions ],
            Normals = anyNormal ? [ .. outNormals ] : [],
            Uvs = anyUv ? [ .. outUvs ] : [],
            Triangles = [ .. outTriangles ],
            TriangleTags = [ .. outTriangleTags ],
            Tags = [ .. tags ],
            Materials = materials.ToArray(),
            TriangleMaterials = [ .. outTriangleMaterials ],
        };
        result.Validate();
        return (result, vertexSource, root, scene);
    }

    enum LayerMapping { None, ByPolygonVertex, ByVertex, AllSame }

    static LayerMapping ParseMapping( FbxNode layer )
    {
        var mapping = layer.Child( "MappingInformationType" );
        if ( mapping is null || mapping.Properties.Count < 1 || mapping.Properties[0] is not string kind )
            return LayerMapping.None;
        return kind switch
        {
            "ByPolygonVertex" => LayerMapping.ByPolygonVertex,
            "ByVertice" or "ByVertex" or "ByControlPoint" => LayerMapping.ByVertex,
            "AllSame" => LayerMapping.AllSame,
            _ => LayerMapping.None, // ByPolygon etc. — not meaningful for smooth normals; skip
        };
    }

    /// <summary>IndexToDirect index array of a layer, or null for Direct reference.</summary>
    static int[]? ParseIndices( FbxNode layer, string indexNodeName )
    {
        var reference = layer.Child( "ReferenceInformationType" );
        var direct = reference is null || reference.Properties.Count < 1
            || reference.Properties[0] is not string kind || kind == "Direct";
        if ( direct )
            return null;
        var indexNode = layer.Child( indexNodeName );
        return indexNode is not null && indexNode.Properties.Count > 0 ? indexNode.AsIntArray( 0 ) : null;
    }

    static float[]? ReadNormalLayer( FbxNode geometry, out LayerMapping mapping, out int[]? indices )
    {
        mapping = LayerMapping.None;
        indices = null;
        var layer = geometry.Child( "LayerElementNormal" );
        var normalsNode = layer?.Child( "Normals" );
        if ( layer is null || normalsNode is null || normalsNode.Properties.Count < 1 )
            return null;
        mapping = ParseMapping( layer );
        indices = ParseIndices( layer, "NormalsIndex" );
        return normalsNode.AsFloatArray( 0 );
    }

    static float[]? ReadUvLayer( FbxNode geometry, out LayerMapping mapping, out int[]? indices )
    {
        mapping = LayerMapping.None;
        indices = null;
        var layer = geometry.Child( "LayerElementUV" );
        var uvNode = layer?.Child( "UV" );
        if ( layer is null || uvNode is null || uvNode.Properties.Count < 1 )
            return null;
        mapping = ParseMapping( layer );
        indices = ParseIndices( layer, "UVIndex" );
        return uvNode.AsFloatArray( 0 );
    }

    static Vector3? SampleLayer3( float[]? values, LayerMapping mapping, int[]? indices,
        int positionIndex, int cornerOrdinal )
    {
        var at = LayerElementAt( values, 3, mapping, indices, positionIndex, cornerOrdinal );
        return at < 0 ? null : new Vector3( values![at * 3], values[at * 3 + 1], values[at * 3 + 2] );
    }

    static Vector2? SampleLayer2( float[]? values, LayerMapping mapping, int[]? indices,
        int positionIndex, int cornerOrdinal )
    {
        var at = LayerElementAt( values, 2, mapping, indices, positionIndex, cornerOrdinal );
        return at < 0 ? null : new Vector2( values![at * 2], values[at * 2 + 1] );
    }

    /// <summary>Resolves the element index of a layer for one polygon corner, -1 when unavailable.</summary>
    static int LayerElementAt( float[]? values, int comps, LayerMapping mapping, int[]? indices,
        int positionIndex, int cornerOrdinal )
    {
        if ( values is null || mapping == LayerMapping.None )
            return -1;

        var direct = mapping switch
        {
            LayerMapping.ByPolygonVertex => cornerOrdinal,
            LayerMapping.ByVertex => positionIndex,
            _ => 0, // AllSame
        };
        if ( indices is not null )
        {
            if ( direct < 0 || direct >= indices.Length )
                return -1;
            direct = indices[direct];
        }
        return direct >= 0 && (direct + 1) * comps <= values.Length ? direct : -1;
    }

    static long Quantize3( Vector3? v ) => v is { } n
        ? ((long)MathF.Round( n.X * 1023f ) & 0xFFFFF) << 40
        | ((long)MathF.Round( n.Y * 1023f ) & 0xFFFFF) << 20
        | (long)MathF.Round( n.Z * 1023f ) & 0xFFFFF
        : -1L;

    static long Quantize2( Vector2? v ) => v is { } uv
        ? ((long)MathF.Round( uv.X * 65535f ) & 0xFFFFFFFF) << 32 | (long)MathF.Round( uv.Y * 65535f ) & 0xFFFFFFFF
        : -1L;
}