using System.Text;
using Sandbox.Diagnostics;

namespace LlmPoc.Llm;

public static class SboxLlmModelLoader
{
	public const string Magic = "SBOXLLM1";
	public const uint SupportedVersion = 1;
	public const byte Float32Dtype = 1;

	private static readonly byte[] MagicBytes = Encoding.ASCII.GetBytes( Magic );

	public static SboxLlmModel LoadFromMounted( string path )
	{
		if ( string.IsNullOrWhiteSpace( path ) )
		{
			throw new ArgumentException( "[LLM:ERROR] Model path cannot be empty.", nameof( path ) );
		}

		if ( !FileSystem.Mounted.FileExists( path ) )
		{
			throw new LlmModelFormatException(
				$"Model file '{path}' is not present in FileSystem.Mounted. " +
				"Confirm it is packaged under Assets using the same relative path." );
		}

		FastTimer timer = FastTimer.StartNew();
		LlmLog.Info( "LOAD", $"Reading '{path}' from FileSystem.Mounted." );
		byte[] bytes = FileSystem.Mounted.ReadAllBytes( path ).ToArray();
		LlmLog.Info( "LOAD", $"Read {bytes.LongLength:N0} bytes in {timer.ElapsedMilliSeconds:N2} ms." );

		timer.Start();
		SboxLlmModel model = Parse( bytes );
		LlmLog.Info(
			"LOAD",
			$"Parsed format v{model.FormatVersion}, {model.TensorCount:N0} tensors, " +
			$"{model.TotalParameterCount:N0} parameters ({model.TotalTensorBytes:N0} bytes) " +
			$"in {timer.ElapsedMilliSeconds:N2} ms." );
		return model;
	}

	public static SboxLlmModel Parse( byte[] bytes )
	{
		if ( bytes is null )
		{
			throw new ArgumentNullException( nameof( bytes ) );
		}

		Cursor cursor = new( bytes );
		for ( int index = 0; index < MagicBytes.Length; index++ )
		{
			byte actual = cursor.ReadByte( "header magic" );
			if ( actual != MagicBytes[index] )
			{
				throw new LlmModelFormatException(
					$"Invalid header magic at byte {index}: expected 0x{MagicBytes[index]:X2}, " +
					$"found 0x{actual:X2}. Expected ASCII '{Magic}'." );
			}
		}

		uint version = cursor.ReadUInt32( "format version" );
		if ( version != SupportedVersion )
		{
			throw new LlmModelFormatException(
				$"Unsupported SBOXLLM format version {version}; expected {SupportedVersion}." );
		}

		uint tensorCount = cursor.ReadUInt32( "tensor count" );
		if ( tensorCount > 1_000_000 )
		{
			throw new LlmModelFormatException(
				$"Unreasonable tensor count {tensorCount:N0} at byte 12." );
		}

		Dictionary<string, Tensor> tensors = new( (int)tensorCount, StringComparer.Ordinal );
		for ( uint tensorIndex = 0; tensorIndex < tensorCount; tensorIndex++ )
		{
			int metadataOffset = cursor.Offset;
			uint nameByteCount = cursor.ReadUInt32( $"tensor {tensorIndex} name length" );
			if ( nameByteCount == 0 || nameByteCount > 16_384 )
			{
				throw new LlmModelFormatException(
					$"Tensor {tensorIndex} has invalid UTF-8 name length {nameByteCount} " +
					$"at offset {metadataOffset}." );
			}

			string name = cursor.ReadUtf8( (int)nameByteCount, $"tensor {tensorIndex} name" );
			if ( tensors.ContainsKey( name ) )
			{
				throw new LlmModelFormatException(
					$"Tensor {tensorIndex} duplicates name '{name}'." );
			}

			byte dtype = cursor.ReadByte( $"{name} dtype" );
			if ( dtype != Float32Dtype )
			{
				throw new LlmModelFormatException(
					$"{name} expected dtype id {Float32Dtype} (little-endian FP32), found {dtype}." );
			}

			byte rank = cursor.ReadByte( $"{name} rank" );
			if ( rank > 16 )
			{
				throw new LlmModelFormatException( $"{name} has unreasonable rank {rank}." );
			}

			ushort reserved = cursor.ReadUInt16( $"{name} reserved field" );
			if ( reserved != 0 )
			{
				throw new LlmModelFormatException(
					$"{name} reserved field must be zero, found {reserved}." );
			}

			int[] shape = new int[rank];
			long elementCount = 1;
			for ( int axis = 0; axis < rank; axis++ )
			{
				uint dimension = cursor.ReadUInt32( $"{name} dimension {axis}" );
				if ( dimension == 0 || dimension > int.MaxValue )
				{
					throw new LlmModelFormatException(
						$"{name} dimension {axis} must be in [1,{int.MaxValue}], found {dimension}." );
				}

				shape[axis] = (int)dimension;
				if ( elementCount > long.MaxValue / dimension )
				{
					throw new LlmModelFormatException(
						$"{name} shape [{string.Join( ",", shape )}] overflows the signed " +
						"64-bit element count." );
				}
				elementCount *= dimension;
			}

			ulong byteCount = cursor.ReadUInt64( $"{name} byte count" );
			if ( (ulong)elementCount > ulong.MaxValue / sizeof( float ) )
			{
				throw new LlmModelFormatException(
					$"{name} FP32 byte count overflows unsigned 64-bit." );
			}
			ulong expectedBytes = (ulong)elementCount * sizeof( float );
			if ( byteCount != expectedBytes )
			{
				throw new LlmModelFormatException(
					$"{name} shape [{string.Join( ",", shape )}] has {elementCount:N0} elements " +
					$"and requires {expectedBytes:N0} FP32 bytes, found {byteCount:N0}." );
			}

			if ( elementCount > int.MaxValue )
			{
				throw new LlmModelFormatException(
					$"{name} has {elementCount:N0} elements, exceeding the managed array limit." );
			}

			float[] data = new float[(int)elementCount];
			for ( int element = 0; element < data.Length; element++ )
			{
				data[element] = cursor.ReadSingle( $"{name} element {element}" );
			}

			Tensor tensor = new( name, shape, data );
			tensors.Add( name, tensor );
			LlmLog.Trace(
				"TENSOR",
				$"Loaded {name} shape={tensor.ShapeText} elements={tensor.ElementCount:N0}." );
		}

		if ( cursor.Remaining != 0 )
		{
			throw new LlmModelFormatException(
				$"Unexpected trailing data: {cursor.Remaining:N0} bytes remain after " +
				$"{tensorCount:N0} tensors at offset {cursor.Offset:N0}." );
		}

		return new SboxLlmModel( (int)version, tensors );
	}

