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

import chalk from "chalk";
import {
    RequestAction,
    printProcessRequestActionResult,
    FullAction,
    ProcessRequestActionResult,
    ExplanationOptions,
    equalNormalizedObject,
    toFullActions,
} from "agent-cache";

import {
    type CommandHandlerContext,
    getCommandResult,
    getRequestId,
} from "../../commandHandlerContext.js";
import { CachedImageWithDetails } from "typechat-utils";
import { Logger } from "telemetry";
import { executeActions } from "../../../execute/actionHandlers.js";
import {
    TypeAgentTranslator,
    TranslatedAction,
} from "../../../translation/agentTranslators.js";
import {
    isMultipleAction,
    isPendingRequest,
} from "../../../translation/multipleActionSchema.js";
import registerDebug from "debug";
import ExifReader from "exifreader";
import { ProfileNames } from "../../../utils/profileNames.js";
import {
    ActionContext,
    ParsedCommandParams,
    SessionContext,
    CompletionGroup,
} from "@typeagent/agent-sdk";
import { CommandHandler } from "@typeagent/agent-sdk/helpers/command";
import { DispatcherName, isUnknownAction } from "../dispatcherUtils.js";
import { getTranslatorForSchema } from "../../../translation/translateRequest.js";
import { getActivityNamespaceSuffix } from "../../../translation/matchRequest.js";
import { addRequestToMemory, addResultToMemory } from "../../memory.js";
import { requestCompletion } from "../../../translation/requestCompletion.js";
import {
    getCannotUseCacheReason,
    getHistoryContext,
    interpretRequest,
    InterpretResult,
} from "../../../translation/interpretRequest.js";

const debugExplain = registerDebug("typeagent:explain");

async function canTranslateWithoutContext(
    requestAction: RequestAction,
    usedTranslators: Map<string, TypeAgentTranslator>,
    logger?: Logger,
) {
    if (requestAction.history === undefined) {
        return;
    }

    // Do the retranslation check, which will also check the action.
    const oldActions: FullAction[] = toFullActions(requestAction.actions);
    const newActions: (FullAction | undefined)[] = [];
    const request = requestAction.request;
    try {
        const translations = new Map<string, TranslatedAction>();
        for (const [schemaName, translator] of usedTranslators) {
            const result = await translator.checkTranslate(request);
            if (!result.success) {
                throw new Error("Failed to translate without history context");
            }
            const newActions = result.data;
            const count = isMultipleAction(newActions)
                ? newActions.parameters.requests.length
                : 1;

            if (count !== oldActions.length) {
                throw new Error("Action count mismatch without context");
            }
            translations.set(schemaName, result.data);
        }

        let index = 0;
        for (const { action } of requestAction.actions) {
            const schemaName = action.schemaName;
            const newTranslatedActions = translations.get(schemaName)!;
            let newAction: TranslatedAction;
            if (isMultipleAction(newTranslatedActions)) {
                const entry = newTranslatedActions.parameters.requests[index];
                if (isPendingRequest(entry)) {
                    throw new Error("Pending request in multiple action");
                }
                newAction = entry.action;
            } else {
                newAction = newTranslatedActions;
            }
            const newSchemaName = usedTranslators
                .get(schemaName)!
                .getSchemaName(newAction.actionName);
            if (newSchemaName === undefined) {
                // Should not happen
                throw new Error(
                    `Internal Error: Unable to match schema name for action '${newAction.actionName}'`,
                );
            }
            newActions.push({
                schemaName: newSchemaName,
                ...newAction,
            });
        }

        debugExplain(
            `With context: ${JSON.stringify(oldActions)}\nWithout context: ${JSON.stringify(newActions)}`,
        );

        if (oldActions.length !== newActions.length) {
            throw new Error("Action count mismatch without context");
        }

        index = 0;
        for (const oldAction of oldActions) {
            const newAction = newActions[index];
            if (newAction === undefined) {
                throw new Error(`Action missing without context`);
            }

            if (newAction.actionName !== oldAction.actionName) {
                throw new Error(`Action Name mismatch without context`);
            }

            if (
                !equalNormalizedObject(
                    newAction.parameters,
                    oldAction.parameters,
                )
            ) {
                throw new Error(`Action parameters mismatch without context`);
            }
            index++;
        }
        logger?.logEvent("contextlessTranslation", {
            request,
            actions: oldActions,
            history: requestAction.history,
            newActions,
        });
    } catch (e: any) {
        logger?.logEvent("contextlessTranslation", {
            requestAction,
            actions: oldActions,
            history: requestAction.history,
            newActions,
            error: e.message,
        });
        throw e;
    }
}

