using AIStudio.Dialogs.Settings;
using AIStudio.Tools.AssistantLaunch;
using AIStudio.Tools.AssistantSessions;

namespace AIStudio.Assistants.GrammarSpelling;

public partial class AssistantGrammarSpelling : AssistantBaseCore<SettingsDialogGrammarSpelling>
{
    protected override Tools.Components Component => Tools.Components.GRAMMAR_SPELLING_ASSISTANT;
    
    protected override string Title => T("Grammar & Spelling Checker");
    
    protected override string Description => T("Check the grammar and spelling of a text.");
    
    // The prompt of this assistant forbids explanations, so the chat gets a prompt of its own, which keeps the rules of the language:
    protected override ChatSystemPromptChoice ChatSystemPrompt => ChatSystemPromptChoice.Custom(this.CreateChatSystemPrompt);

    protected override string SystemPrompt => 
        $"""
        You are an expert in languages and their rules. For example, you know how US and UK English or German in
        Germany and German in Austria differ. You receive text as input. You check the spelling and grammar of
        this text according to the rules of {this.SystemPromptLanguage()}. You never add information. You
        never ask the user for additional information. You do not attempt to improve the wording of the text.
        Your response includes only the corrected text. Do not explain your changes. If no changes are needed,
        you return the text unchanged.
        """;

    /// <summary>
    /// Creates the system prompt of a chat which continues a check of spelling and grammar, for follow-up questions about it.
    /// </summary>
    private string CreateChatSystemPrompt() =>
        $"""
        You are an expert in languages and their rules. For example, you know how US and UK English or German in
        Germany and German in Austria differ. The user had a text checked for spelling and grammar according to the
        rules of {this.SystemPromptLanguage()}: the conversation so far holds the text and its corrected version.
        Now you answer the user's follow-up questions about the corrections, e.g., which rule applies or why a sentence
        changed, and you check further texts or revise the corrected text when the user asks for it. When you correct a
        text, you correct only its spelling and grammar and keep its wording, unless the user asks otherwise.

        Answer questions in the language the user writes in.
        """;
    
    protected override bool AllowProfiles => false;
    
    protected override bool ShowDedicatedProgress => true;
    
    protected override Func<string> Result2Copy => () => this.correctedText;
    
    protected override IReadOnlyList<IButtonData> FooterButtons =>
    [
        new SendToButton
        {
            UseResultingContentBlockData = false,
            GetText = () => string.IsNullOrWhiteSpace(this.correctedText) ? this.inputText : this.correctedText
        },
    ];
    
    protected override string SubmitText => T("Proofread");

    protected override Func<Task> SubmitAction => this.ProofreadText;

    protected override string SendToChatVisibleUserPromptPrefix => T("Check the following text for grammar and spelling mistakes:");

    protected override string SendToChatVisibleUserPromptContent => this.inputText;
    
    protected override void ResetForm()
    {
        this.inputText = string.Empty;
        this.correctedText = string.Empty;
        if (!this.MightPreselectValues())
        {
            this.selectedTargetLanguage = CommonLanguages.AS_IS;
            this.customTargetLanguage = string.Empty;
        }
    }

    protected override bool MightPreselectValues()
    {
        if (this.SettingsManager.ConfigurationData.GrammarSpelling.PreselectOptions)
        {
            this.selectedTargetLanguage = this.SettingsManager.ConfigurationData.GrammarSpelling.PreselectedTargetLanguage;
            this.customTargetLanguage = this.SettingsManager.ConfigurationData.GrammarSpelling.PreselectedOtherLanguage;
            return true;
        }
        
        return false;
    }
    
    #region Overrides of AssistantBase

    protected override void ApplyTextLaunchPayload(TextLaunchPayload payload) => this.inputText = payload.Text;

    #endregion

    private string inputText = string.Empty;
    private CommonLanguages selectedTargetLanguage;
    private string customTargetLanguage = string.Empty;
    private string correctedText = string.Empty;
    private static readonly AssistantSessionStateKey<string> INPUT_TEXT_STATE_KEY = new(nameof(inputText));
    private static readonly AssistantSessionStateKey<CommonLanguages> SELECTED_TARGET_LANGUAGE_STATE_KEY = new(nameof(selectedTargetLanguage));
    private static readonly AssistantSessionStateKey<string> CUSTOM_TARGET_LANGUAGE_STATE_KEY = new(nameof(customTargetLanguage));
    private static readonly AssistantSessionStateKey<string> CORRECTED_TEXT_STATE_KEY = new(nameof(correctedText));

    /// <inheritdoc />
    protected override void CaptureCustomAssistantSessionState(AssistantSessionStateWriter state)
    {
        state.Set(INPUT_TEXT_STATE_KEY, this.inputText);
        state.Set(SELECTED_TARGET_LANGUAGE_STATE_KEY, this.selectedTargetLanguage);
        state.Set(CUSTOM_TARGET_LANGUAGE_STATE_KEY, this.customTargetLanguage);
        state.Set(CORRECTED_TEXT_STATE_KEY, this.correctedText);
    }

    /// <inheritdoc />
    protected override void RestoreCustomAssistantSessionState(AssistantSessionStateReader state)
    {
        state.Restore(INPUT_TEXT_STATE_KEY, value => this.inputText = value);
        state.Restore(SELECTED_TARGET_LANGUAGE_STATE_KEY, value => this.selectedTargetLanguage = value);
        state.Restore(CUSTOM_TARGET_LANGUAGE_STATE_KEY, value => this.customTargetLanguage = value);
        state.Restore(CORRECTED_TEXT_STATE_KEY, value => this.correctedText = value);
    }

    private string? ValidateText(string text)
    {
        if(string.IsNullOrWhiteSpace(text))
            return T("Please provide a text as input. You might copy the desired text from a document or a website.");
        
        return null;
    }
    
    private string? ValidateCustomLanguage(string language)
    {
        if(this.selectedTargetLanguage == CommonLanguages.OTHER && string.IsNullOrWhiteSpace(language))
            return T("Please provide a custom language.");
        
        return null;
    }
    
    private string SystemPromptLanguage()
    {
        var lang = this.selectedTargetLanguage switch
        {
            CommonLanguages.AS_IS => "the source language",
            CommonLanguages.OTHER => this.customTargetLanguage,
            
            _ => $"{this.selectedTargetLanguage.Name}",
        };

        if (string.IsNullOrWhiteSpace(lang))
            return "the source language";

        return lang;
    }

    private async Task ProofreadText()
    {
        await this.Form!.ValidateAsync();
        if (!this.InputIsValid)
            return;
        
        this.CreateChatThread();
        var time = this.AddUserRequest(this.inputText);
        
        this.correctedText = await this.AddAIResponseAsync(time);
        if (!this.IsAssistantComponentDisposed)
            await this.JsRuntime.GenerateAndShowDiff(this.inputText, this.correctedText);
    }

    protected override async Task OnAssistantSessionRenderedAsync(AssistantSessionSnapshot snapshot)
    {
        if (!snapshot.IsActive && !string.IsNullOrWhiteSpace(this.inputText) && !string.IsNullOrWhiteSpace(this.correctedText))
            await this.JsRuntime.GenerateAndShowDiff(this.inputText, this.correctedText);
    }
}