	private sealed class Cursor
	{
		private readonly byte[] _bytes;

		public int Offset { get; private set; }
		public int Remaining => _bytes.Length - Offset;

		public Cursor( byte[] bytes )
		{
			_bytes = bytes;
		}

		public byte ReadByte( string label )
		{
			Require( 1, label );
			return _bytes[Offset++];
		}

		public ushort ReadUInt16( string label )
		{
			Require( 2, label );
			int start = Offset;
			Offset += 2;
			return (ushort)(_bytes[start] | (_bytes[start + 1] << 8));
		}

		public uint ReadUInt32( string label )
		{
			Require( 4, label );
			int start = Offset;
			Offset += 4;
			return (uint)(
				_bytes[start]
				| (_bytes[start + 1] << 8)
				| (_bytes[start + 2] << 16)
				| (_bytes[start + 3] << 24) );
		}

		public ulong ReadUInt64( string label )
		{
			uint low = ReadUInt32( label );
			uint high = ReadUInt32( label );
			return low | ((ulong)high << 32);
		}

		public float ReadSingle( string label )
		{
			uint bits = ReadUInt32( label );
			return BitConverter.Int32BitsToSingle( unchecked( (int)bits ) );
		}

		public string ReadUtf8( int byteCount, string label )
		{
			Require( (ulong)byteCount, label );
			int start = Offset;
			Offset += byteCount;
			try
			{
				return Encoding.UTF8.GetString( _bytes, start, byteCount );
			}
			catch ( DecoderFallbackException error )
			{
				throw new LlmModelFormatException(
					$"{label} is not valid UTF-8 at file offset {start}.", error );
			}
		}

		private void Require( ulong byteCount, string label )
		{
			if ( byteCount > (ulong)Remaining )
			{
				throw new LlmModelFormatException(
					$"Truncated {label}: requested {byteCount:N0} bytes at offset {Offset:N0}, " +
					$"but only {Remaining:N0} remain." );
			}
		}
	}
}