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

import {
    ExecutableAction,
    FullAction,
    getFullActionName,
    PromptEntity,
} from "agent-cache";
import {
    type CommandHandlerContext,
    getCommandResult,
    getRequestId,
} from "../context/commandHandlerContext.js";
import registerDebug from "debug";
import { getAppAgentName } from "../translation/agentTranslators.js";
import {
    ActionResult,
    ActionContext,
    ParsedCommandParams,
    ParameterDefinitions,
    AppAction,
} from "@typeagent/agent-sdk";
import {
    createActionResult,
    actionResultToString,
    createActionResultFromError,
} from "@typeagent/agent-sdk/helpers/action";
import {
    displayError,
    displayStatus,
    displayWarn,
} from "@typeagent/agent-sdk/helpers/display";
import { IncrementalJsonValueCallBack } from "typechat-utils";
import { ProfileNames } from "../utils/profileNames.js";
import { UnknownAction } from "../context/dispatcher/schema/dispatcherActionSchema.js";
import {
    DispatcherName,
    isUnknownAction,
} from "../context/dispatcher/dispatcherUtils.js";
import { isPendingRequestAction } from "../translation/pendingRequest.js";
import {
    isSwitchEnabled,
    translatePendingRequestAction,
} from "../translation/translateRequest.js";
import { validateAction } from "@typeagent/action-schema";
import {
    PendingAction,
    resolveEntities,
    toPendingActions,
} from "./pendingActions.js";
import { getActionContext } from "./actionContext.js";
import {
    addActionResultToMemory,
    addResultToMemory,
} from "../context/memory.js";
import { setActivityContext } from "./activityContext.js";
import { tryGetActionSchema } from "../translation/actionSchemaFileCache.js";
import { processFlow } from "./flowInterpreter.js";

const debugActions = registerDebug("typeagent:dispatcher:actions");
const debugCommandExecError = registerDebug(
    "typeagent:dispatcher:command:exec:error",
);
export function getSchemaNamePrefix(
    schemaName: string,
    systemContext: CommandHandlerContext,
) {
    const config = systemContext.agents.getActionConfig(schemaName);
    return `[${config.emojiChar} ${schemaName}] `;
}

function getStreamingActionContext(
    appAgentName: string,
    actionIndex: number,
    systemContext: CommandHandlerContext,
    fullAction: FullAction,
) {
    const actionContext = systemContext.streamingActionContext;
    systemContext.streamingActionContext = undefined;

    if (
        actionContext === undefined ||
        actionContext.actionIndex !== actionIndex
    ) {
        actionContext?.closeActionContext();
        return undefined;
    }

    // If we are reusing the streaming action context, we need to update the action.
    systemContext.clientIO.setDisplayInfo(
        getRequestId(systemContext),
        appAgentName,
        actionIndex,
        fullAction,
    );
    return actionContext;
}