function getExplainerOptions(
    requestAction: RequestAction,
    context: CommandHandlerContext,
): ExplanationOptions | undefined {
    if (!context.session.explanation) {
        // Explanation is disabled
        return undefined;
    }

    if (
        !context.session.getConfig().explainer.filter.multiple &&
        requestAction.actions.length > 1
    ) {
        // filter multiple
        return undefined;
    }

    const usedTranslators = new Map<string, TypeAgentTranslator>();
    const actions = requestAction.actions;
    const activeSchemas = new Set(context.agents.getActiveSchemas());
    for (const { action } of actions) {
        if (isUnknownAction(action)) {
            // can't explain unknown actions
            return undefined;
        }

        const schemaName = action.schemaName;
        if (context.agents.getActionConfig(schemaName).cached === false) {
            // explanation disable at the translator level
            return undefined;
        }

        // TODO: This does not support activities.
        usedTranslators.set(
            schemaName,
            getTranslatorForSchema(context, schemaName, activeSchemas),
        );
    }
    const { list, value, translate } =
        context.session.getConfig().explainer.filter.reference;

    // In NFA mode, skip old construction system validation
    const isNFAMode = context.agentCache.isUsingNFAGrammar();

    return {
        namespaceSuffix: getActivityNamespaceSuffix(
            context,
            requestAction.history?.activityContext,
        ),
        // In NFA mode, skip contextless translation check (not needed for grammar rules)
        checkExplainable:
            translate && !isNFAMode
                ? (requestAction: RequestAction) =>
                      canTranslateWithoutContext(
                          requestAction,
                          usedTranslators,
                          context.logger,
                      )
                : undefined,
        valueInRequest: isNFAMode ? false : value,
        noReferences: list,
    };
}

async function requestExplain(
    context: CommandHandlerContext,
    attachments: CachedImageWithDetails[] | undefined,
    translationResult: InterpretResult,
) {
    // Make sure the current requestId is captured
    const requestId = getRequestId(context);

    const { fromCache, fromUser, requestAction } = translationResult;

    const notifyExplained = (error?: string) => {
        context.clientIO.notify(
            requestId,
            "explained",
            {
                time: new Date().toLocaleTimeString(),
                fromCache,
                fromUser,
                error,
            },
            DispatcherName,
        );
    };

    const notifyExplainedResult = (result: ProcessRequestActionResult) => {
        const explanationResult = result.explanationResult.explanation;
        const error = explanationResult.success
            ? undefined
            : explanationResult?.message;
        notifyExplained(error);

        // Notify about grammar result (success or rejection)
        if (result.grammarResult) {
            context.clientIO.notify(
                requestId,
                "grammarRule",
                {
                    success: result.grammarResult.success,
                    message: result.grammarResult.message,
                    rule: result.grammarResult.generatedRule,
                    time: new Date().toLocaleTimeString(),
                },
                DispatcherName,
            );
        }
    };

    const cannotUseCacheReason = getCannotUseCacheReason(
        context,
        attachments,
        requestAction.history,
    );
    if (cannotUseCacheReason !== undefined) {
        notifyExplained(`cannot not use cache (${cannotUseCacheReason})`);
    }
    if (fromCache && !fromUser) {
        // If it is from cache, and not from the user, explanation is not necessary.
        notifyExplained();
        return;
    }

    const options = getExplainerOptions(requestAction, context);
    if (options === undefined) {
        return;
    }

    const processRequestActionP = context.agentCache.processRequestAction(
        requestAction,
        true,
        options,
    );

    if (context.explanationAsynchronousMode) {
        processRequestActionP
            .then(notifyExplainedResult)
            .catch((e) => notifyExplained(e.message));
    } else {
        console.log(
            chalk.grey(`Generating explanation for '${requestAction}'`),
        );
        const processRequestActionResult = await processRequestActionP;
        notifyExplainedResult(processRequestActionResult);

        // Print result details (includes grammar result if using NFA)
        printProcessRequestActionResult(processRequestActionResult);
    }
}

