using System;
using System.Collections.Generic;
using System.Security.Cryptography;
using System.Text;

namespace GlinerPoc.Packaging;

/// <summary>
/// C# parser/validator for the SBGLI1 binary manifest produced by the
/// deterministic Python exporter (format version 1). Mirrors the on-disk
/// layout exactly; every failure throws with an actionable message.
///
/// Layout (little-endian):
///   header 128 B: magic(8) ver flags headerBytes tensorCount chunkCount
///                 shardCount dtype reserved ckptSha(32) rcSha(32) zero(24)
///   tensor records, chunk records, shard records, recordsSha(32)
/// </summary>
public sealed class Sbgli1Manifest
{
	public const string Magic = "SBGLI1";
	public const int FormatVersion = 1;
	public const int HeaderBytes = 128;
	public const int StorageDtypeF32 = 1;

	public sealed class TensorRecord
	{
		public string Name;
		public int[] Shape;
		public int Rank;
		public long ElementCount;
		public long ByteCount;
		public int FirstChunkIndex;
		public int ChunkCount;
	}

	public sealed class ChunkRecord
	{
		public int TensorIndex;
		public int ChunkIndex;
		public long ElemOffset;
		public long ElemCount;
		public int RowStart;
		public int RowCount;
		public int ShardIndex;
		public long ByteOffsetInShard;
		public string Sha256;
	}

	public sealed class ShardRecord
	{
		public int Index;
		public long ByteCount;
		public string Sha256;
	}

	public int FormatVersionField;
	public int TensorCount;
	public int ChunkCount;
	public int ShardCount;
	public string CheckpointSha256;
	public string RuntimeConfigSha256;

	public List<TensorRecord> Tensors = new();
	public List<ChunkRecord> Chunks = new();
	public List<ShardRecord> Shards = new();

	private readonly Dictionary<string, int> _tensorIndexByName = new( StringComparer.Ordinal );

	public static Sbgli1Manifest Parse( byte[] data )
	{
		var m = new Sbgli1Manifest();
		int pos = 0;

		string magic = Encoding.ASCII.GetString( data, 0, 8 ).TrimEnd( '\0' );
		if ( magic != Magic )
			throw new InvalidOperationException( $"[GLI:ERROR] SBGLI1 magic expected '{Magic}', found '{magic}'." );
		pos = 8;

		m.FormatVersionField = (int)ReadU32( data, ref pos );
		int flags = (int)ReadU32( data, ref pos );
		int headerBytes = (int)ReadU32( data, ref pos );
		if ( m.FormatVersionField != FormatVersion )
			throw new InvalidOperationException( $"[GLI:ERROR] SBGLI1 format version expected {FormatVersion}, found {m.FormatVersionField}." );
		if ( (flags & 1) == 0 )
			throw new InvalidOperationException( "[GLI:ERROR] SBGLI1 big-endian data is unsupported." );
		if ( headerBytes != HeaderBytes )
			throw new InvalidOperationException( $"[GLI:ERROR] SBGLI1 header size expected {HeaderBytes}, found {headerBytes}." );

		m.TensorCount = (int)ReadU32( data, ref pos );
		m.ChunkCount = (int)ReadU32( data, ref pos );
		m.ShardCount = (int)ReadU32( data, ref pos );
		int dtype = (int)ReadU32( data, ref pos );
		int reserved = (int)ReadU32( data, ref pos );
		if ( dtype != StorageDtypeF32 )
			throw new InvalidOperationException( $"[GLI:ERROR] SBGLI1 storage dtype expected {StorageDtypeF32} (FP32), found {dtype}." );
		if ( reserved != 0 )
			throw new InvalidOperationException( "[GLI:ERROR] SBGLI1 reserved header field must be zero." );

		m.CheckpointSha256 = Hex( data, pos, 32 );
		pos += 32;
		m.RuntimeConfigSha256 = Hex( data, pos, 32 );
		pos += 32;
		for ( int i = 0; i < 24; i++ )
		{
			if ( data[pos + i] != 0 )
				throw new InvalidOperationException( "[GLI:ERROR] SBGLI1 reserved header tail must be zero." );
		}
		pos += 24;
		if ( pos != HeaderBytes )
			throw new InvalidOperationException( "[GLI:ERROR] SBGLI1 header size mismatch." );

		int recordsStart = pos;

		for ( int i = 0; i < m.TensorCount; i++ )
		{
			var t = new TensorRecord();
			int nameLen = data[pos];
			pos += 1;
			t.Name = Encoding.UTF8.GetString( data, pos, nameLen );
			pos += nameLen;
			t.Rank = data[pos];
			pos += 1;
			t.Shape = new int[t.Rank];
			for ( int d = 0; d < t.Rank; d++ )
			{
				t.Shape[d] = (int)ReadU32( data, ref pos );
			}
			t.ElementCount = (long)ReadU64( data, ref pos );
			t.ByteCount = (long)ReadU64( data, ref pos );
			t.FirstChunkIndex = (int)ReadU32( data, ref pos );
			t.ChunkCount = (int)ReadU32( data, ref pos );

			if ( t.Rank <= 0 || t.Rank > 8 )
				throw new InvalidOperationException( $"[GLI:ERROR] Tensor '{t.Name}' has invalid rank {t.Rank}." );
			if ( t.ElementCount <= 0 || t.ByteCount != t.ElementCount * 4 )
				throw new InvalidOperationException( $"[GLI:ERROR] Tensor '{t.Name}' element/byte count mismatch ({t.ElementCount} / {t.ByteCount})." );
			if ( m._tensorIndexByName.ContainsKey( t.Name ) )
				throw new InvalidOperationException( $"[GLI:ERROR] Duplicate tensor record '{t.Name}'." );
			m._tensorIndexByName[t.Name] = i;
			m.Tensors.Add( t );
		}

		for ( int i = 0; i < m.ChunkCount; i++ )
		{
			var c = new ChunkRecord
			{
				TensorIndex = (int)ReadU32( data, ref pos ),
				ChunkIndex = (int)ReadU32( data, ref pos ),
				ElemOffset = (long)ReadU64( data, ref pos ),
				ElemCount = (long)ReadU64( data, ref pos ),
				RowStart = (int)ReadU32( data, ref pos ),
				RowCount = (int)ReadU32( data, ref pos ),
				ShardIndex = (int)ReadU32( data, ref pos ),
				ByteOffsetInShard = (long)ReadU64( data, ref pos ),
				Sha256 = Hex( data, pos, 32 ),
			};
			pos += 32;
			if ( c.TensorIndex < 0 || c.TensorIndex >= m.TensorCount )
				throw new InvalidOperationException( $"[GLI:ERROR] Chunk {i} references tensor {c.TensorIndex} out of bounds." );
			if ( c.ShardIndex < 0 || c.ShardIndex >= m.ShardCount )
				throw new InvalidOperationException( $"[GLI:ERROR] Chunk {i} references shard {c.ShardIndex} out of bounds." );
			var owner = m.Tensors[c.TensorIndex];
			if ( c.ElemOffset < 0 || c.ElemOffset + c.ElemCount > owner.ElementCount )
				throw new InvalidOperationException( $"[GLI:ERROR] Chunk {i} element range outside tensor '{owner.Name}'." );
			m.Chunks.Add( c );
		}

		for ( int i = 0; i < m.ShardCount; i++ )
		{
			var s = new ShardRecord
			{
				Index = (int)ReadU32( data, ref pos ),
				ByteCount = (long)ReadU64( data, ref pos ),
				Sha256 = Hex( data, pos, 32 ),
			};
			pos += 32;
			m.Shards.Add( s );
		}

		int recordsEnd = pos;
		string recordsSha = Convert.ToHexString( SHA256.HashData(
			data[recordsStart..recordsEnd] ) ).ToLowerInvariant();
		string storedRecordsSha = Hex( data, pos, 32 );
		pos += 32;
		if ( recordsSha != storedRecordsSha )
			throw new InvalidOperationException(
				$"[GLI:ERROR] SBGLI1 records region SHA-256 mismatch (expected {storedRecordsSha}, computed {recordsSha})." );
		if ( pos != data.Length )
			throw new InvalidOperationException( $"[GLI:ERROR] SBGLI1 manifest has {data.Length - pos} trailing bytes." );
		return m;
	}