// REVIEW: don't export this
export async function executeAction(
    executableAction: ExecutableAction,
    context: ActionContext<CommandHandlerContext>,
    actionIndex: number,
): Promise<ActionResult> {
    const action = executableAction.action;
    if (debugActions.enabled) {
        debugActions(
            `Executing action: ${JSON.stringify(action, undefined, 2)}`,
        );
    }

    const schemaName = action.schemaName;
    // For nested action calls (e.g., from TaskFlow scripts), agentContext may be
    // the agent's own context rather than CommandHandlerContext. In that case,
    // use _systemContext which exposes the CommandHandlerContext.
    const sessionCtx = context.sessionContext as any;
    const systemContext: CommandHandlerContext =
        sessionCtx._systemContext ?? sessionCtx.agentContext;

    const appAgentName = getAppAgentName(schemaName);
    const requestId = getRequestId(systemContext);
    const appAgent = systemContext.agents.getAppAgent(appAgentName);

    // Update the last action translator.
    systemContext.lastActionSchemaName = schemaName;

    // Reuse the same streaming action context if one is available.

    const { actionContext, closeActionContext } =
        getStreamingActionContext(
            appAgentName,
            actionIndex,
            systemContext,
            action,
        ) ??
        getActionContext(
            appAgentName,
            systemContext,
            requestId,
            actionIndex,
            action,
        );

    const prefix = getSchemaNamePrefix(action.schemaName, systemContext);
    displayStatus(
        `${prefix}Executing action ${getFullActionName(executableAction)}`,
        context,
    );
    actionContext.profiler = systemContext.commandProfiler?.measure(
        ProfileNames.executeAction,
        true,
        actionIndex,
    );
    let result: ActionResult;
    try {
        // Check if this action has a registered flow program
        const flowDef = systemContext.agents.getFlow(
            schemaName,
            action.actionName,
        );
        if (flowDef !== undefined) {
            const flowParams = (action.parameters ?? {}) as Record<
                string,
                unknown
            >;
            // Pass the outer context (ActionContext<CommandHandlerContext>) so the
            // flow interpreter can access the full agent registry.
            result = await processFlow(
                flowDef,
                flowParams,
                context,
                actionIndex,
            );
        } else {
            if (appAgent.executeAction === undefined) {
                throw new Error(
                    `Agent '${appAgentName}' does not support executeAction.`,
                );
            }
            result =
                (await appAgent.executeAction(action, actionContext)) ??
                createActionResult(
                    `Action ${getFullActionName(executableAction)} completed.`,
                );
        }
    } catch (e: any) {
        if (e.name === "AbortError") {
            throw e;
        }
        result = createActionResultFromError(e.message);
    }
    actionContext.profiler?.stop();
    actionContext.profiler = undefined;

    if (debugActions.enabled) {
        debugActions(actionResultToString(result));
    }

    // Display the action result.
    if (result.error !== undefined) {
        if (!("fallbackToReasoning" in result) || !result.fallbackToReasoning) {
            displayError(result.error, actionContext);
        }
    } else {
        if (result.displayContent !== undefined) {
            actionContext.actionIO.appendDisplay(
                result.displayContent,
                "block",
            );
        }
        if (result.dynamicDisplayId !== undefined) {
            systemContext.clientIO.setDynamicDisplay(
                requestId,
                schemaName,
                actionIndex,
                result.dynamicDisplayId,
                result.dynamicDisplayNextRefreshMs!,
            );
        }
        if (result.pendingChoice !== undefined) {
            const pc = result.pendingChoice;
            systemContext.pendingChoiceRoutes.set(pc.choiceId, {
                agentName: appAgentName,
                requestId,
                actionIndex,
            });
            systemContext.clientIO.requestChoice(
                requestId,
                pc.choiceId,
                pc.type,
                pc.message,
                pc.type === "multiChoice" ? pc.choices : [],
                schemaName,
            );
        }
    }

    closeActionContext();
    return result;
}

async function canExecute(
    actions: ExecutableAction[],
    context: ActionContext<CommandHandlerContext>,
): Promise<boolean> {
    const sessionCtx = context.sessionContext as any;
    const systemContext: CommandHandlerContext =
        sessionCtx._systemContext ?? sessionCtx.agentContext;
    const unknown: UnknownAction[] = [];
    const disabled = new Set<string>();
    for (const { action } of actions) {
        if (isUnknownAction(action)) {
            unknown.push(action);
        }
        if (
            action.schemaName &&
            !systemContext.agents.isActionActive(action.schemaName)
        ) {
            disabled.add(action.schemaName);
        }
    }

    if (unknown.length > 0) {
        const unknownRequests = unknown.map(
            (action) => action.parameters.request,
        );
        const lines = [
            `Unable to determine ${actions.length > 1 ? "one or more actions in" : "action for"} the request.`,
            ...unknownRequests.map((s) => `- ${s}`),
        ];
        addResultToMemory(systemContext, lines.join("\n"), DispatcherName);

        const config = systemContext.session.getConfig();
        if (!isSwitchEnabled(config)) {
            lines.push("");
            lines.push("Switching agents is disabled");
        } else {
            const entries = await Promise.all(
                unknownRequests.map((request) =>
                    systemContext.agents.semanticSearchActionSchema(
                        request,
                        1,
                        () => true, // don't filter
                    ),
                ),
            );
            const schemaNames = new Set(
                entries
                    .filter((e) => e !== undefined)
                    .map((e) => e![0].item.actionSchemaFile.schemaName)
                    .filter(
                        (schemaName) =>
                            !systemContext.agents.isSchemaActive(schemaName),
                    ),
            );

            if (schemaNames.size > 0) {
                lines.push("");
                lines.push(
                    `Possible agent${schemaNames.size > 1 ? "s" : ""} to handle the request${unknownRequests.length > 1 ? "s" : ""} are not active: ${Array.from(schemaNames).join(", ")}`,
                );
            }
        }

        displayError(lines, context);
        return false;
    }

    if (disabled.size > 0) {
        const message = `Not executed. Action disabled for ${Array.from(disabled.values()).join(", ")}`;
        addResultToMemory(systemContext, message, DispatcherName);
        displayWarn(message, context);
        return false;
    }

    return true;
}

