#nullable enable
using System;
using System.Collections.Generic;
using System.IO;
using System.IO.Compression;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;

namespace HumanoidMocapX.Inference;

/// <summary>Reads tensor data from the pinned PyTorch ZIP checkpoint format without Python,
/// importing modules, invoking pickle callables, or deserializing executable objects.</summary>
public sealed partial class TorchCheckpoint : IDisposable
{
    public sealed record TensorInfo(string Name,string Storage,string Dtype,long StorageLength,long Offset,int[] Shape,long[] Stride);
    sealed record Symbol(string Module,string Name);
    sealed record StorageRef(string Key,string Dtype,long Count);
    sealed record TensorRef(StorageRef Storage,long Offset,int[] Shape,long[] Stride);
    sealed class ScriptObject(Symbol type)
    {
        internal Symbol Type { get; }=type;
        internal Dictionary<object,object?>? Fields;
    }
    readonly ZipArchive archive;
    readonly string prefix;
    readonly Dictionary<string,TensorInfo> tensors=new(StringComparer.Ordinal);
    public IReadOnlyDictionary<string,TensorInfo> Tensors=>tensors;
    const int MaximumMetadataBytes=32*1024*1024;
    const int StorageReadBytes=64*1024;
    const long MaximumTensorElements=1_000_000_000;

    public TorchCheckpoint(string path,string? stateKey=null)
    {
        archive=ZipFile.OpenRead(path);
        try
        {
            var root=ReadMetadata(archive,out prefix) as Dictionary<object,object?> ?? throw new InvalidDataException("Checkpoint root must be a dictionary.");
            // Some inference checkpoints (including ACE) keep separately named
            // component dictionaries. Select one explicitly; never traverse optimizer state.
            var key=stateKey??(root.ContainsKey("state_dict")?"state_dict":null);
            var state=key is null?root:root.TryGetValue(key,out var value)?value as Dictionary<object,object?>:null;
            if(state is null)throw new InvalidDataException($"Checkpoint component '{key}' is missing or is not a dictionary.");
            foreach(var pair in state)
            {
                if(pair.Value is not TensorRef t)continue;
                if(pair.Key is not string name)throw new InvalidDataException("Tensor names must be strings.");
                var info=new TensorInfo(name,t.Storage.Key,t.Storage.Dtype,t.Storage.Count,t.Offset,t.Shape,t.Stride);
                Validate(info);tensors.Add(name,info);
            }
            if(tensors.Count==0)throw new InvalidDataException("No tensor weights in checkpoint.");
        }
        catch{archive.Dispose();throw;}
    }
    static object? ReadMetadata(ZipArchive archive,out string prefix,bool scriptObjects=false)
    {
        var entries=archive.Entries.Where(e=>e.FullName.EndsWith("/data.pkl",StringComparison.Ordinal)).ToArray();
        if(entries.Length!=1)throw new InvalidDataException("Expected exactly one PyTorch data.pkl entry.");
        var entry=entries[0];prefix=entry.FullName.Substring(0,entry.FullName.Length-8);
        if(entry.Length>MaximumMetadataBytes)throw new InvalidDataException("Checkpoint metadata exceeds the 32 MiB limit.");
        if(archive.GetEntry(prefix+"byteorder") is { } byteorder)
        {using var r=new StreamReader(byteorder.Open());if(r.ReadToEnd().Trim()!="little")throw new NotSupportedException("Only little-endian checkpoints are supported.");}
        using var source=entry.Open();using var memory=new MemoryStream();source.CopyTo(memory);
        return new DataReader(memory.ToArray(),scriptObjects:scriptObjects).Read();
    }
    static long Count(TensorInfo tensor)
    {long count=1;foreach(var d in tensor.Shape){if(d<0)throw new InvalidDataException("Negative tensor dimension.");count=checked(count*d);if(count>MaximumTensorElements)throw new InvalidDataException("Tensor exceeds element limit.");}return count;}
    void Validate(TensorInfo tensor)
    {
        if(tensor.Shape.Length>8||tensor.Shape.Length!=tensor.Stride.Length||tensor.Offset<0||tensor.StorageLength<0||tensor.StorageLength>MaximumTensorElements)
            throw new InvalidDataException("Invalid tensor dimensions/storage.");
        if(tensor.Storage.Length==0||tensor.Storage.Any(c=>!char.IsAsciiDigit(c)))throw new InvalidDataException("Unexpected tensor storage key.");
        var count=Count(tensor);long last=tensor.Offset;
        for(var i=0;i<tensor.Shape.Length;i++)
        {if(tensor.Stride[i]<0)throw new InvalidDataException("Negative tensor stride.");last=checked(last+Math.Max(0,tensor.Shape[i]-1)*tensor.Stride[i]);}
        if(count>0&&last>=tensor.StorageLength)throw new InvalidDataException("Tensor view exceeds its storage.");
        var storage=archive.GetEntry(prefix+"data/"+tensor.Storage)??throw new InvalidDataException("Missing tensor storage.");
        if(storage.Length!=checked(tensor.StorageLength*ElementSize(tensor.Dtype)))throw new InvalidDataException("Tensor storage byte length mismatch.");
    }
    static int ElementSize(string dtype)=>dtype switch
    {"FloatStorage" or "IntStorage" or "UInt32Storage"=>4,"DoubleStorage" or "LongStorage"=>8,"HalfStorage" or "BFloat16Storage"=>2,"ByteStorage" or "BoolStorage"=>1,_=>throw new NotSupportedException("Unsupported checkpoint dtype: "+dtype)};

