using System.Security.Cryptography;
using System.Text;
using Sandbox.Diagnostics;

namespace LlmPoc.Llm;

[AssetType( Name = "TinyStories Model", Extension = "llmmdl", Category = "LLM" )]
public sealed partial class TinyStoriesModelResource : GameResource
{
	public const int ExpectedModelByteCount = 27_855_830;
	public const string ExpectedModelSha256 =
		"12a5dfb8733e92b4a4eace12c8d8c9f022e2b956d0ca1e088d83acb70e1b12e9";
	public const string ExpectedConfigSha256 =
		"3682184758bea455d1dfb49d330e8627b93f3369e390540265a18afb8472321a";
	public const string ExpectedTokenizerSha256 =
		"1fe93b6152957cf9cfd6d89002467f789ce8b3f3e000b3a2edf27c808ddd0b9e";

	[Property]
	public string ConfigJson { get; set; } = "";

	[Property]
	public string TokenizerJson { get; set; } = "";

	[Property]
	public TinyStoriesModelBlob ModelData { get; set; } = new();

	[Property]
	public int ModelByteCount { get; set; }

	[Property]
	public string ModelSha256 { get; set; } = "";

	[Property]
	public string ConfigSha256 { get; set; } = "";

	[Property]
	public string TokenizerSha256 { get; set; } = "";
}

/// <summary>
/// Binary payload for the byte-identical SBOXLLM1 model file. BlobData keeps
/// the large payload out of the resource JSON and in the compiled resource blob.
/// </summary>
public sealed class TinyStoriesModelBlob : BlobData
{
	private const int MaximumPayloadBytes = 64 * 1024 * 1024;

	public byte[] Bytes { get; set; } = Array.Empty<byte>();

	public override int Version => 1;

	public override void Serialize( ref Writer writer )
	{
		writer.Stream.WriteArray<byte>( Bytes );
	}

	public override void Deserialize( ref Reader reader )
	{
		Bytes = reader.Stream.ReadArray<byte>( MaximumPayloadBytes );
	}
}

public sealed record TinyStoriesRuntimeArtifacts(
	SboxLlmModel Model,
	TinyStoriesConfig Config,
	Gpt2ByteBpeTokenizer Tokenizer,
	TinyStoriesModelResource Resource );

