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

import {
    CommandHandlerTable,
    getCommandInterface,
} from "@typeagent/agent-sdk/helpers/command";
import { RequestCommandHandler } from "./handlers/requestCommandHandler.js";
import { TranslateCommandHandler } from "./handlers/translateCommandHandler.js";
import { ExplainCommandHandler } from "./handlers/explainCommandHandler.js";
import {
    ActionContext,
    ActionResult,
    AppAgent,
    AppAgentManifest,
    TypeAgentAction,
} from "@typeagent/agent-sdk";
import { CommandHandlerContext } from "../commandHandlerContext.js";
import {
    createActionResultFromMarkdownDisplay,
    createActionResultFromTextDisplay,
    createActionResultNoDisplay,
} from "@typeagent/agent-sdk/helpers/action";
import { DispatcherActions } from "./schema/dispatcherActionSchema.js";
import {
    ClarifyRequestAction,
    ClarifyUnresolvedReference,
} from "./schema/clarifyActionSchema.js";
import { loadAgentJsonTranslator } from "../../translation/agentTranslators.js";
import { lookupAndAnswer } from "../../search/search.js";
import {
    LookupAction,
    LookupActivity,
    LookupAndAnswerConversation,
} from "./schema/lookupActionSchema.js";
import { translateRequest } from "../../translation/translateRequest.js";
import { ActivityActions } from "./schema/activityActionSchema.js";
import { ClarifyEntityAction } from "../../execute/pendingActions.js";
import { MatchCommandHandler } from "./handlers/matchCommandHandler.js";
import { DispatcherEmoji } from "./dispatcherUtils.js";
import { getHistoryContext } from "../../translation/interpretRequest.js";
import { ReasoningAction } from "./schema/reasoningActionSchema.js";
import {
    executeReasoningAction as executeClaudeReasoningAction,
    executeReasoning as executeClaudeReasoning,
} from "../../reasoning/claude.js";
import {
    executeReasoningAction as executeCopilotReasoningAction,
    executeReasoning as executeCopilotReasoning,
} from "../../reasoning/copilot.js";
import { ReasonCommandHandler } from "./handlers/reasonCommandHandler.js";

const dispatcherHandlers: CommandHandlerTable = {
    description: "Type Agent Dispatcher Commands",
    commands: {
        request: new RequestCommandHandler(),
        match: new MatchCommandHandler(),
        translate: new TranslateCommandHandler(),
        reason: new ReasonCommandHandler(),
        explain: new ExplainCommandHandler(),
    },
};

async function executeDispatcherAction(
    action: TypeAgentAction<
        | DispatcherActions
        | ClarifyRequestAction
        | LookupAction
        | LookupActivity
        | ActivityActions
        | ClarifyEntityAction
        | ReasoningAction
    >,
    context: ActionContext<CommandHandlerContext>,
) {
    switch (action.schemaName) {
        case "dispatcher.clarify":
            switch (action.actionName) {
                case "clarifyMultiplePossibleActionName":
                case "clarifyMissingParameter":
                case "clarifyMultipleAgentMatches":
                    return clarifyRequestAction(action, context);
                case "clarifyUnresolvedReference":
                    const result = await clarifyWithLookup(action, context);
                    // If we fail to clarify with lookup, just ask the user.
                    return result ?? clarifyRequestAction(action, context);
                case "clarifyEntities":
                    return clarifyEntityAction(action, context);
            }
            break;
        case "dispatcher.lookup":
            switch (action.actionName) {
                case "lookupAndAnswerConversation": {
                    const lookupResult = await lookupAndAnswer(action, context);
                    if (lookupResult.error !== undefined) {
                        // Conversation lookup found nothing — fall back to reasoning
                        const systemContext =
                            context.sessionContext.agentContext;
                        if (!systemContext.noReasoning) {
                            const engine =
                                systemContext.session.getConfig().execution
                                    .reasoning;
                            const reason =
                                engine === "copilot"
                                    ? executeCopilotReasoning
                                    : executeClaudeReasoning;
                            return reason(
                                action.parameters.originalRequest,
                                context,
                            );
                        }
                    }
                    return lookupResult;
                }
                case "startLookup":
                    const location =
                        action.parameters.lookup.source === "internet"
                            ? action.parameters.lookup.site !== undefined
                                ? `on internet sites ${action.parameters.lookup.site.join(", ")}`
                                : `on the internet`
                            : `in the conversation`;

                    const displayText = `Ok. What do you want to look up ${location}?`;
                    const result = createActionResultFromTextDisplay(
                        displayText,
                        displayText,
                    );
                    // TODO: formalize the schema for activityContext
                    result.activityContext = {
                        activityName: "lookup",
                        description: `Looking up ${location}`,
                        state: {
                            ...action.parameters.lookup,
                        },
                    };
                    return result;
            }
            break;
        case "dispatcher.activity":
            switch (action.actionName) {
                case "exitActivity":
                    const result =
                        createActionResultFromTextDisplay("Ok.  What's next?");

                    const systemContext = context.sessionContext.agentContext;
                    const endAction =
                        systemContext.activityContext?.activityEndAction;
                    if (endAction !== undefined) {
                        result.additionalActions = [endAction];
                    }
                    result.activityContext = null; // clear the activity context.

                    return result;
            }
            break;
        case "dispatcher.reasoning":
            if (action.actionName === "reasoningAction") {
                const config =
                    context.sessionContext.agentContext.session.getConfig();
                const engine = config.execution.reasoning;

                switch (engine) {
                    case "claude":
                        return executeClaudeReasoningAction(action, context);
                    case "copilot":
                        return executeCopilotReasoningAction(action, context);
                    default:
                        throw new Error(
                            `Unsupported reasoning engine: ${engine}`,
                        );
                }
            }
            break;
    }

    throw new Error(
        `Unknown dispatcher action: ${action.schemaName}.${action.actionName}`,
    );
}