    /// <summary>Materializes one tensor, respecting noncontiguous views. Never loads the whole
    /// checkpoint into memory. Callers choose which model's tensors to retain.</summary>
    public float[] ReadFloat(string name,CancellationToken cancellation=default)
    {
        var tensor=tensors.TryGetValue(name,out var found)?found:throw new KeyNotFoundException(name);
        var count=checked((int)Count(tensor));var elementSize=ElementSize(tensor.Dtype);
        long contiguousStride=1;var contiguous=true;
        for(var dim=tensor.Shape.Length-1;dim>=0;dim--)
        {if(tensor.Shape[dim]>1&&tensor.Stride[dim]!=contiguousStride)contiguous=false;contiguousStride*=tensor.Shape[dim];}
        if(contiguous&&tensor.Dtype=="FloatStorage"&&BitConverter.IsLittleEndian
            &&tensor.Offset==0&&count==tensor.StorageLength&&count<=int.MaxValue/sizeof(float))
        {
            // Whole FP32 storage already has the output layout; avoid a second model-sized allocation.
            cancellation.ThrowIfCancellationRequested();var direct=new float[count];
            ReadStorage(tensor,MemoryMarshal.AsBytes(direct.AsSpan()),cancellation);
            cancellation.ThrowIfCancellationRequested();return direct;
        }
        var storage=ReadStorage(tensor,cancellation);var output=new float[count];
        if(contiguous&&tensor.Dtype=="FloatStorage"&&BitConverter.IsLittleEndian)
        {
            cancellation.ThrowIfCancellationRequested();
            Buffer.BlockCopy(storage,checked((int)(tensor.Offset*4)),output,0,checked(count*4));
            cancellation.ThrowIfCancellationRequested();return output;
        }
        for(var i=0;i<count;i++)
        {
            if((i&16383)==0)cancellation.ThrowIfCancellationRequested();
            var at=checked((int)(StorageIndex(tensor,i)*elementSize));
            output[i]=tensor.Dtype switch
            {
                "FloatStorage"=>BitConverter.ToSingle(storage,at),"DoubleStorage"=>(float)BitConverter.ToDouble(storage,at),
                "HalfStorage"=>(float)BitConverter.UInt16BitsToHalf(BitConverter.ToUInt16(storage,at)),
                "BFloat16Storage"=>BitConverter.Int32BitsToSingle(BitConverter.ToUInt16(storage,at)<<16),
                "IntStorage"=>BitConverter.ToInt32(storage,at),"LongStorage"=>BitConverter.ToInt64(storage,at),
                "UInt32Storage"=>BitConverter.ToUInt32(storage,at),
                "ByteStorage" or "BoolStorage"=>storage[at],_=>throw new NotSupportedException(tensor.Dtype)
            };
        }
        return output;
    }