export class RequestCommandHandler implements CommandHandler {
    public readonly description = "Translate and explain a request";
    public readonly parameters = {
        args: {
            request: {
                description: "Request to translate",
                implicitQuotes: true,
                optional: true, // can be optional since the user can supply images and no text // TODO: revisit
                skipKnowledgeExtraction: true,
            },
        },
    } as const;
    public async run(
        context: ActionContext<CommandHandlerContext>,
        params: ParsedCommandParams<typeof this.parameters>,
        attachments?: string[],
    ) {
        const systemContext = context.sessionContext.agentContext;
        const profiler = systemContext.commandProfiler?.measure(
            ProfileNames.request,
        );
        try {
            // Don't handle the request if it contains the separator
            const request =
                params.args.request === undefined ? "" : params.args.request;
            if (request.includes(RequestAction.Separator)) {
                throw new Error(
                    `Invalid translation request with translation separator '${RequestAction.Separator}'.  Use @explain if you want to explain a translation.`,
                );
            }

            // store attachments for later reuse
            const cachedAttachments: CachedImageWithDetails[] = [];
            if (
                attachments &&
                systemContext.session.sessionDirPath !== undefined
            ) {
                for (let i = 0; i < attachments?.length; i++) {
                    const [attachmentName, tags]: [string, ExifReader.Tags] =
                        await systemContext.session.storeUserSuppliedFile(
                            attachments![i],
                        );

                    cachedAttachments.push(
                        new CachedImageWithDetails(
                            tags,
                            attachmentName,
                            attachments![i],
                        ),
                    );
                }
            }

            // Make sure we clear any left over streaming context
            systemContext.streamingActionContext?.closeActionContext();
            systemContext.streamingActionContext = undefined;

            // Translate to action

            // Get the history context before adding the request to memory
            const history = getHistoryContext(systemContext);
            if (systemContext.userRequestKnowledgeExtraction === true) {
                addRequestToMemory(systemContext, request, cachedAttachments);
            }
            let interpretResult: InterpretResult;
            try {
                interpretResult = await interpretRequest(
                    context,
                    request,
                    cachedAttachments,
                    history,
                );
            } catch (e: any) {
                if (systemContext.userRequestKnowledgeExtraction === true) {
                    addResultToMemory(
                        systemContext,
                        `Error translating request '${request}': ${e.message}`,
                        DispatcherName,
                    );
                }
                systemContext?.logger?.logEvent("request:exception", {
                    request,
                    message: e.message,
                    stack: e.stack,
                });
                throw e;
            }

            const { requestAction, tokenUsage } = interpretResult;

            if (tokenUsage) {
                const commandResult = getCommandResult(systemContext);
                if (commandResult !== undefined) {
                    commandResult.tokenUsage = tokenUsage;
                }
            }

            // Execute the actions
            await executeActions(
                requestAction.actions,
                requestAction.history?.entities,
                context,
            );

            await requestExplain(
                systemContext,
                cachedAttachments,
                interpretResult,
            );
        } finally {
            profiler?.stop();
        }
    }
    public async getCompletion(
        context: SessionContext<CommandHandlerContext>,
        params: ParsedCommandParams<typeof this.parameters>,
        names: string[],
    ): Promise<CompletionGroup[]> {
        const completions: CompletionGroup[] = [];
        for (const name of names) {
            if (name === "request") {
                const requestPrefix = params.args.request;
                completions.push(
                    ...(await requestCompletion(
                        requestPrefix,
                        context.agentContext,
                    )),
                );
            }
        }
        return completions;
    }
}