/// <summary>
/// The single production loading path for config, tokenizer and model weights.
/// Development parity artifacts deliberately remain outside this class.
/// </summary>
public static class TinyStoriesRuntimeArtifactLoader
{
	public static TinyStoriesRuntimeArtifacts LoadFromResourceLibrary( string path )
	{
		if ( string.IsNullOrWhiteSpace( path ) )
		{
			throw new ArgumentException(
				"[LLM:ERROR] Runtime model resource path cannot be empty.", nameof( path ) );
		}

		if ( !ResourceLibrary.TryGet<TinyStoriesModelResource>( path, out var resource ) ||
			resource is null )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Packaged TinyStories model resource '{path}' could not be " +
				"loaded through ResourceLibrary." );
		}

		return Load( resource, path );
	}

	public static TinyStoriesRuntimeArtifacts Load(
		TinyStoriesModelResource resource,
		string sourceLabel = "TinyStoriesModelResource" )
	{
		if ( resource is null )
		{
			throw new ArgumentNullException( nameof( resource ) );
		}

		FastTimer timer = FastTimer.StartNew();
		byte[] modelBytes = resource.ModelData?.Bytes ?? Array.Empty<byte>();
		ValidatePayloads( resource, modelBytes, sourceLabel );

		TinyStoriesConfig config = DeserializeRequired<TinyStoriesConfig>(
			resource.ConfigJson, "config", sourceLabel );
		config.Validate();
		if ( config.MaximumPositions <= 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Model maximum positions must be positive, found " +
				$"{config.MaximumPositions}." );
		}
		if ( config.ActivationFunction != "gelu_new" )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Runtime requires gelu_new, found " +
				$"'{config.ActivationFunction}'." );
		}

		SboxLlmModel model = SboxLlmModelLoader.Parse( modelBytes );
		Gpt2ByteBpeTokenizer tokenizer = Gpt2ByteBpeTokenizer.LoadFromJson(
			resource.TokenizerJson, sourceLabel );
		if ( tokenizer.VocabularySize != config.VocabularySize )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Tokenizer vocabulary has {tokenizer.VocabularySize} entries, " +
				$"model config expects {config.VocabularySize}." );
		}

		model.GetRequiredTensor( "transformer.wte.weight" )
			.RequireShape( config.VocabularySize, config.HiddenSize );
		model.GetRequiredTensor( "transformer.wpe.weight" )
			.RequireShape( config.MaximumPositions, config.HiddenSize );
		model.GetRequiredTensor( "transformer.ln_f.weight" )
			.RequireShape( config.HiddenSize );
		model.GetRequiredTensor( "transformer.ln_f.bias" )
			.RequireShape( config.HiddenSize );
		model.GetRequiredTensor( "lm_head.weight" )
			.RequireShape( config.VocabularySize, config.HiddenSize );

		LlmLog.Info(
			"LOAD",
			$"Packaged GameResource PASS source='{sourceLabel}' bytes={modelBytes.Length:N0} " +
			$"sha256={TinyStoriesModelResource.ExpectedModelSha256} " +
			$"tensors={model.TensorCount:N0} parameters={model.TotalParameterCount:N0} " +
			$"elapsed_ms={timer.ElapsedMilliSeconds:N2}." );
		LlmLog.Info(
			"PARITY",
			$"packaged_sources config_sha256={TinyStoriesModelResource.ExpectedConfigSha256} " +
			$"tokenizer_sha256={TinyStoriesModelResource.ExpectedTokenizerSha256} " +
			$"model_sha256={TinyStoriesModelResource.ExpectedModelSha256} PASS" );
		return new TinyStoriesRuntimeArtifacts( model, config, tokenizer, resource );
	}

	private static void ValidatePayloads(
		TinyStoriesModelResource resource,
		byte[] modelBytes,
		string sourceLabel )
	{
		RequireExact( "model byte count metadata", TinyStoriesModelResource.ExpectedModelByteCount,
			resource.ModelByteCount, sourceLabel );
		RequireExact( "model blob byte count", TinyStoriesModelResource.ExpectedModelByteCount,
			modelBytes.Length, sourceLabel );
		RequireExact( "model SHA-256 metadata", TinyStoriesModelResource.ExpectedModelSha256,
			resource.ModelSha256, sourceLabel );
		RequireExact( "config SHA-256 metadata", TinyStoriesModelResource.ExpectedConfigSha256,
			resource.ConfigSha256, sourceLabel );
		RequireExact( "tokenizer SHA-256 metadata", TinyStoriesModelResource.ExpectedTokenizerSha256,
			resource.TokenizerSha256, sourceLabel );

		RequireExact( "model blob SHA-256", TinyStoriesModelResource.ExpectedModelSha256,
			ComputeSha256( modelBytes ), sourceLabel );
		RequireExact( "config content SHA-256", TinyStoriesModelResource.ExpectedConfigSha256,
			ComputeSha256( Encoding.UTF8.GetBytes( resource.ConfigJson ?? "" ) ), sourceLabel );
		RequireExact( "tokenizer content SHA-256", TinyStoriesModelResource.ExpectedTokenizerSha256,
			ComputeSha256( Encoding.UTF8.GetBytes( resource.TokenizerJson ?? "" ) ), sourceLabel );
	}

	private static T DeserializeRequired<T>( string json, string label, string sourceLabel )
	{
		if ( string.IsNullOrWhiteSpace( json ) )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Packaged resource '{sourceLabel}' has empty {label} JSON." );
		}
		T value = Json.Deserialize<T>( json );
		return value ?? throw new InvalidOperationException(
			$"[LLM:ERROR] Packaged resource '{sourceLabel}' {label} JSON deserialized to null." );
	}

	public static string ComputeSha256( byte[] bytes )
	{
		return Convert.ToHexString( SHA256.HashData( bytes ) ).ToLowerInvariant();
	}

	private static void RequireExact<T>( string label, T expected, T actual, string sourceLabel )
	{
		if ( !EqualityComparer<T>.Default.Equals( expected, actual ) )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Packaged resource '{sourceLabel}' {label} expected " +
				$"'{expected}', found '{actual}'." );
		}
	}
}