    /// <summary>Reads integer tensor values exactly, including unsigned skeleton indices.
    /// Rejects floating-point storage; missing names and cancellation propagate.</summary>
    public long[] ReadInt64(string name,CancellationToken cancellation=default)
    {
        var tensor=tensors.TryGetValue(name,out var found)?found:throw new KeyNotFoundException(name);
        if(tensor.Dtype is not ("ByteStorage" or "BoolStorage" or "IntStorage" or "UInt32Storage" or "LongStorage"))
            throw new InvalidDataException("Expected integer tensor storage: "+name);
        var storage=ReadStorage(tensor,cancellation);var size=ElementSize(tensor.Dtype);
        var output=new long[checked((int)Count(tensor))];
        for(var i=0;i<output.Length;i++)
        {
            if((i&16383)==0)cancellation.ThrowIfCancellationRequested();
            var at=checked((int)(StorageIndex(tensor,i)*size));
            output[i]=tensor.Dtype switch
            {
                "IntStorage"=>BitConverter.ToInt32(storage,at),"UInt32Storage"=>BitConverter.ToUInt32(storage,at),
                "LongStorage"=>BitConverter.ToInt64(storage,at),_=>storage[at]
            };
        }
        return output;
    }

    byte[] ReadStorage(TensorInfo tensor,CancellationToken cancellation)
    {
        cancellation.ThrowIfCancellationRequested();
        var storage=new byte[checked((int)(tensor.StorageLength*ElementSize(tensor.Dtype)))];
        ReadStorage(tensor,storage,cancellation);return storage;
    }

    void ReadStorage(TensorInfo tensor,Span<byte> storage,CancellationToken cancellation)
    {
        cancellation.ThrowIfCancellationRequested();
        using var input=archive.GetEntry(prefix+"data/"+tensor.Storage)!.Open();
        var offset=0;
        while(offset<storage.Length)
        {cancellation.ThrowIfCancellationRequested();var n=input.Read(storage.Slice(offset,Math.Min(StorageReadBytes,storage.Length-offset)));if(n==0)throw new EndOfStreamException();offset+=n;}
    }

    static long StorageIndex(TensorInfo tensor,int index)
    {
        long result=tensor.Offset;
        for(var dimension=tensor.Shape.Length-1;dimension>=0;dimension--)
        {result+=(index%tensor.Shape[dimension])*tensor.Stride[dimension];index/=tensor.Shape[dimension];}
        return result;
    }
    public void Dispose()=>archive.Dispose();