export type ActionExecutionError = {
    error: string;
    failedAction: ExecutableAction;
    fallbackToReasoning?: boolean;
};

export async function executeActions(
    actions: ExecutableAction[],
    entities: PromptEntity[] | undefined,
    context: ActionContext<CommandHandlerContext>,
): Promise<ActionExecutionError | undefined> {
    const sessionCtx = context.sessionContext as any;
    const systemContext: CommandHandlerContext =
        sessionCtx._systemContext ?? sessionCtx.agentContext;
    const commandResult = getCommandResult(systemContext);
    if (commandResult !== undefined) {
        commandResult.actions = actions.map(({ action }) => action);
    }

    // Even if the action is not executed, resolve the entities for the commandResult.
    const actionQueue: PendingAction[] = await toPendingActions(
        context,
        actions,
        entities,
    );

    if (!(await canExecute(actions, context))) {
        return;
    }

    debugActions(`Executing actions: ${JSON.stringify(actions, undefined, 2)}`);
    let actionIndex = 0;
    while (actionQueue.length !== 0) {
        systemContext.currentAbortSignal?.throwIfAborted();
        const pending = actionQueue.shift()!;
        const executableAction = pending.executableAction;

        const action = executableAction.action;

        if (isPendingRequestAction(action)) {
            const translationResult = await translatePendingRequestAction(
                action,
                context,
                actionIndex,
            );

            const requestAction = translationResult.requestAction;
            actionQueue.unshift(
                ...(await toPendingActions(
                    context,
                    requestAction.actions,
                    requestAction.history?.entities,
                )),
            );
            continue;
        }
        const appAgentName = getAppAgentName(action.schemaName);
        // resolve result entities.
        const resultEntityResolver = pending.resultEntityResolver;
        let resolvedEntities = pending.resolvedEntities;
        if (resultEntityResolver !== undefined) {
            const resultResolvedEntities = await resolveEntities(
                systemContext.agents,
                action,
                resultEntityResolver,
            );
            if (resultResolvedEntities !== undefined) {
                if (resolvedEntities === undefined) {
                    resolvedEntities = [];
                }
                resolvedEntities.push(...resultResolvedEntities);
            }
        }
        const result = await executeAction(
            executableAction,
            context,
            actionIndex,
        );

        // add the action result to memory whether it has error or not.
        if (
            systemContext.actionResultEntityStorage ||
            systemContext.actionResultKnowledgeExtraction
        ) {
            addActionResultToMemory(
                systemContext,
                executableAction,
                resolvedEntities,
                action.schemaName,
                result,
            );
        }

        if (result.error !== undefined) {
            // Stop executing further action on error.
            return { error: result.error, failedAction: executableAction };
        }

        const resultEntityId = executableAction.resultEntityId;
        if (resultEntityId !== undefined) {
            if (result.resultEntity === undefined) {
                throw new Error(
                    `Action ${getFullActionName(
                        executableAction,
                    )} did not return a result entity.`,
                );
            }
            if (resultEntityResolver === undefined) {
                throw new Error(
                    `Internal error: resultEntityResolver is undefined`,
                );
            }
            resultEntityResolver.setResultEntity(
                `\${result-${resultEntityId}}`,
                {
                    ...result.resultEntity,
                    sourceAppAgentName: appAgentName,
                },
            );
        }

        if (result.activityContext !== undefined) {
            if (actionQueue.length > 0) {
                throw new Error(
                    `Cannot change activity context when there are pending actions.`,
                );
            }

            debugActions(
                `Result activity context: ${JSON.stringify(result.activityContext, undefined, 2)}`,
            );
            const prevActivityContext = systemContext.activityContext;
            const openLocalView = setActivityContext(
                action.schemaName,
                result.activityContext,
                systemContext,
            );
            if (openLocalView !== undefined) {
                if (openLocalView) {
                    const port =
                        systemContext.agents.getLocalHostPort(appAgentName);
                    if (port !== undefined) {
                        await systemContext.clientIO.openLocalView(
                            getRequestId(systemContext),
                            port,
                        );
                    }
                } else if (prevActivityContext !== undefined) {
                    const port = systemContext.agents.getLocalHostPort(
                        prevActivityContext.appAgentName,
                    );
                    if (port !== undefined) {
                        await systemContext.clientIO.closeLocalView(
                            getRequestId(systemContext),
                            port,
                        );
                    }
                }
            }
            if (systemContext.activityContext !== undefined) {
                debugActions(
                    `Starting activity: ${JSON.stringify(systemContext.activityContext, undefined, 2)}`,
                );
            } else if (prevActivityContext !== undefined) {
                // Activity context cleared.
                debugActions(
                    `Stopped activity: ${JSON.stringify(prevActivityContext, undefined, 2)}`,
                );
            }
        }

        if (result.additionalActions !== undefined) {
            try {
                const actions = getAdditionalExecutableActions(
                    result.additionalActions,
                    action.schemaName,
                    systemContext,
                );
                // REVIEW: assume that the agent will fill the entities already?  Also, current format doesn't support resultEntityIds.
                actionQueue.unshift(
                    ...(await toPendingActions(context, actions, undefined)),
                );
            } catch (e) {
                throw new Error(
                    `${action.schemaName}.${action.actionName} returned an invalid action: ${e}`,
                );
            }
        }
        actionIndex++;
    }
    return undefined;
}

