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

import { KnowproContext } from "./knowproContext.js";
import { dateTime, readJsonFile } from "typeagent";
import { error, Result, success } from "typechat";
import {
    BatchCallback,
    Comparison,
    SearchRequest,
    searchRequestDef,
} from "./types.js";
import {
    compareArray,
    compareNumberArray,
    compareObject,
    compareStringArray,
    dateRangeToTimeRange,
    getCommandArgs,
    isStem,
    isUndefinedOrEmpty,
    queryError,
    runTestBatch,
} from "./common.js";
import { getLangSearchResult } from "./knowproCommands.js";
import { parseTypedArguments } from "interactive-app";
import { execSearchRequest } from "./knowproCommands.js";
import * as kp from "knowpro";
import { TestRunReport } from "./logging.js";

export type LangSearchResults = {
    searchText: string;
    cmd?: string | undefined;
    searchQueryExpr: kp.querySchema.SearchQuery;
    compiledQueryExpr?: kp.SearchQueryExpr[] | undefined;
    results: LangSearchResult[];
};

export type LangSearchResult = {
    messageMatches: kp.MessageOrdinal[];
    entityMatches?: kp.SemanticRefOrdinal[] | undefined;
    topicMatches?: kp.SemanticRefOrdinal[] | undefined;
    actionMatches?: kp.SemanticRefOrdinal[] | undefined;
};

export async function runSearchBatch(
    context: KnowproContext,
    batchFilePath: string,
    destFilePath?: string,
    cb?: BatchCallback<Result<LangSearchResults>>,
    stopOnError: boolean = false,
): Promise<Result<LangSearchResults[]>> {
    return runTestBatch(
        batchFilePath,
        async (cmd, args) => {
            const response = await getSearchResults(context, args);
            if (response.success) {
                response.data.cmd = cmd;
            }
            return response;
        },
        destFilePath,
        cb,
        stopOnError,
    );
}

async function getSearchResults(
    context: KnowproContext,
    args: string[],
): Promise<Result<LangSearchResults>> {
    const response = await execSearchRequest(context, args);
    if (!response.searchResults.success) {
        return response.searchResults;
    }
    const results = collectLangSearchResults(
        response.debugContext.searchText,
        response.searchResults.data,
        response.debugContext,
    );
    return success(results);
}

function collectLangSearchResults(
    searchText: string,
    searchResults: kp.ConversationSearchResult[],
    debugContext: kp.LanguageSearchDebugContext,
): LangSearchResults {
    return {
        searchText,
        searchQueryExpr: debugContext.searchQuery!,
        compiledQueryExpr: debugContext.searchQueryExpr,
        results: searchResults.map((cr) => {
            const lr: LangSearchResult = {
                messageMatches: cr.messageMatches.map((m) => m.messageOrdinal),
            };
            getKnowledgeResults(cr, lr);
            return lr;
        }),
    };
}

export async function verifyLangSearchResultsBatch(
    context: KnowproContext,
    batchFilePath: string,
    cb?: BatchCallback<Result<Comparison<LangSearchResults>>>,
    stopOnError: boolean = false,
): Promise<Result<Comparison<LangSearchResults>[]>> {
    const expectedResults =
        await readJsonFile<LangSearchResults[]>(batchFilePath);
    if (!expectedResults || expectedResults.length === 0) {
        return error("No results in file");
    }

    const results: Comparison<LangSearchResults>[] = [];
    for (let i = 0; i < expectedResults.length; ++i) {
        const expected = expectedResults[i];
        const args = getCommandArgs(expected.cmd);
        if (args.length === 0) {
            continue;
        }
        let response = await getSearchResults(context, args);
        if (response.success) {
            const actual = response.data;
            const error = compareLangSearchResults(actual, expected);
            let comparisonResult: Comparison<LangSearchResults> = {
                actual,
                expected,
                error,
            };
            results.push(comparisonResult);
            if (cb) {
                cb(success(comparisonResult), i, expectedResults.length);
            }
        } else {
            response = queryError(expected.cmd!, response);
            if (cb) {
                cb(response, i, expectedResults.length);
            }
            if (stopOnError) {
                return response;
            }
        }
    }
    return success(results);
}

