// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.


using TypeAgent.ExamplesLib.CommandLine;
using TypeAgent.KnowPro;
using TypeAgent.KnowPro.Lang;

namespace KnowProConsole;

public class MemoryCommands : ICommandModule
{
    KnowProConsoleContext _kpContext;

    public MemoryCommands(KnowProConsoleContext context)
    {
        _kpContext = context;
    }

    public IList<Command> GetCommands()
    {
        return [
            SearchTermsDef(),
            MessagesDef(),
            AliasesDef(),
            SearchDef(),
            SearchRagDef(),
            AnswerRagDef(),
            AnswerDef()
        ];
    }

    private Command SearchTermsDef()
    {
        Command cmd = new("kpSearchTerms")
        {
        };
        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.SearchTermsAsync);
        return cmd;
    }

    private Task SearchTermsAsync(ParseResult result, CancellationToken cancellationToken)
    {
        IConversation conversation = EnsureConversation();
        KnowProWriter.WriteError("Not impl");
        return Task.CompletedTask;
    }

    private Command MessagesDef()
    {
        Command command = new("kpMessages");
        command.SetAction(this.MessagesAsync);
        return command;
    }
    private async Task MessagesAsync(ParseResult parseResult, CancellationToken cancellationToken)
    {
        IConversation conversation = EnsureConversation();

        await KnowProWriter.WriteMessagesAsync(conversation);
    }

    private Command SemanticRefsDef()
    {
        Command command = new("kpSemanticRefs");
        command.SetAction(this.SemanticRefsAsync);
        return command;
    }
    private async Task SemanticRefsAsync(ParseResult parseResult, CancellationToken cancellationToken)
    {
        IConversation conversation = EnsureConversation();

        await KnowProWriter.WriteSemanticRefsAsync(conversation);
    }

    private Command AliasesDef()
    {
        Command command = new("kpAlias")
        {
            Options.Arg<string>("term"),
            Options.Arg<string>("alias")
        };
        command.SetAction(this.AliasesAsync);
        return command;
    }
    private async Task AliasesAsync(ParseResult parseResult, CancellationToken cancellationToken)
    {
        IConversation conversation = EnsureConversation();

        //await KnowProWriter.WriteSemanticRefsAsync(conversation);
        NamedArgs namedArgs = new NamedArgs(parseResult);
        var term = namedArgs.Get("term");
        var alias = namedArgs.Get("alias");

        var aliases = conversation.SecondaryIndexes.TermToRelatedTermsIndex.Aliases;
        if (string.IsNullOrEmpty(term))
        {
            // Display all
            var relatedTerms = await aliases.GetTermsAsync(cancellationToken);
            KnowProWriter.WriteList(relatedTerms, ListType.Ol);
        }
        else if (!string.IsNullOrEmpty(alias))
        {
            await aliases.AddTermAsync(term, alias, cancellationToken);
        }
        else
        {
            var relatedTerms = await aliases.LookupTermAsync(term, cancellationToken);
            if (!relatedTerms.IsNullOrEmpty())
            {
                relatedTerms!.ForEach(KnowProWriter.WriteTerm);
            }
        }
    }

    private Command SearchDef()
    {
        Command command = new("kpSearch")
        {
            Options.Arg<string>("query"),
        };
        command.SetAction(this.SearchAsync);
        return command;
    }
    private Task SearchAsync(ParseResult args, CancellationToken cancellationToken)
    {
        IConversation conversation = EnsureConversation();
        KnowProWriter.WriteError("Not impl");
        return Task.CompletedTask;
    }

    private Command SearchRagDef()
    {
        Command command = new("kpSearchRag", "Text similarity search.")
        {
            Options.Arg<string>("query"),
            Options.Arg<int>("maxMatches", 25),
            Options.Arg<double>("minScore", 0.7),
            Options.Arg<int>("budget", 16 * 1024)
        };
        command.SetAction(this.SearchRagAsync);
        return command;
    }

    private async Task SearchRagAsync(ParseResult args, CancellationToken cancellationToken)
    {
        IConversation conversation = EnsureConversation();

        var namedArgs = new NamedArgs(args);
        var query = namedArgs.Get<string>("query");
        if (string.IsNullOrEmpty(query))
        {
            return;
        }
        var matches = await conversation.SearchRagAsync(
            query,
            namedArgs.Get<int>("maxMatches"),
            namedArgs.Get<double>("minScore"),
            namedArgs.Get<int>("budget"),
            cancellationToken
        );
        await KnowProWriter.WriteConversationSearchResultsAsync(
            conversation,
            matches,
            true,
            true,
            true
        );
    }

    private IConversation EnsureConversation()
    {
        return (_kpContext.Conversation is not null)
            ? _kpContext.Conversation!
            : throw new InvalidOperationException("No conversation loaded");
    }

    private Command AnswerRagDef()
    {
        Command command = new("kpAnswerRag", "Answer using classic RAG.")
        {
            Options.Arg<string>("query"),
            Options.Arg<int>("maxMatches", 25),
            Options.Arg<double>("minScore", 0.7),
            Options.Arg<int>("budget", 16 * 1024)
        };
        command.SetAction(this.AnswerRagAsync);
        return command;
    }

    private async Task AnswerRagAsync(ParseResult args, CancellationToken cancellationToken)
    {

        IConversation conversation = EnsureConversation();

        NamedArgs namedArgs = new NamedArgs(args);
        string? query = namedArgs.Get<string>("query");
        if (string.IsNullOrEmpty(query))
        {
            return;
        }

        AnswerResponse answer = await conversation.AnswerQuestionRagAsync(
            query,
            namedArgs.Get<double>("minScore"),
            namedArgs.Get<int>("budget"),
            new AnswerContextOptions()
            {
                MessagesTopK = namedArgs.Get<int>("maxMatches"),
            },
            null,
            cancellationToken
        ).ConfigureAwait(false);

        KnowProWriter.WriteLine();
        if (answer.Type == AnswerType.Answered)
        {
            KnowProWriter.WriteLine(ConsoleColor.Green, answer.Answer);
        }
        else
        {
            KnowProWriter.WriteLine(ConsoleColor.Yellow, answer.WhyNoAnswer);
        }
        KnowProWriter.WriteLine();
    }

    private Command AnswerDef()
    {
        Command cmd = new("kpAnswer")
        {
            Args.Arg<string>("query"),
            Options.Arg<bool>("debug", false)
        };
        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.AnswerAsync);
        return cmd;
    }

    private async Task AnswerAsync(ParseResult args)
    {
        IConversation conversation = EnsureConversation();

        NamedArgs namedArgs = new NamedArgs(args);
        string? query = namedArgs.Get<string>("query");
        if (string.IsNullOrEmpty(query))
        {
            return;
        }
        if (namedArgs.Get<bool>("debug"))
        {
            LangSearchDebugContext? debugContext = new LangSearchDebugContext();
            var searchResults = await conversation.SearchAsync(query, null, null, debugContext);
            KnowProWriter.WriteLine(ConsoleColor.DarkGray, debugContext.ToJson());
            foreach (var searchResult in searchResults)
            {
                AnswerContext context = await AnswerContext.FromSearchResultAsync(conversation, searchResult, new AnswerContextOptions() { MessagesTopK = 25, EntitiesTopK = 25 });
                if (namedArgs.Get<bool>("debug"))
                {
                    KnowProWriter.WriteLine(ConsoleColor.Cyan, context.ToJson());
                }

                var answer = await conversation.AnswerQuestionAsync(query, context);
                KnowProWriter.WriteJson(answer);
            }
        }
        else
        {
            AnswerResponse answer = await conversation.AnswerQuestionAsync(
                query,
                null,
                null,
                new AnswerContextOptions()
                {
                    MessagesTopK = 25,
                    EntitiesTopK = 25
                },
                null
            ).ConfigureAwait(false);

            KnowProWriter.WriteLine();
            if (answer.Type == AnswerType.Answered)
            {
                KnowProWriter.WriteLine(ConsoleColor.Green, answer.Answer);
            }
            else
            {
                KnowProWriter.WriteLine(ConsoleColor.Yellow, answer.WhyNoAnswer);
            }
            KnowProWriter.WriteLine();
        }
    }

}
