@using Sandbox
@using Sandbox.UI
@using LlmPoc.Llm
@inherits PanelComponent
@namespace LlmPoc.UI

<root>
	<div class="chat-shell">
	<div class="header">
		<div class="heading">
			<div class="title">Native Local LLM POC</div>
			<div class="subtitle">TinyStories-Instruct-1M | native in-process C# inference</div>
		</div>
		<div class="model-status @StatusClass">@ModelStatus</div>
	</div>

	<div @ref="MessageList" class="conversation">
		@foreach ( ChatMessage message in Messages )
		{
			<div class="message @message.RoleClass">
				<div class="role">@message.Role</div>
				<div class="text">@message.Text</div>
			</div>
		}

		@if ( RequestInFlight )
		{
			<div class="message assistant generating">
				<div class="role">TinyStories</div>
				<div class="text">Generating...</div>
			</div>
		}
	</div>

	@if ( !string.IsNullOrEmpty( VisibleError ) )
	{
		<div class="error-status">@VisibleError</div>
	}

	<div class="input-row">
		<TextEntry @ref="InputBox"
			class="prompt-input"
			MaxLength="@(512)"
			Placeholder="Type a prompt..."
			onsubmit="@OnInputSubmit" />
		<button class="send-button @(CanSend ? "" : "disabled")" disabled="@(!CanSend)" onclick="@OnInputSubmit">
			@(RequestInFlight ? "Working..." : "Send")
		</button>
	</div>

	<div class="milestone-note">
		Native greedy generation. Each prompt is independent; visible history is not model memory.
	</div>
	</div>
</root>