export async function* runLangSearchBatch(
    conversation: kp.IConversation,
    queryTranslator: kp.SearchQueryTranslator,
    searchQueries: Iterable<string>,
): AsyncIterableIterator<Result<LangSearchResults>> {
    for (const searchText of searchQueries) {
        const debugContext: kp.LanguageSearchDebugContext = {};
        const searchResult = await getLangSearchResult(
            conversation,
            queryTranslator,
            searchText,
            undefined,
            undefined,
            debugContext,
        );
        if (searchResult.success) {
            yield success({
                searchText,
                searchQueryExpr: debugContext.searchQuery!,
                results: searchResult.data.map((cr) => {
                    const lr: LangSearchResult = {
                        messageMatches: cr.messageMatches.map(
                            (m) => m.messageOrdinal,
                        ),
                    };
                    getKnowledgeResults(cr, lr);
                    return lr;
                }),
            });
        } else {
            yield searchResult;
        }
    }

    function getKnowledgeResults(
        cr: kp.ConversationSearchResult,
        lr: LangSearchResult,
    ) {
        lr.entityMatches = getKnowledgeResult(cr, "entity");
        lr.topicMatches = getKnowledgeResult(cr, "topic");
        lr.actionMatches = getKnowledgeResult(cr, "action");
    }

    function getKnowledgeResult(
        cr: kp.ConversationSearchResult,
        type: kp.KnowledgeType,
    ) {
        return cr.knowledgeMatches
            .get(type)
            ?.semanticRefMatches.map((sr) => sr.semanticRefOrdinal);
    }
}

export function createSearchTestReport(
    results: Comparison<LangSearchResults>[],
    srcDataPath: string,
    dateRange: dateTime.DateRange,
): TestRunReport<Comparison<LangSearchResults>> {
    const errors = results.filter(
        (c) => c.error !== undefined && c.error.length > 0,
    );
    return {
        name: "search",
        timeRange: dateRangeToTimeRange(dateRange),
        srcData: srcDataPath,
        errors,
        countRun: results.length,
    };
}

function getKnowledgeResults(
    cr: kp.ConversationSearchResult,
    lr: LangSearchResult,
) {
    lr.entityMatches = getMatchedSemanticRefOrdinals(cr, "entity");
    lr.topicMatches = getMatchedSemanticRefOrdinals(cr, "topic");
    lr.actionMatches = getMatchedSemanticRefOrdinals(cr, "action");
}

function getMatchedSemanticRefOrdinals(
    cr: kp.ConversationSearchResult,
    type: kp.KnowledgeType,
) {
    return cr.knowledgeMatches
        .get(type)
        ?.semanticRefMatches.map((sr) => sr.semanticRefOrdinal);
}

function compareLangSearchResults(
    lr1: LangSearchResults,
    lr2: LangSearchResults,
): string | undefined {
    let error = compareSearchQuery(lr1.searchQueryExpr, lr2.searchQueryExpr);
    if (error !== undefined && error.length > 0) {
        return error;
    }
    if (lr1.results.length !== lr2.results.length) {
        return `Number results ${lr1.results.length} != ${lr2.results.length}`;
    }
    for (let i = 0; i < lr1.results.length; ++i) {
        const error = compareLangSearchResult(lr1.results[i], lr2.results[i]);
        if (error !== undefined && error.length > 0) {
            return error;
        }
    }
    return undefined;
}

function compareLangSearchResult(
    lr1: LangSearchResult,
    lr2: LangSearchResult,
): string | undefined {
    let error = compareNumberArray(
        lr1.messageMatches,
        lr2.messageMatches,
        "message",
    );
    if (error !== undefined) {
        return error;
    }
    error = compareNumberArray(lr1.entityMatches, lr2.entityMatches, "entity");
    if (error !== undefined) {
        return error;
    }
    error = compareNumberArray(lr1.topicMatches, lr2.topicMatches, "topic");
    if (error !== undefined) {
        return error;
    }
    error = compareNumberArray(lr1.actionMatches, lr2.actionMatches, "action");
    if (error !== undefined) {
        return error;
    }
    return undefined;
}

