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

using Microsoft.CodeAnalysis.CSharp.Syntax;
using TypeAgent.Common;
using TypeAgent.ExamplesLib.CommandLine;
using TypeAgent.KnowPro;
using TypeAgent.KnowPro.Answer;
using TypeAgent.KnowPro.Lang;
using static TypeAgent.ExamplesLib.KnowProWriter;

namespace KnowProConsole;

public class TestCommands : ICommandModule
{
    KnowProConsoleContext _kpContext;

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

    public IList<Command> GetCommands()
    {
        return [
            SearchTermsDef(),
            SearchPropertyTermsDef(),
            SearchMessagesTermsDef(),
            TestEmbeddingsDef(),
            SearchQueryTermsDef(),
            SearchLangDef(),
            KnowledgeDef(),
            BuildIndexDef(),
        ];
    }

    private Command SearchTermsDef()
    {
        Command cmd = new("kpTestSearchTerms")
        {
            Options.Arg<bool>("andTerms", "'And' all terms. Set to TRUE, otherwise search terms are 'Or'.", false),
            Options.Arg<bool>("displayAsc", "Display results in ascending order.", true),
            Options.Arg<bool>("distinct", "Show distinct results.", true),
            Options.Arg<string>("endDate", "Ending at this date (ISO format)."),
            Options.Arg<int?>("endMinute", "Ending minute."),
            Options.Arg<bool>("exact", "Exact match only. No related terms.", false),
            Options.Arg<string>("ktype", "Filter results to a specific knowledge type [entity | topic | action | tag]."),
            Options.Arg<int>("maxToDisplay", "Maximum number of results to display", 25),
            Options.Arg<string>("startDate", "Starting at this date (ISO format)."),
            Options.Arg<int?>("startMinute", "Starting minute."),
        };

        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.SearchTermsAsync);
        return cmd;
    }

    private async Task SearchTermsAsync(ParseResult result, CancellationToken cancellationToken)
    {
        IConversation conversation = EnsureConversation();

        //// Hard coded test for now
        //SearchTermGroup searchGroup = new SearchTermGroup(SearchTermBooleanOp.Or)
        //{
        //    "Children of Time",
        //    "book"
        //};
        //await TestSearchKnowledgeAsync(conversation, searchGroup, cancellationToken);

        //searchGroup = new SearchTermGroup(SearchTermBooleanOp.OrMax, searchGroup.Terms);
        //await TestSearchKnowledgeAsync(conversation, searchGroup, cancellationToken);

        //searchGroup = new SearchTermGroup(SearchTermBooleanOp.And, searchGroup.Terms);
        //await TestSearchKnowledgeAsync(conversation, searchGroup, cancellationToken);

        //searchGroup = new SearchTermGroup(SearchTermBooleanOp.OrMax)
        //{
        //    "Children of Physics",
        //    "book"
        //};

        // get the named args
        NamedArgs namedArgs = new NamedArgs(result);

        WriterOptions options = new WriterOptions()
        {
            MaxToDisplay = namedArgs.Get("maxToDisplay") is not null ? namedArgs.Get<int>("maxToDisplay") : WriterOptions.Default.MaxToDisplay,
            Ascending = namedArgs.Get<bool>("displayAsc"),
            Distinct = namedArgs.Get<bool>("distinct")
        };

        var convTimeRange = await conversation.GetStartTimestampRangeAsync();

        await TestSearchKnowledgeAsync(conversation, SearchSeletExpressionFromCommandArgs(result, convTimeRange), options, cancellationToken);
    }

    /// <summary>
    /// Conerts the supplied parsed command line arguments into a search select expression.
    /// </summary>
    /// <param name="parsedArgs">The parsed command arguments.</param>
    /// <param name="convTimeRange">The conversation timerange</param>
    /// <returns>A select expression representing the supplied argumentns</returns>
    internal static SearchSelectExpr SearchSeletExpressionFromCommandArgs(ParseResult parsedArgs, TimestampRange? convTimeRange)
    {
        // get the named args
        NamedArgs namedArgs = new NamedArgs(parsedArgs);

        // get the logic operator
        SearchTermBooleanOp logicOperator = SearchTermBooleanOp.Or;
        if (namedArgs.Get("andTerms") is not null && bool.Parse(namedArgs.Get("andTerms")!))
        {
            logicOperator = SearchTermBooleanOp.And;
        }

        // are we doing an exact match?
        bool exactMatch = namedArgs.Get("exact") is not null && bool.Parse(namedArgs.Get("exact")!);

        SearchTermGroup stg = new SearchTermGroup(logicOperator)
        {
            // add unmatched tokens as search terms
            { parsedArgs.UnmatchedTokens, exactMatch }
        };

        WhenFilter? when = null;
        if (namedArgs.Get("ktype") is not null)
        {
            when = new WhenFilter();
            when.KnowledgeType = KnowledgeType.Parse(namedArgs.Get("ktype")!);
        }

        // start and end dates
        TimestampRange? timestampRange = null;
        if (namedArgs.Get("startDate") is not null || namedArgs.Get("endDate") is not null)
        {
            timestampRange = new TimestampRange()
            {
                StartTimestamp = namedArgs.Get("startDate") ?? convTimeRange?.StartTimestamp ?? string.Empty,
                EndTimestamp = namedArgs.Get("endDate") ?? convTimeRange?.EndTimestamp ?? string.Empty
            };

            DateRange dateRange = new(timestampRange.Value);

            if (namedArgs.Get("startMinute") is not null)
            {
                dateRange.Start = dateRange.Start.AddMinutes(namedArgs.Get<int>("startMinute"));
            }

            if (namedArgs.Get("endMinute") is not null && dateRange.HasEnd)
            {
                dateRange.End = dateRange.End!.Value.AddMinutes(namedArgs.Get<int>("endMinute"));
            }

            when ??= new WhenFilter();
            when.DateRange = dateRange;
        }

        SearchSelectExpr select = new SearchSelectExpr(stg, when);

        return select;
    }

    private Command SearchPropertyTermsDef()
    {
        Command cmd = new("kpTestPropertyTerms")
        {
        };
        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.SearchPropertyTermsAsync);
        return cmd;
    }

    private async Task SearchPropertyTermsAsync(ParseResult result, CancellationToken cancellationToken)
    {
        IConversation conversation = EnsureConversation();

        // Hard coded test for now
        SearchTermGroup searchGroup = new SearchTermGroup(SearchTermBooleanOp.Or)
        {
            { "genre", "sci-fi" },
            { KnowledgePropertyName.EntityName, "Children of Time" },
        };

        // get the named args
        NamedArgs namedArgs = new NamedArgs(result);

        WriterOptions options = new WriterOptions()
        {
            MaxToDisplay = namedArgs.Get("MaxToDisplay") is not null ? namedArgs.Get<int>("MaxToDisplay") : WriterOptions.Default.MaxToDisplay,
            Ascending = namedArgs.Get<bool>("displayAsc"),
            Distinct = namedArgs.Get<bool>("distinct")
        };

        await TestSearchKnowledgeAsync(conversation, searchGroup, options, cancellationToken);
    }

    private Command SearchMessagesTermsDef()
    {
        Command cmd = new("kpTestMessageTerms")
        {
        };
        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.SearchMessagesAsync);
        return cmd;
    }

    private async Task SearchMessagesAsync(ParseResult result, CancellationToken cancellationToken)
    {
        IConversation conversation = EnsureConversation();

        var lengths = await conversation.Messages.GetMessageLengthAsync([34], cancellationToken);
        KnowProWriter.WriteJson(lengths);

        var message = await conversation.Messages.GetAsync(34, cancellationToken);
        KnowProWriter.WriteMessage(message);

        var length = await conversation.Messages.GetMessageLengthAsync(34, cancellationToken);
        KnowProWriter.WriteJson(length);

        // Hard coded test for now
        SearchSelectExpr select = new(
            new SearchTermGroup(SearchTermBooleanOp.Or)
            {
                { "genre", "sci-fi" },
                { KnowledgePropertyName.EntityName, "Children of Time" }
            }
        );

        ConversationSearchResult? searchResults = await conversation.SearchAsync(
            select,
            null,
            cancellationToken
        );
        await KnowProWriter.WriteConversationSearchResultsAsync(
            conversation,
            searchResults,
            true,
            true
        );

        DateRange? conversationDateRange = await conversation.GetDateRangeAsync();
        if (conversationDateRange is not null)
        {
            select.When = new WhenFilter()
            {
                DateRange = new()
                {
                    Start = conversationDateRange.Value.Start,
                    End = conversationDateRange.Value.Start.AddMinutes(10)
                }
            };

        }
        searchResults = await conversation.SearchAsync(
            select,
            new SearchOptions()
            {
                ExactMatch = false,
                MaxCharsInBudget = 1024
            },
            cancellationToken
        );
        await KnowProWriter.WriteConversationSearchResultsAsync(conversation, searchResults, true, true);
    }

    private Command TestEmbeddingsDef()
    {
        Command cmd = new("kpTestEmbeddings")
        {
            Options.Arg<string>("text", "text to embed"),
            Options.Arg<bool>("add", "add to index")
        };
        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.TestEmbeddingsAsync);
        return cmd;
    }

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

        var settings = AzureModelApiSettings.EmbeddingSettingsFromEnv();
        var model = new OpenAITextEmbeddingModel(settings);
        var modelWithCache = new TextEmbeddingModelWithCache(
            model,
            new TextEmbeddingCache(1024)
        );
        modelWithCache.Cache.PersistentCache = conversation.SecondaryIndexes.TermToRelatedTermsIndex.FuzzyIndex as IReadOnlyCache<string, Embedding>;
        NamedArgs namedArgs = new(args);
        string? text = namedArgs.Get("text");// ?? "The quick brown fox";
        if (!string.IsNullOrEmpty(text))
        {
            var result = await model.GenerateAsync(text, cancellationToken);
            KnowProWriter.WriteLine(result.Length);
            return;
        }

        IList<string> allTerms = await conversation.SemanticRefIndex.GetTermsAsync(cancellationToken);
        var fuzzyIndex = conversation.SecondaryIndexes.TermToRelatedTermsIndex.FuzzyIndex;
        if (namedArgs.Get<bool>("add"))
        {
            await fuzzyIndex.ClearAsync(cancellationToken);
            ProgressBar progress = new ProgressBar(allTerms.Count);
            foreach (var batch in allTerms.Batch(16))
            {
                await fuzzyIndex.AddTermsAsync(batch, cancellationToken);
                progress.Advance(batch.Count);
                await Task.Delay(1000, cancellationToken);
            }
        }

        foreach (var term in allTerms)
        {
            KnowProWriter.WriteLine(ConsoleColor.Cyan, term);
            _kpContext.Stopwatch.Restart();
            var matches = await fuzzyIndex.LookupTermAsync(term, 10, 0, cancellationToken);
            _kpContext.Stopwatch.Stop();

            KnowProWriter.WriteTiming(_kpContext.Stopwatch);

            matches.ForEach(KnowProWriter.WriteTerm);
        }
    }

    private Command BuildIndexDef()
    {
        Command cmd = new("kpTestBuildIndex")
        {
            Options.Arg<bool>("related", "index related terms", true),
            Options.Arg<bool>("messages", "index messages", true),
            Options.Arg<bool>("semanticRefs", "index semantic refs", true)
        };
        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.BuildIndexAsync);
        return cmd;
    }

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

        NamedArgs namedArgs = new NamedArgs(args);

        var cachingModel = conversation.Settings.RelatedTermIndexSettings.EmbeddingIndexSetting?.EmbeddingModel as TextEmbeddingModelWithCache;
        try
        {
            if (cachingModel is not null)
            {
                cachingModel.CacheEnabled = false;
            }
            if (namedArgs.Get<bool>("related"))
            {
                await conversation.SecondaryIndexes.TermToRelatedTermsIndex.FuzzyIndex.ClearAsync(cancellationToken);

                await conversation.RebuildRelatedTermsIndexAsync(cancellationToken);
            }
            if (namedArgs.Get<bool>("messages"))
            {
                await conversation.SecondaryIndexes.MessageIndex.ClearAsync(cancellationToken);

                await conversation.UpdateMessageIndexAsync(false, cancellationToken);
            }
            if (namedArgs.Get<bool>("semanticRefs"))
            {
                await conversation.SemanticRefIndex.ClearAsync(cancellationToken);
                await conversation.SecondaryIndexes.PropertyToSemanticRefIndex.ClearAsync(cancellationToken);
                await conversation.SecondaryIndexes.TermToRelatedTermsIndex.FuzzyIndex.ClearAsync(cancellationToken);

                await conversation.UpdateSemanticRefIndexAsync(cancellationToken);
            }
            KnowProWriter.WriteLine();
        }
        catch (Exception ex)
        {
            KnowProWriter.WriteError(ex);
        }
        finally
        {
            if (cachingModel is not null)
            {
                cachingModel.CacheEnabled = true;
            }
        }
    }

    internal async Task TestSearchKnowledgeAsync(IConversation conversation, SearchTermGroup searchGroup, WriterOptions options, CancellationToken cancellationToken)
    {
        await TestSearchKnowledgeAsync(conversation, new SearchSelectExpr(searchGroup), options, cancellationToken);
    }

    internal async Task TestSearchKnowledgeAsync(IConversation conversation, SearchSelectExpr searchExpresion, WriterOptions options, CancellationToken cancellationToken)
    {
        KnowProWriter.WriteLine(searchExpresion);

        var results = await conversation.SearchKnowledgeAsync(
            searchExpresion,
            null,
            cancellationToken
        ).ConfigureAwait(false);

        await KnowProWriter.WriteKnowledgeSearchResultsAsync(_kpContext.Conversation!, results, options);
    }

    private Command SearchQueryTermsDef()
    {
        Command cmd = new("kpTestSearchQuery")
        {
            Args.Arg<string>("query")
        };
        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.SearchQueryTermsAsync);
        return cmd;
    }

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

        NamedArgs namedArgs = new NamedArgs(args);
        var query = namedArgs.Get("query");
        if (string.IsNullOrEmpty(query))
        {
            return;
        }
        var model = new OpenAIChatModel();
        SearchQueryTranslator translator = new SearchQueryTranslator(model);
        var result = await translator.TranslateAsync(query, null, cancellationToken);
        KnowProWriter.WriteJson(result);
    }

    private Command SearchLangDef()
    {
        Command cmd = new("kpTestSearchLang")
        {
            Args.Arg<string>("query")
        };
        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.SearchLangAsync);
        return cmd;
    }

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

        NamedArgs namedArgs = new NamedArgs(args);
        var query = namedArgs.Get("query");
        if (string.IsNullOrEmpty(query))
        {
            return;
        }
        if (conversation is IMemory memory)
        {
            IList<ConversationSearchResult> results = await memory.SearchAsync(query, null, null, null, cancellationToken);
            foreach (var result in results)
            {
                await KnowProWriter.WriteConversationSearchResultsAsync(
                    conversation,
                    result,
                    true,
                    false
                );
            }
        }
    }

    private Command KnowledgeDef()
    {
        Command cmd = new("kpTestKnowledge")
        {
            Args.Arg<string>("text")
        };
        cmd.TreatUnmatchedTokensAsErrors = false;
        cmd.SetAction(this.KnowledgeAsync);
        return cmd;
    }

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

        NamedArgs namedArgs = new NamedArgs(args);
        var text = namedArgs.Get("text");
        if (string.IsNullOrEmpty(text))
        {
            return;
        }
        var model = new OpenAIChatModel();
        KnowledgeExtractor exctractor = new KnowledgeExtractor(model);
        var result = await exctractor.ExtractAsync(text, cancellationToken);
        if (result is not null)
        {
            KnowProWriter.WriteJson(result);
        }
    }

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

}