	public TensorRecord GetRequiredTensor( string name )
	{
		if ( !_tensorIndexByName.TryGetValue( name, out int index ) )
			throw new InvalidOperationException( $"[GLI:ERROR] Required tensor '{name}' is missing from the SBGLI1 manifest." );
		return Tensors[index];
	}

	/// <summary>
	/// Copies the exact bytes of one tensor from the provided shard payloads.
	/// shardData must be indexed by shard record index. Every chunk is
	/// SHA-256-verified before copying.
	/// </summary>
	public byte[] ReadTensorBytes( string name, IReadOnlyList<byte[]> shardData )
	{
		var t = GetRequiredTensor( name );
		var result = new byte[t.ByteCount];
		int written = 0;
		for ( int ci = t.FirstChunkIndex; ci < t.FirstChunkIndex + t.ChunkCount; ci++ )
		{
			var c = Chunks[ci];
			byte[] shard = shardData[c.ShardIndex];
			long off = c.ByteOffsetInShard;
			int len = (int)(c.ElemCount * 4);
			if ( off < 0 || off + len > shard.Length )
				throw new InvalidOperationException(
					$"[GLI:ERROR] Chunk {ci} of '{name}' reads [{off}, {off + len}) outside shard {c.ShardIndex} ({shard.Length} B)." );
			byte[] payload = shard[(int)off..(int)(off + len)];
			string sha = Convert.ToHexString( SHA256.HashData( payload ) ).ToLowerInvariant();
			if ( sha != c.Sha256 )
				throw new InvalidOperationException( $"[GLI:ERROR] Chunk {ci} of '{name}' SHA-256 mismatch." );
			Buffer.BlockCopy( payload, 0, result, written, len );
			written += len;
		}
		if ( written != t.ByteCount )
			throw new InvalidOperationException( $"[GLI:ERROR] Tensor '{name}' reconstructed {written} of {t.ByteCount} bytes." );
		return result;
	}

	private static ulong ReadU32( byte[] data, ref int pos )
	{
		ulong v = BitConverter.ToUInt32( data, pos );
		pos += 4;
		return v;
	}

	private static ulong ReadU64( byte[] data, ref int pos )
	{
		ulong v = BitConverter.ToUInt64( data, pos );
		pos += 8;
		return v;
	}

	private static string Hex( byte[] data, int offset, int count )
	{
		return Convert.ToHexString( data, offset, count ).ToLowerInvariant();
	}
}