export function compareSearchQuery(
    s1: kp.querySchema.SearchQuery,
    s2: kp.querySchema.SearchQuery,
): string | undefined {
    if (s1.searchExpressions.length !== s2?.searchExpressions.length) {
        return `searchQuery.searchExpressions.length: ${s1.searchExpressions.length} !== ${s2.searchExpressions.length}`;
    }
    for (let i = 0; i < s1.searchExpressions.length; ++i) {
        const error = compareSearchExpr(
            s1.searchExpressions[i],
            s2.searchExpressions[i],
        );
        if (error !== undefined) {
            return error;
        }
    }
    return undefined;
}

export function compareSearchExpr(
    s1: kp.querySchema.SearchExpr,
    s2: kp.querySchema.SearchExpr,
): string | undefined {
    if (s1.filters.length !== s2.filters.length) {
        return `SearchExpr.filters.length: ${s1.filters.length} !== ${s2.filters.length}`;
    }

    for (let i = 0; i < s1.filters.length; ++i) {
        const f1 = s1.filters[i];
        const f2 = s2.filters[i];

        let error = compareObject(
            f1.entitySearchTerms,
            f2.entitySearchTerms,
            "entitySearchTerms",
        );
        if (error !== undefined) {
            return error;
        }
        error = compareActionTerm(f1.actionSearchTerm, f2.actionSearchTerm);
        if (error !== undefined) {
            return error;
        }
        error = compareStringArray(
            f1.searchTerms,
            f2.searchTerms,
            "searchTerms",
        );
        if (error !== undefined) {
            return error;
        }
        error = compareObject(f1.timeRange, f2.timeRange, "timeRange");
        if (error !== undefined) {
            return error;
        }
    }
    return undefined;
}

function compareActionTerm(
    x?: kp.querySchema.ActionTerm,
    y?: kp.querySchema.ActionTerm,
    strict: boolean = true,
): string | undefined {
    if (x === undefined && y === undefined) {
        return undefined;
    }
    let error = compareStringArray(
        x?.actionVerbs?.words,
        y?.actionVerbs?.words,
        "verbs",
        true,
        strict ? undefined : (xv, yv) => isStem(xv, yv) || isStem(yv, xv),
    );
    if (error !== undefined) {
        return error;
    }
    error = compareEntityTerms(
        x?.actorEntities,
        y?.actorEntities,
        "actorEntities",
        strict,
    );
    if (error !== undefined) {
        return error;
    }
    error = compareEntityTerms(
        x?.additionalEntities,
        y?.additionalEntities,
        "additionalEntities",
        strict,
    );
    if (error !== undefined) {
        return error;
    }
    error = compareEntityTerms(
        x?.targetEntities,
        y?.targetEntities,
        "targetEntities",
        strict,
    );
    if (error !== undefined) {
        return error;
    }
    if (x?.isInformational !== y?.isInformational) {
        return `isInformational: ${x?.isInformational} !== ${y?.isInformational}`;
    }
    return undefined;
}

function compareEntityTerms(
    x: kp.querySchema.EntityTerm[] | "*" | undefined,
    y: kp.querySchema.EntityTerm[] | "*" | undefined,
    label: string,
    strict: boolean = true,
): string | undefined {
    if (isUndefinedOrEmpty(x) && isUndefinedOrEmpty(y)) {
        return undefined;
    }
    if (!Array.isArray(x) || !Array.isArray(y)) {
        return compareObject(x, y, label);
    }
    return compareArray(x, y, label, (xt, yt) =>
        compareEntityTerm(xt, yt, label, strict),
    );
}

function compareEntityTerm(
    x: kp.querySchema.EntityTerm,
    y: kp.querySchema.EntityTerm,
    label: string,
    strict: boolean,
): string | undefined {
    if (strict) {
        return compareObject(x, y, label);
    }

    if (x.name !== y.name) {
        return `${label}.name: ${x.name} !== ${y.name}`;
    }

    let error = compareTypes(x.type, y.type, label);
    if (error !== undefined) {
        return error;
    }

    if (x.isNamePronoun !== y.isNamePronoun) {
        return `${label}.isNamePronoun: ${x.isNamePronoun} !== ${y.isNamePronoun}`;
    }

    return compareObject(x.facets, y.facets, label + ".facets");

    function compareTypes(
        xType: string[] | undefined,
        yType: string[] | undefined,
        label: string,
    ): string | undefined {
        let error = compareStringArray(xType, yType, label + ".type");
        if (error === undefined) {
            return undefined;
        }
        if (xType !== undefined && isTypeVariation(xType, yType)) {
            return undefined;
        }
        if (yType !== undefined && isTypeVariation(yType, xType)) {
            return undefined;
        }
        return error;
    }
}