@code
{
	[Property]
	public bool RunIntegrationSelfTest { get; set; }

	[Property]
	public string IntegrationTestPrompt { get; set; } =
		TinyStoriesLanguageModelServiceRegression.GoldenPrompt;

	[Property]
	public int IntegrationTestMaxNewTokens { get; set; } = 12;

	[Property]
	public bool ValidateGoldenIntegrationTest { get; set; } = true;

	private readonly List<ChatMessage> Messages = new();
	private TinyStoriesLanguageModelService ModelService;
	private TextEntry InputBox;
	private Panel MessageList;
	private bool RequestInFlight;
	private bool HasStarted;
	private bool SelfTestStarted;
	private string UiError = "";
	private LanguageModelGenerationResult LastResult;

	private bool CanSend => !RequestInFlight && ModelService?.CanGenerate == true;

	private string ModelStatus
	{
		get
		{
			if ( ModelService is null ) return "Error";
			if ( RequestInFlight || ModelService.State == LanguageModelServiceState.Generating )
			{
				return "Generating...";
			}
			if ( LastResult is not null && ModelService.State == LanguageModelServiceState.Ready )
			{
				return $"Ready | {LastResult.GeneratedTokenIds.Length} tokens in " +
					$"{LastResult.TotalGenerationMilliseconds:N0} ms";
			}
			return ModelService.State switch
			{
				LanguageModelServiceState.NotLoaded => "Model not loaded",
				LanguageModelServiceState.Loading => "Loading model...",
				LanguageModelServiceState.Ready => "Ready",
				LanguageModelServiceState.Error => "Error",
				_ => ModelService.State.ToString()
			};
		}
	}

	private string StatusClass => ModelService?.State switch
	{
		LanguageModelServiceState.Ready => "ready",
		LanguageModelServiceState.Loading => "loading",
		LanguageModelServiceState.Generating => "loading",
		LanguageModelServiceState.Error => "error",
		_ => "not-ready"
	};

	private string VisibleError => !string.IsNullOrEmpty( UiError )
		? UiError
		: ModelService?.ErrorMessage ?? "";

	protected override void OnStart()
	{
		HasStarted = true;
		ModelService = GetComponent<TinyStoriesLanguageModelService>();
		if ( ModelService is null )
		{
			UiError = "Native model service is missing from the chat GameObject.";
			LlmLog.Error( UiError );
			return;
		}

		Messages.Add( new ChatMessage(
			"System",
			"Native in-process generation is active. Each prompt starts a new independent story.",
			"system" ) );
		LlmLog.Info( "INIT", $"Chat UI started with real service {ModelService.DisplayName}." );
		_ = InitializeModelAsync();
	}

	private async Task InitializeModelAsync()
	{
		try
		{
			await ModelService.InitializeAsync();
		}
		catch
		{
			UiError = ModelService.ErrorMessage;
		}
		StateHasChanged();
	}

	protected override void OnTreeFirstBuilt()
	{
		InputBox?.Focus();
		if ( HasStarted && RunIntegrationSelfTest && !SelfTestStarted && InputBox is not null )
		{
			SelfTestStarted = true;
			string prompt = IntegrationTestPrompt ?? "";
			InputBox.Text = prompt;
			InputBox.Text = "";
			_ = SendAsync(
				prompt,
				IntegrationTestMaxNewTokens,
				ValidateGoldenIntegrationTest );
		}
	}

	private void OnInputSubmit()
	{
		if ( InputBox is null || RequestInFlight )
		{
			return;
		}
		if ( ModelService is null )
		{
			UiError = "Native model service is unavailable.";
			StateHasChanged();
			return;
		}
		if ( !ModelService.CanGenerate )
		{
			UiError = ModelService.State == LanguageModelServiceState.Loading
				? "Model is still loading."
				: ModelService.ErrorMessage;
			StateHasChanged();
			return;
		}

		string prompt = InputBox.Text ?? "";
		if ( string.IsNullOrWhiteSpace( prompt ) )
		{
			UiError = "Enter a prompt before sending.";
			StateHasChanged();
			return;
		}

		InputBox.Text = "";
		_ = SendAsync( prompt, ModelService.DefaultMaxNewTokens, validateGolden: false );
	}

	private async Task SendAsync( string prompt, int maxNewTokens, bool validateGolden )
	{
		if ( RequestInFlight ) return;
		UiError = "";
		Messages.Add( new ChatMessage( "You", prompt, "user" ) );
		RequestInFlight = true;
		StateHasChanged();
		MessageList?.TryScrollToBottom();
		LlmLog.Info(
			"GEN",
			$"UI submitted prompt chars={prompt.Length} max_new_tokens={maxNewTokens}; " +
			"generating_state=true." );

		try
		{
			LanguageModelGenerationResult result = await ModelService.GenerateAsync(
				prompt, maxNewTokens );
			LastResult = result;
			if ( validateGolden )
			{
				TinyStoriesLanguageModelServiceRegression.ValidateGolden(
					result, ModelService.State );
			}
			string visibleText = string.IsNullOrEmpty( result.GeneratedText )
				? "[EOS]"
				: result.GeneratedText;
			Messages.Add( new ChatMessage( "TinyStories", visibleText, "assistant" ) );
			LlmLog.Info(
				"GEN",
				$"Real native response displayed chars={visibleText.Length} " +
				$"generated_tokens={result.GeneratedTokenIds.Length} " +
				$"ids=[{string.Join( ",", result.GeneratedTokenIds )}] " +
				$"text='{EscapeVisible( visibleText )}' PASS." );
		}
		catch ( Exception error )
		{
			UiError = !string.IsNullOrWhiteSpace( ModelService?.ErrorMessage )
				? ModelService.ErrorMessage
				: error.Message;
			LlmLog.Error( $"UI native generation request failed: {error}" );
		}
		finally
		{
			RequestInFlight = false;
			StateHasChanged();
			MessageList?.TryScrollToBottom();
			InputBox?.Focus();
			LlmLog.Info( "GEN", "UI generating_state=false." );
		}
	}

	protected override int BuildHash()
	{
		return HashCode.Combine(
			Messages.Count,
			RequestInFlight,
			UiError,
			ModelService?.State ?? LanguageModelServiceState.NotLoaded,
			ModelService?.ErrorMessage ?? "",
			LastResult?.GeneratedTokenIds.Length ?? 0 );
	}

	private sealed record ChatMessage( string Role, string Text, string RoleClass );

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