using Sandbox.Diagnostics;

namespace LlmPoc.Llm;

public sealed class TinyStoriesModelForwardResult
{
	public int SequenceLength { get; init; }
	public Tensor FinalLayerNorm { get; init; }
	public Tensor Logits { get; init; }
	public double EmbeddingMilliseconds { get; init; }
	public double[] LayerMilliseconds { get; init; }
	public string[] AttentionTypes { get; init; }
	public double AllLayersMilliseconds { get; init; }
	public double FinalLayerNormMilliseconds { get; init; }
	public double LmHeadMilliseconds { get; init; }
	public double TotalMilliseconds { get; init; }
}

/// <summary>
/// Production-oriented full-context forward orchestration. It deliberately reuses
/// the parity-proven embedding, transformer-layer, final-norm, and LM-head kernels.
/// </summary>
public static class TinyStoriesModelForward
{
	public static TinyStoriesModelForwardResult ForwardLastTokenLogits(
		SboxLlmModel model,
		TinyStoriesConfig config,
		IReadOnlyList<int> tokenIds,
		bool logDiagnostics = false )
	{
		if ( model is null ) throw new ArgumentNullException( nameof( model ) );
		if ( config is null ) throw new ArgumentNullException( nameof( config ) );
		if ( tokenIds is null ) throw new ArgumentNullException( nameof( tokenIds ) );
		if ( tokenIds.Count == 0 )
		{
			throw new ArgumentException(
				"[LLM:ERROR] Model forward requires at least one token.", nameof( tokenIds ) );
		}
		if ( tokenIds.Count > config.MaximumPositions )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Model forward sequence length {tokenIds.Count} exceeds " +
				$"maximum context length {config.MaximumPositions}; context truncation is disabled." );
		}
		for ( int index = 0; index < tokenIds.Count; index++ )
		{
			int tokenId = tokenIds[index];
			if ( tokenId < 0 || tokenId >= config.VocabularySize )
			{
				throw new IndexOutOfRangeException(
					$"[LLM:ERROR] Model forward token ID {tokenId} at index {index} is outside " +
					$"[0,{config.VocabularySize})." );
			}
		}

		FastTimer totalTimer = FastTimer.StartNew();
		FastTimer embeddingTimer = FastTimer.StartNew();
		Tensor hiddenStates = TinyStoriesForwardStages.CombineEmbeddings(
			model, config, tokenIds, logDiagnostics );
		double embeddingMilliseconds = embeddingTimer.ElapsedMilliSeconds;

		double[] layerMilliseconds = new double[config.LayerCount];
		string[] attentionTypes = new string[config.LayerCount];
		double allLayersMilliseconds = 0;
		for ( int layerIndex = 0; layerIndex < config.LayerCount; layerIndex++ )
		{
			GptNeoLayerForwardResult layer = GptNeoTransformerLayer.Forward(
				model,
				config,
				hiddenStates,
				layerIndex,
				logDiagnostics,
				validateInputMutation: false );
			hiddenStates = layer.Output;
			layerMilliseconds[layerIndex] = layer.ElapsedMilliseconds;
			attentionTypes[layerIndex] = layer.AttentionType;
			allLayersMilliseconds += layer.ElapsedMilliseconds;
		}

		FastTimer finalLayerNormTimer = FastTimer.StartNew();
		Tensor finalLayerNorm = TinyStoriesForwardStages.ApplyFinalLayerNorm(
			model, config, hiddenStates, "forward.final_ln", logDiagnostics );
		double finalLayerNormMilliseconds = finalLayerNormTimer.ElapsedMilliSeconds;

		FastTimer lmHeadTimer = FastTimer.StartNew();
		Tensor logits = TinyStoriesModelHead.ProjectLastPositionNoBias(
			finalLayerNorm,
			model.GetRequiredTensor( "lm_head.weight" ),
			"forward.final_logits_last_position" );
		double lmHeadMilliseconds = lmHeadTimer.ElapsedMilliSeconds;
		logits.RequireShape( config.VocabularySize );

		if ( logDiagnostics )
		{
			LlmLog.Trace(
				"PERF",
				$"stage=model_forward sequence={tokenIds.Count} embeddings_ms={embeddingMilliseconds:N4} " +
				$"layers_ms={allLayersMilliseconds:N4} final_ln_ms={finalLayerNormMilliseconds:N4} " +
				$"lm_head_ms={lmHeadMilliseconds:N4} total_ms={totalTimer.ElapsedMilliSeconds:N4}" );
		}

		return new TinyStoriesModelForwardResult
		{
			SequenceLength = tokenIds.Count,
			FinalLayerNorm = finalLayerNorm,
			Logits = logits,
			EmbeddingMilliseconds = embeddingMilliseconds,
			LayerMilliseconds = layerMilliseconds,
			AttentionTypes = attentionTypes,
			AllLayersMilliseconds = allLayersMilliseconds,
			FinalLayerNormMilliseconds = finalLayerNormMilliseconds,
			LmHeadMilliseconds = lmHeadMilliseconds,
			TotalMilliseconds = totalTimer.ElapsedMilliSeconds
		};
	}
}