function isTypeVariation(xType: string[], yType?: string[]): boolean {
    if (xType.length === 1 && isUndefinedOrEmpty(yType)) {
        const xValue = xType[0];
        switch (xValue.toLowerCase()) {
            default:
                break;

            case "person":
                return true;
        }
    }
    return false;
}

export function compareSearchQueryScope(
    s1: kp.querySchema.SearchQuery,
    s2: kp.querySchema2.SearchQuery,
): string | undefined {
    if (s1.searchExpressions.length !== s2?.searchExpressions.length) {
        return `searchQuery.searchExpressions.length: ${s1.searchExpressions.length} !== ${s2.searchExpressions.length}`;
    }
    for (let i = 0; i < s1.searchExpressions.length; ++i) {
        const error = compareSearchExprScope(
            s1.searchExpressions[i],
            s2.searchExpressions[i],
        );
        if (error !== undefined) {
            return error;
        }
    }
    return undefined;
}

export function compareSearchExprScope(
    s1: kp.querySchema.SearchExpr,
    s2: kp.querySchema2.SearchExpr,
    strict: boolean = false,
): string | undefined {
    if (s1.filters.length !== s2.filters.length) {
        return `SearchExpr.filters.length: ${s1.filters.length} !== ${s2.filters.length}`;
    }

    for (let i = 0; i < s1.filters.length; ++i) {
        const f1 = s1.filters[i];
        const f2 = s2.filters[i];

        let error = compareEntityTerms(
            f1.entitySearchTerms,
            f2.entitySearchTerms,
            "entitySearchTerms",
            strict,
        );
        if (error !== undefined) {
            return error;
        }
        error = compareActionTerm(
            f1.actionSearchTerm,
            f2.actionSearchTerm,
            strict,
        );
        if (error !== undefined) {
            return error;
        }
        error = compareStringArray(
            f1.searchTerms,
            f2.searchTerms,
            "searchTerms",
        );
        if (error !== undefined) {
            return error;
        }
        error = compareObject(
            f1.timeRange,
            f2.scopeSubQuery?.timeRange,
            "timeRange",
        );
        if (error !== undefined) {
            return error;
        }
    }
    return undefined;
}

export interface ScopeQueryComparison
    extends Comparison<
        kp.querySchema2.SearchQuery,
        kp.querySchema.SearchQuery
    > {
    query: string;
}

export async function compareSearchQueryTranslations(
    queryTranslator: kp.SearchQueryTranslator,
    query: string,
): Promise<Result<ScopeQueryComparison>> {
    if (!queryTranslator.translate2) {
        return error("not supported");
    }
    const result = await queryTranslator.translate(query);
    if (!result.success) {
        return result;
    }
    const resultScope = await queryTranslator.translate2!(query);
    if (!resultScope.success) {
        return resultScope;
    }

    let cmp: ScopeQueryComparison = {
        actual: resultScope.data,
        expected: result.data,
        query,
    };
    cmp.error = compareSearchQueryScope(result.data, resultScope.data);
    return success(cmp);
}

export async function compareSearchQueryTranslationsBatch(
    context: KnowproContext,
    batchFilePath: string,
    cb?: BatchCallback<Result<ScopeQueryComparison>>,
    stopOnError: boolean = false,
): Promise<Result<ScopeQueryComparison[]>> {
    const queryTranslator = context.getConversationQueryTranslator();
    return runTestBatch(
        batchFilePath,
        (_, args: string[]) => {
            const request = parseTypedArguments<SearchRequest>(
                args,
                searchRequestDef(),
            );
            return compareSearchQueryTranslations(
                queryTranslator,
                request.query,
            );
        },
        undefined,
        cb,
        stopOnError,
    );
}