    sealed class DataReader
    {
        readonly BinaryReader reader;
        readonly List<object?> stack=new();
        readonly Dictionary<int,object?> memo=new();
        readonly bool scriptObjects;
        static readonly object Mark=new();
        int operations;
        public int Position=>checked((int)reader.BaseStream.Position);
        public DataReader(byte[] bytes,int offset=0,bool scriptObjects=false){reader=new BinaryReader(new MemoryStream(bytes),Encoding.UTF8);reader.BaseStream.Position=offset;this.scriptObjects=scriptObjects;}
        object? Pop(){if(stack.Count==0)throw new InvalidDataException("Empty pickle stack.");var value=stack[^1];stack.RemoveAt(stack.Count-1);return value;}
        object? Peek()=>stack.Count==0?throw new InvalidDataException("Empty pickle stack."):stack[^1];
        object?[] MarkItems()
        {
            var index=stack.LastIndexOf(Mark);if(index<0)throw new InvalidDataException("Missing pickle mark.");
            var values=stack.Skip(index+1).ToArray();stack.RemoveRange(index,stack.Count-index);return values;
        }
        string Line()
        {var bytes=new List<byte>();byte b;while((b=reader.ReadByte())!=10){if(bytes.Count>=4096)throw new InvalidDataException("Pickle identifier too long.");bytes.Add(b);}return Encoding.UTF8.GetString(bytes.ToArray());}
        string Text(int length)
        {if(length<0||length>MaximumMetadataBytes||length>reader.BaseStream.Length-reader.BaseStream.Position)throw new InvalidDataException("Invalid pickle string length.");return Encoding.UTF8.GetString(reader.ReadBytes(length));}
        static long Integer(object? value)=>value switch{int i=>i,long l=>l,_=>throw new InvalidDataException("Expected integer.")};
        static object?[] Tuple(object? value)=>value as object?[]??throw new InvalidDataException("Expected tuple.");
        static Dictionary<object,object?> Dict(object? value)=>value as Dictionary<object,object?>??throw new InvalidDataException("Expected dictionary.");
        static void Set(Dictionary<object,object?> dictionary,object? key,object? value)
        {
            // Training checkpoints contain integer-keyed optimizer states. Keep those
            // as inert data, with numeric keys distinct from strings such as "0".
            var normalized=key switch{string s=>(object)s,int i=>(long)i,long l=>l,_=>throw new InvalidDataException("Only string and integer dictionary keys are supported.")};
            dictionary[normalized]=value;
        }
        public object? Read(bool requireEnd=true)
        {
            while(reader.BaseStream.Position<reader.BaseStream.Length)
            {
                if(++operations>2_000_000||stack.Count>100_000||memo.Count>100_000)throw new InvalidDataException("Checkpoint metadata complexity limit.");
                var op=reader.ReadByte();
                switch(op)
                {
                    case 0x80:var protocol=reader.ReadByte();if(protocol>5)throw new NotSupportedException("Pickle protocol "+protocol);break;
                    case (byte)'.':if(stack.Count!=1||(requireEnd&&reader.BaseStream.Position!=reader.BaseStream.Length))throw new InvalidDataException("Invalid pickle termination.");return Pop();
                    case (byte)'(':stack.Add(Mark);break;
                    case (byte)')':stack.Add(System.Array.Empty<object?>());break;
                    case (byte)'}':stack.Add(new Dictionary<object,object?>());break;
                    case (byte)']':stack.Add(new List<object?>());break;
                    case (byte)'N':stack.Add(null);break;
                    case 0x88:stack.Add(true);break;
                    case 0x89:stack.Add(false);break;
                    case (byte)'K':stack.Add((int)reader.ReadByte());break;
                    case (byte)'M':stack.Add((int)reader.ReadUInt16());break;
                    case (byte)'J':stack.Add(reader.ReadInt32());break;
                    case (byte)'G':var f=reader.ReadBytes(8);if(f.Length!=8)throw new EndOfStreamException();System.Array.Reverse(f);stack.Add(BitConverter.ToDouble(f));break;
                    case (byte)'X':stack.Add(Text(reader.ReadInt32()));break;
                    case 0x8c:stack.Add(Text(reader.ReadByte()));break;
                    case 0x95:var frame=reader.ReadUInt64();if(frame>(ulong)(reader.BaseStream.Length-reader.BaseStream.Position))throw new InvalidDataException("Truncated pickle frame.");break;
                    case (byte)'q':memo[reader.ReadByte()]=Peek();break;
                    case (byte)'r':memo[reader.ReadInt32()]=Peek();break;
                    case 0x94:memo[memo.Count]=Peek();break;
                    case (byte)'h':stack.Add(memo[reader.ReadByte()]);break;
                    case (byte)'j':stack.Add(memo[reader.ReadInt32()]);break;
                    case (byte)'t':stack.Add(MarkItems());break;
                    case 0x85:stack.Add(new[]{Pop()});break;
                    case 0x86:var second=Pop();stack.Add(new[]{Pop(),second});break;
                    case 0x87:var third=Pop();second=Pop();stack.Add(new[]{Pop(),second,third});break;
                    case 0x81:
                        var constructorArguments=Tuple(Pop());var objectType=Pop() as Symbol;
                        if(!scriptObjects||constructorArguments.Length!=0||objectType is null||!objectType.Module.StartsWith("__torch__",StringComparison.Ordinal))
                            throw new InvalidDataException("Unsupported checkpoint object construction.");
                        stack.Add(new ScriptObject(objectType));break;
                    case (byte)'s':var value=Pop();var key=Pop();Set(Dict(Peek()),key,value);break;
                    case (byte)'u':var pairs=MarkItems();if(pairs.Length%2!=0)throw new InvalidDataException("Odd dictionary entries.");var d=Dict(Peek());for(var i=0;i<pairs.Length;i+=2)Set(d,pairs[i],pairs[i+1]);break;
                    case (byte)'a':value=Pop();((List<object?>)Peek()!).Add(value);break;
                    case (byte)'e':var items=MarkItems();((List<object?>)Peek()!).AddRange(items);break;
                    case (byte)'c':stack.Add(AllowedSymbol(Line(),Line()));break;
                    case 0x93:var name=Pop() as string??throw new InvalidDataException();var module=Pop() as string??throw new InvalidDataException();stack.Add(AllowedSymbol(module,name));break;
                    case (byte)'Q':
                        var storage=Tuple(Pop());
                        if(storage.Length<5||storage[0] as string!="storage"||storage[1] is not Symbol type||type.Module!="torch"||storage[2] is not string storageKey)
                            throw new InvalidDataException("Unsupported persistent pickle reference.");
                        if(storage.Length>5&&storage[5] is not null)throw new NotSupportedException("Legacy storage aliases are not supported.");
                        stack.Add(new StorageRef(storageKey,type.Name,Integer(storage[4])));break;
                    case (byte)'R':
                        var arguments=Tuple(Pop());var symbol=Pop() as Symbol??throw new InvalidDataException("Unsupported pickle callable.");
                        if(symbol==new Symbol("collections","OrderedDict")&&arguments.Length==0)stack.Add(new Dictionary<object,object?>());
                        else if(symbol==new Symbol("collections","Counter")&&arguments.Length<=1)
                            stack.Add(arguments.Length==0?new Dictionary<object,object?>():new Dictionary<object,object?>(Dict(arguments[0])));
                        else if(symbol.Module=="torch._utils"&&(symbol.Name=="_rebuild_tensor_v2"||symbol.Name=="_rebuild_tensor")&&arguments.Length>=4&&arguments[0] is StorageRef sr)
                            stack.Add(new TensorRef(sr,Integer(arguments[1]),Tuple(arguments[2]).Select(x=>checked((int)Integer(x))).ToArray(),Tuple(arguments[3]).Select(Integer).ToArray()));
                        else if(scriptObjects&&symbol==new Symbol("torch.jit._pickle","restore_type_tag")&&arguments.Length==2
                            &&arguments[1] as string=="List[str]"&&arguments[0] is List<object?> strings&&strings.All(x=>x is string))
                            stack.Add(strings);
                        else if(scriptObjects&&symbol==new Symbol("torch.jit._pickle","build_intlist")&&arguments.Length==1
                            &&arguments[0] is List<object?> integers&&integers.All(x=>x is int or long))stack.Add(integers);
                        else throw new InvalidDataException("Unsupported checkpoint construction: "+symbol);
                        break;
                    case (byte)'b':
                        var metadata=Dict(Pop());
                        if(scriptObjects&&Peek() is ScriptObject instance)
                        {
                            if(instance.Fields is not null)throw new InvalidDataException("Repeated checkpoint object state.");
                            instance.Fields=metadata;break;
                        }
                        _ = Dict(Peek());
                        if(metadata.Keys.Any(k=>!Equals(k,"_metadata")))throw new InvalidDataException("Unsupported checkpoint object state.");
                        break;
                    default:throw new NotSupportedException($"Checkpoint pickle opcode 0x{op:X2} at {reader.BaseStream.Position-1} is unsupported.");
                }
            }
            throw new EndOfStreamException("Missing pickle STOP.");
        }
        Symbol AllowedSymbol(string module,string name)
        {
            if(scriptObjects&&(module=="__torch__"||module.StartsWith("__torch__.",StringComparison.Ordinal)
                ||module=="torch.jit._pickle"&&name is "restore_type_tag" or "build_intlist"))return new(module,name);
            if(module=="collections"&&(name=="OrderedDict"||name=="Counter")||module=="torch._utils"&&(name=="_rebuild_tensor_v2"||name=="_rebuild_tensor"))return new(module,name);
            if(module=="torch"){_ = ElementSize(name);return new(module,name);}
            throw new InvalidDataException("Executable/unsupported checkpoint symbol rejected: "+module+"."+name);
        }
    }
}