using Sandbox.Diagnostics;

namespace LlmPoc.Llm;

public sealed class GreedyGenerationStepResult
{
	public int Step { get; init; }
	public int InputSequenceLength { get; init; }
	public int NewTokenPosition { get; init; }
	public int TokenId { get; init; }
	public string DecodedToken { get; init; }
	public float Top1Logit { get; init; }
	public int Top2TokenId { get; init; }
	public float Top2Logit { get; init; }
	public float Top1Top2Margin { get; init; }
	public LogitRank[] TopFive { get; init; }
	public bool EosReached { get; init; }
	public double ForwardMilliseconds { get; init; }
	public double LmHeadMilliseconds { get; init; }
	public double StepMilliseconds { get; init; }
	public double[] LayerMilliseconds { get; init; }
}

public sealed class GreedyGenerationStepObservation
{
	public GreedyGenerationStepResult Step { get; init; }
	public int[] InputTokenIds { get; init; }
	public TinyStoriesModelForwardResult Forward { get; init; }
}

public sealed class GreedyGenerationResult
{
	public int[] PromptTokenIds { get; init; }
	public int[] GeneratedTokenIds { get; init; }
	public int[] FullSequenceTokenIds { get; init; }
	public string GeneratedText { get; init; }
	public string StopReason { get; init; }
	public bool EosReached { get; init; }
	public GreedyGenerationStepResult[] Steps { get; init; }
	public double TotalForwardMilliseconds { get; init; }
	public double TotalGenerationMilliseconds { get; init; }
}

/// <summary>
/// Deterministic correctness baseline: every new token recomputes the complete
/// model over the complete current context. No K/V state is retained.
/// </summary>
public static class TinyStoriesGreedyGenerator
{
	public static GreedyGenerationResult Generate(
		SboxLlmModel model,
		TinyStoriesConfig config,
		Gpt2ByteBpeTokenizer tokenizer,
		IReadOnlyList<int> promptTokenIds,
		int maxNewTokens,
		Action<GreedyGenerationStepObservation> observer = null,
		bool logLifecycle = true )
	{
		if ( model is null ) throw new ArgumentNullException( nameof( model ) );
		if ( config is null ) throw new ArgumentNullException( nameof( config ) );
		if ( tokenizer is null ) throw new ArgumentNullException( nameof( tokenizer ) );
		if ( promptTokenIds is null ) throw new ArgumentNullException( nameof( promptTokenIds ) );
		if ( promptTokenIds.Count == 0 )
		{
			throw new ArgumentException(
				"[LLM:ERROR] Greedy generation requires at least one prompt token.",
				nameof( promptTokenIds ) );
		}
		if ( maxNewTokens <= 0 )
		{
			throw new ArgumentOutOfRangeException(
				nameof( maxNewTokens ), maxNewTokens,
				"[LLM:ERROR] Greedy generation maxNewTokens must be positive." );
		}
		if ( promptTokenIds.Count > config.MaximumPositions - maxNewTokens )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Greedy generation prompt length {promptTokenIds.Count} plus " +
				$"maxNewTokens {maxNewTokens} would exceed maximum context length " +
				$"{config.MaximumPositions}; context truncation is disabled." );
		}

		int[] promptCopy = promptTokenIds.ToArray();
		List<int> context = new( promptCopy.Length + maxNewTokens );
		context.AddRange( promptCopy );
		List<int> generated = new( maxNewTokens );
		List<GreedyGenerationStepResult> steps = new( maxNewTokens );
		double totalForwardMilliseconds = 0;
		double totalGenerationMilliseconds = 0;
		bool eosReached = false;

		if ( logLifecycle )
		{
			LlmLog.Info(
				"GEN",
				$"starting prompt_tokens={promptCopy.Length} max_new_tokens={maxNewTokens} " +
				$"maximum_context={config.MaximumPositions} strategy=greedy " +
				"full_recompute=true kv_cache=false" );
		}

		for ( int stepIndex = 0; stepIndex < maxNewTokens; stepIndex++ )
		{
			FastTimer stepTimer = FastTimer.StartNew();
			TinyStoriesModelForwardResult forward =
				TinyStoriesModelForward.ForwardLastTokenLogits( model, config, context );
			int tokenId = TinyStoriesModelHead.ArgmaxFinite(
				forward.Logits.Data, $"generation.step{stepIndex}.logits" );
			LogitRank[] topFive = TinyStoriesModelHead.TopKFinite(
				forward.Logits.Data, 5, $"generation.step{stepIndex}.logits" );
			string decodedToken = tokenizer.Decode( new[] { tokenId } );
			float margin = topFive[0].Logit - topFive[1].Logit;
			bool stepEos = tokenId == config.EosTokenId;
			GreedyGenerationStepResult step = new()
			{
				Step = stepIndex,
				InputSequenceLength = context.Count,
				NewTokenPosition = context.Count,
				TokenId = tokenId,
				DecodedToken = decodedToken,
				Top1Logit = topFive[0].Logit,
				Top2TokenId = topFive[1].TokenId,
				Top2Logit = topFive[1].Logit,
				Top1Top2Margin = margin,
				TopFive = topFive,
				EosReached = stepEos,
				ForwardMilliseconds = forward.TotalMilliseconds,
				LmHeadMilliseconds = forward.LmHeadMilliseconds,
				StepMilliseconds = stepTimer.ElapsedMilliSeconds,
				LayerMilliseconds = forward.LayerMilliseconds
			};

			GreedyGenerationStepObservation observation = new()
			{
				Step = step,
				InputTokenIds = context.ToArray(),
				Forward = forward
			};
			observer?.Invoke( observation );
			if ( observer is null && logLifecycle )
			{
				LlmLog.Info(
					"GEN",
					$"step={stepIndex} context={context.Count} token={tokenId} " +
					$"piece='{EscapeVisible( decodedToken )}' margin={margin:G9} " +
					$"forward_ms={forward.TotalMilliseconds:N4}" );
			}
			LlmLog.Trace(
				"GEN",
				$"step={stepIndex} sequence={context.Count} new_token_position={context.Count} " +
				$"top5=[{string.Join( ",", topFive.Select( item => $"{item.TokenId}:{item.Logit:G9}" ) )}]" );

			context.Add( tokenId );
			generated.Add( tokenId );
			steps.Add( step );
			totalForwardMilliseconds += forward.TotalMilliseconds;
			totalGenerationMilliseconds += step.StepMilliseconds;
			if ( stepEos )
			{
				eosReached = true;
				break;
			}
		}

		for ( int index = 0; index < promptCopy.Length; index++ )
		{
			if ( promptCopy[index] != promptTokenIds[index] )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Greedy generation mutated caller prompt token {index}." );
			}
		}

		return new GreedyGenerationResult
		{
			PromptTokenIds = promptCopy,
			GeneratedTokenIds = generated.ToArray(),
			FullSequenceTokenIds = context.ToArray(),
			GeneratedText = tokenizer.Decode( generated ),
			StopReason = eosReached ? "eos" : "max_new_tokens",
			EosReached = eosReached,
			Steps = steps.ToArray(),
			TotalForwardMilliseconds = totalForwardMilliseconds,
			TotalGenerationMilliseconds = totalGenerationMilliseconds
		};
	}

	private static string EscapeVisible( string value )
	{
		return value
			.Replace( "\\", "\\\\" )
			.Replace( "\r", "\\r" )
			.Replace( "\n", "\\n" )
			.Replace( "\t", "\\t" )
			.Replace( "'", "\\'" );
	}
}