using System.Text.Json.Serialization;

namespace LlmPoc.Llm;

public sealed class TinyStoriesConfig
{
	[JsonPropertyName( "vocab_size" )]
	public int VocabularySize { get; set; }

	[JsonPropertyName( "max_position_embeddings" )]
	public int MaximumPositions { get; set; }

	[JsonPropertyName( "hidden_size" )]
	public int HiddenSize { get; set; }

	[JsonPropertyName( "num_layers" )]
	public int LayerCount { get; set; }

	[JsonPropertyName( "num_heads" )]
	public int AttentionHeadCount { get; set; }

	[JsonPropertyName( "window_size" )]
	public int AttentionWindowSize { get; set; }

	[JsonPropertyName( "activation_function" )]
	public string ActivationFunction { get; set; }

	[JsonPropertyName( "layer_norm_epsilon" )]
	public float LayerNormEpsilon { get; set; }

	[JsonPropertyName( "embed_dropout" )]
	public float EmbeddingDropout { get; set; }

	[JsonPropertyName( "bos_token_id" )]
	public int BosTokenId { get; set; }

	[JsonPropertyName( "eos_token_id" )]
	public int EosTokenId { get; set; }

	[JsonPropertyName( "tie_word_embeddings" )]
	public bool TieWordEmbeddings { get; set; }

	[JsonPropertyName( "attention_layers" )]
	public string[] AttentionLayers { get; set; }

	public void Validate()
	{
		if ( VocabularySize <= 0 || HiddenSize <= 0 || LayerCount <= 0 || AttentionHeadCount <= 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Invalid model config: vocab={VocabularySize}, hidden={HiddenSize}, " +
				$"layers={LayerCount}, heads={AttentionHeadCount}." );
		}
		if ( HiddenSize % AttentionHeadCount != 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Hidden size {HiddenSize} is not divisible by " +
				$"attention head count {AttentionHeadCount}." );
		}
		if ( AttentionLayers is null || AttentionLayers.Length != LayerCount )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Expected {LayerCount} attention layer modes, " +
				$"found {AttentionLayers?.Length ?? 0}." );
		}
	}
}

public sealed class TensorManifestDocument
{
	[JsonPropertyName( "format" )]
	public string Format { get; set; }

	[JsonPropertyName( "version" )]
	public int Version { get; set; }

	[JsonPropertyName( "dtype" )]
	public string Dtype { get; set; }

	[JsonPropertyName( "tensor_count" )]
	public int TensorCount { get; set; }

	[JsonPropertyName( "tensors" )]
	public List<TensorManifestEntry> Tensors { get; set; }
}

public sealed class TensorManifestEntry
{
	[JsonPropertyName( "name" )]
	public string Name { get; set; }

	[JsonPropertyName( "shape" )]
	public int[] Shape { get; set; }

	[JsonPropertyName( "dtype" )]
	public string Dtype { get; set; }

	[JsonPropertyName( "elements" )]
	public long Elements { get; set; }

	[JsonPropertyName( "bytes" )]
	public long Bytes { get; set; }
}

public sealed class LlmReferenceData
{
	[JsonPropertyName( "prompt" )]
	public string Prompt { get; set; }

	[JsonPropertyName( "input_token_ids" )]
	public int[] InputTokenIds { get; set; }

	[JsonPropertyName( "greedy_next_token_id" )]
	public int GreedyNextTokenId { get; set; }

	[JsonPropertyName( "greedy_next_token_text" )]
	public string GreedyNextTokenText { get; set; }

	[JsonPropertyName( "generated_token_ids" )]
	public int[] GeneratedTokenIds { get; set; }

	[JsonPropertyName( "generated_text" )]
	public string GeneratedText { get; set; }

	[JsonPropertyName( "full_sequence_token_ids" )]
	public int[] FullSequenceTokenIds { get; set; }

	[JsonPropertyName( "eos_token_id" )]
	public int EosTokenId { get; set; }

	[JsonPropertyName( "bos_token_id" )]
	public int BosTokenId { get; set; }
}

public static class ModelManifestValidator
{
	public static void Validate( SboxLlmModel model, TensorManifestDocument manifest )
	{
		if ( model is null )
		{
			throw new ArgumentNullException( nameof( model ) );
		}
		if ( manifest is null )
		{
			throw new ArgumentNullException( nameof( manifest ) );
		}
		if ( manifest.Format != "SBOXLLM" || manifest.Version != SboxLlmModelLoader.SupportedVersion )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Manifest format/version expected SBOXLLM/" +
				$"{SboxLlmModelLoader.SupportedVersion}, found {manifest.Format}/{manifest.Version}." );
		}
		if ( manifest.Dtype != "float32" )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Manifest dtype expected float32, found '{manifest.Dtype}'." );
		}
		if ( manifest.Tensors is null )
		{
			throw new InvalidOperationException( "[LLM:ERROR] Manifest tensors array is missing." );
		}
		if ( manifest.TensorCount != manifest.Tensors.Count || model.TensorCount != manifest.TensorCount )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Tensor count mismatch: binary={model.TensorCount}, " +
				$"manifest field={manifest.TensorCount}, manifest entries={manifest.Tensors.Count}." );
		}

		for ( int index = 0; index < manifest.Tensors.Count; index++ )
		{
			TensorManifestEntry expected = manifest.Tensors[index];
			Tensor actual = model.GetRequiredTensor( expected.Name );
			if ( expected.Dtype != "float32" )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] {expected.Name} manifest dtype expected float32, found {expected.Dtype}." );
			}
			actual.RequireShape( expected.Shape );
			if ( expected.Elements < 0 || expected.Bytes < 0 ||
				actual.ElementCount != expected.Elements ||
				expected.Bytes / sizeof( float ) != expected.Elements ||
				expected.Bytes % sizeof( float ) != 0 )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] {expected.Name} manifest counts disagree: " +
					$"binary elements={actual.ElementCount}, manifest elements={expected.Elements}, " +
					$"manifest bytes={expected.Bytes}." );
			}
		}
	}
}