function getAdditionalExecutableActions(
    actions: AppAction[],
    sourceSchemaName: string,
    context: CommandHandlerContext,
) {
    const appAgentName = getAppAgentName(sourceSchemaName);
    const executableActions: ExecutableAction[] = [];
    for (const newAction of actions) {
        const fullAction = (
            newAction.schemaName !== undefined
                ? newAction
                : {
                      ...newAction,
                      schemaName: sourceSchemaName,
                  }
        ) as FullAction;

        if (appAgentName !== DispatcherName) {
            // For non-dispatcher, action can only be triggered within the same agent,
            // with the exception of the reasoning action which is a meta-action that
            // any agent may request to hand off complex tasks to the reasoning loop.
            const actionAppAgentName = getAppAgentName(fullAction.schemaName);
            const isReasoningAction =
                actionAppAgentName === DispatcherName &&
                fullAction.actionName === "reasoningAction";
            if (actionAppAgentName !== appAgentName && !isReasoningAction) {
                throw new Error(
                    `Cannot invoke actions from other agent '${actionAppAgentName}'.`,
                );
            }
        }

        const actionInfo = tryGetActionSchema(fullAction, context.agents);
        if (actionInfo === undefined) {
            throw new Error(
                `Action not found ${fullAction.schemaName}.${fullAction.actionName}`,
            );
        }
        validateAction(actionInfo, fullAction);

        executableActions.push({ action: fullAction });
    }
    return executableActions;
}

export function startStreamPartialAction(
    schemaName: string,
    actionName: string,
    context: CommandHandlerContext,
    actionIndex: number,
): IncrementalJsonValueCallBack {
    const appAgentName = getAppAgentName(schemaName);
    const appAgent = context.agents.getAppAgent(appAgentName);
    if (appAgent.streamPartialAction === undefined) {
        // The config declared that there are streaming action, but the agent didn't implement it.
        throw new Error(
            `Agent '${appAgentName}' does not support streamPartialAction.`,
        );
    }

    const actionContextWithClose = getActionContext(
        appAgentName,
        context,
        getRequestId(context),
        actionIndex,
        {
            schemaName,
            actionName,
        },
    );

    context.streamingActionContext = actionContextWithClose;

    return (name: string, value: any, delta?: string) => {
        appAgent.streamPartialAction!(
            actionName,
            name,
            value,
            delta,
            actionContextWithClose.actionContext,
        );
    };
}

export async function executeCommand(
    commands: string[],
    params: ParsedCommandParams<ParameterDefinitions> | undefined,
    appAgentName: string,
    context: CommandHandlerContext,
    attachments?: string[],
): Promise<void> {
    const appAgent = context.agents.getAppAgent(appAgentName);
    if (appAgent.executeCommand === undefined) {
        throw new Error(
            `Agent '${appAgentName}' does not support executeCommand.`,
        );
    }

    // update the last action name
    const { actionContext, closeActionContext } = getActionContext(
        appAgentName,
        context,
        getRequestId(context),
        undefined,
        commands,
    );

    try {
        actionContext.profiler = context.commandProfiler?.measure(
            ProfileNames.executeCommand,
            true,
        );

        try {
            return await appAgent.executeCommand(
                commands,
                params,
                actionContext,
                attachments,
            );
        } catch (e: any) {
            if (e.name === "AbortError") {
                throw e;
            }
            displayError(`ERROR: ${e.message}`, actionContext);
            debugCommandExecError(e.stack);
        }
    } finally {
        actionContext.profiler?.stop();
        actionContext.profiler = undefined;
        closeActionContext();
    }
}