function clarifyRequestAction(
    action: ClarifyRequestAction,
    context: ActionContext<CommandHandlerContext>,
) {
    const { request, clarifyingQuestion } = action.parameters;
    context.actionIO.appendDisplay({
        type: "text",
        speak: true,
        content: clarifyingQuestion,
    });

    const result = createActionResultNoDisplay(clarifyingQuestion);
    result.additionalInstructions = [
        `Asked the user to clarify the request '${request}'`,
    ];
    return result;
}

// Cache the lookup clarify translator per systemContext since it is expensive
// to create and only depends on the agents + promptLogger for the session.
const lookupClarifyTranslatorCache = new WeakMap<
    CommandHandlerContext,
    ReturnType<typeof loadAgentJsonTranslator>
>();

function getLookupClarifyTranslator(systemContext: CommandHandlerContext) {
    const cached = lookupClarifyTranslatorCache.get(systemContext);
    if (cached !== undefined) {
        return cached;
    }
    const agents = systemContext.agents;
    const actionConfigs = [
        agents.getActionConfig("dispatcher.lookup"),
        agents.getActionConfig("dispatcher"),
    ];
    const translator = loadAgentJsonTranslator(
        actionConfigs,
        [],
        agents,
        undefined,
        undefined,
        undefined,
        systemContext.promptLogger,
    );
    lookupClarifyTranslatorCache.set(systemContext, translator);
    return translator;
}

async function clarifyWithLookup(
    action: ClarifyUnresolvedReference,
    context: ActionContext<CommandHandlerContext>,
): Promise<ActionResult | undefined> {
    const systemContext = context.sessionContext.agentContext;
    const agents = systemContext.agents;
    if (
        !agents.isSchemaActive("dispatcher.lookup") ||
        !agents.isActionActive("dispatcher.lookup")
    ) {
        // lookup is disabled either for translation or action. Just ask the user.
        return undefined;
    }

    const translator = getLookupClarifyTranslator(systemContext);

    const question = `What is ${action.parameters.reference}?`;
    const result = await translator.translate(question);

    if (!result.success) {
        return undefined;
    }
    const lookupAction = result.data as LookupAndAnswerConversation;
    if (lookupAction.actionName !== "lookupAndAnswerConversation") {
        return undefined;
    }
    const lookupResult = await lookupAndAnswer(
        lookupAction as LookupAndAnswerConversation,
        context,
    );

    if (
        lookupResult.error !== undefined ||
        lookupResult.historyText === undefined
    ) {
        return undefined;
    }

    // TODO: This translation can probably more scoped based on the `actionName` field.
    const history = getHistoryContext(systemContext);
    if (history === undefined) {
        // Can't do clarify without history.
        return undefined;
    }

    history.promptSections.push({
        role: "assistant",
        content: lookupResult.historyText,
    });

    const translationResult = await translateRequest(
        context,
        action.parameters.request,
        history,
    );

    if (!translationResult) {
        // undefined means not found or not translated
        // null means cancelled because of replacement parse error.
        return undefined;
    }

    if (translationResult.requestAction.actions.length > 1) {
        // REVIEW: Expect only one action?.
        return undefined;
    }

    return {
        additionalActions: [translationResult.requestAction.actions[0].action],
        entities: [],
    };
}

function clarifyEntityAction(
    action: ClarifyEntityAction,
    context: ActionContext<CommandHandlerContext>,
): ActionResult {
    const { type, name, result } = action.parameters;

    if (result.entities.length < 2) {
        throw new Error(
            "ClarifyEntityAction should not be called with empty or single entity.",
        );
    }

    const question = [
        `Multiple ${type.toLowerCase()} entities named '${name}' were found`,
        ...result.entities.map((entity) => `- ${entity.name}`),
        "", // markdown needs an extra line to start a new paragraph
        `Please clarify which one you meant.`,
    ];

    return createActionResultFromMarkdownDisplay(
        question,
        undefined,
        result.entities,
    );
}

export const dispatcherManifest: AppAgentManifest = {
    emojiChar: DispatcherEmoji,
    description: "Built-in agent to dispatch requests",
    schema: {
        description: "",
        schemaType: "DispatcherActions",
        schemaFile: "./src/context/dispatcher/schema/dispatcherActionSchema.ts",
        injected: true,
        cached: false,
    },
    subActionManifests: {
        clarify: {
            schema: {
                description: "Action that helps you clarify your request.",
                schemaFile:
                    "./src/context/dispatcher/schema/clarifyActionSchema.ts",
                schemaType: "ClarifyRequestAction",
                injected: true,
                cached: false,
            },
        },
        lookup: {
            schema: {
                description:
                    "Action that helps you look up information to answer user questions.",
                schemaFile:
                    "./src/context/dispatcher/schema/lookupActionSchema.ts",
                schemaType: {
                    action: "LookupAction",
                    activity: "LookupActivity",
                },
                injected: true,
                cached: false,
            },
        },
        activity: {
            transient: true,
            schema: {
                description: "Action that manages activity context.",
                schemaFile:
                    "./src/context/dispatcher/schema/activityActionSchema.ts",
                schemaType: "ActivityActions",
                injected: true,
                cached: false,
            },
        },
        reasoning: {
            schema: {
                description:
                    "Action that helps you reason through requests that require multiple steps.",
                schemaFile:
                    "./src/context/dispatcher/schema/reasoningActionSchema.ts",
                schemaType: "ReasoningAction",
                injected: true,
                cached: false,
            },
        },
    },
};

export const dispatcherAgent: AppAgent = {
    executeAction: executeDispatcherAction,
    ...getCommandInterface(dispatcherHandlers),
};
