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

import {
    AgentMessageKind,
    AgentThreadHandle,
    AppAgentEvent,
    AppAgent,
    AppAgentManifest,
    DisplayAppendMode,
    SessionContext,
    DisplayContent,
    type GrammarValidationRequest,
    type GrammarValidationResult,
} from "@typeagent/agent-sdk";
import { IAgentMessage, RequestId } from "@typeagent/dispatcher-types";
import { CommandHandlerContext } from "../context/commandHandlerContext.js";
import { IndexData } from "image-memory";
import { IndexManager } from "../context/indexManager.js";
import { validateGrammarPatternsImpl } from "../validation/grammarValidationService.mjs";
import registerDebug from "debug";

const debug = registerDebug("typeagent:dispatcher:sessionContext");
const debugClientCountWarn = registerDebug(
    "typeagent:dispatcher:clientCount:warn",
);
import { randomUUID } from "node:crypto";

export function createSessionContext<T = unknown>(
    name: string,
    agentContext: T,
    context: CommandHandlerContext,
    allowDynamicAgent: boolean,
    sessionContextId: string,
): SessionContext<T> {
    const sessionDirPath = context.session.getSessionDirPath();
    const storageProvider = context.storageProvider;
    const storage = sessionDirPath
        ? storageProvider!.getStorage(name, sessionDirPath)
        : undefined;
    const instanceStorageDir = context.instanceDir ?? context.persistDir;
    const instanceStorage = instanceStorageDir
        ? storageProvider!.getStorage(name, instanceStorageDir)
        : undefined;
    const dynamicAgentNames = new Set<string>();
    const addDynamicAgent = allowDynamicAgent
        ? (agentName: string, manifest: AppAgentManifest, appAgent: AppAgent) =>
              // acquire the lock to prevent change the state while we are processing a command or removing dynamic agent.
              // WARNING: deadlock if this is call because we are processing a request
              context.commandLock(async () => {
                  await context.agents.addDynamicAgent(
                      agentName,
                      manifest,
                      appAgent,
                  );
                  dynamicAgentNames.add(agentName);
                  // Update the enable state to reflect the new agent
                  context.agents.setState(context, context.session.getConfig());
              })
        : () => {
              throw new Error("Permission denied: cannot add dynamic agent");
          };

    const removeDynamicAgent = allowDynamicAgent
        ? (agentName: string) =>
              // acquire the lock to prevent change the state while we are processing a command or adding dynamic agent.
              // WARNING: deadlock if this is called while we are processing a request
              context.commandLock(async () => {
                  if (!dynamicAgentNames.delete(agentName)) {
                      throw new Error(
                          `Permission denied: dynamic agent '${agentName}' not added by this agent`,
                      );
                  }
                  return context.agents.removeAgent(
                      agentName,
                      context.agentCache.grammarStore,
                  );
              })
        : () => {
              throw new Error("Permission denied: cannot remove dynamic agent");
          };

    const forceCleanupDynamicAgent = allowDynamicAgent
        ? (agentName: string) =>
              context.commandLock(async () => {
                  dynamicAgentNames.delete(agentName);
                  return context.agents.forceCleanupAgent(
                      agentName,
                      context.agentCache.grammarStore,
                  );
              })
        : () => {
              throw new Error(
                  "Permission denied: cannot force cleanup dynamic agent",
              );
          };

    const sessionContext: SessionContext<T> = {
        get agentContext() {
            return agentContext;
        },
        get sessionStorage() {
            return storage;
        },
        get instanceStorage() {
            return instanceStorage;
        },
        get sessionContextId() {
            return sessionContextId;
        },
        notify(
            event: AppAgentEvent,
            message: string | DisplayContent,
            notificationId?: string,
        ) {
            context.clientIO.notify(notificationId, event, message, name);
        },
        beginAgentThread(kind: AgentMessageKind): AgentThreadHandle {
            const clientRequestId = `agent-${name}-${randomUUID()}`;
            const requestId: RequestId = {
                requestId: "",
                clientRequestId,
            };
            let completed = false;
            const buildMessage = (content: DisplayContent): IAgentMessage => ({
                message: content,
                requestId,
                source: name,
                kind,
            });
            const completedError = () =>
                new Error(
                    `Agent thread ${clientRequestId} is completed; call beginAgentThread() to start a new thread.`,
                );
            return {
                kind,
                setDisplay(content: DisplayContent) {
                    if (completed) throw completedError();
                    context.clientIO.setDisplay(buildMessage(content));
                },
                appendDisplay(
                    content: DisplayContent,
                    mode: DisplayAppendMode = "block",
                ) {
                    if (completed) throw completedError();
                    context.clientIO.appendDisplay(buildMessage(content), mode);
                },
                complete() {
                    completed = true;
                },
            };
        },
        async toggleTransientAgent(subAgentName: string, enable: boolean) {
            if (!subAgentName.startsWith(`${name}.`)) {
                throw new Error(`Invalid sub agent name: ${subAgentName}`);
            }
            const state = context.agents.getTransientState(subAgentName);
            if (state === undefined) {
                throw new Error(
                    `Transient sub agent not found: ${subAgentName}`,
                );
            }

            if (state === enable) {
                return;
            }

            // acquire the lock to prevent change the state while we are processing a command.
            // WARNING: deadlock if this is call because we are processing a request
            return context.commandLock(async () => {
                context.agents.toggleTransient(subAgentName, enable);
                // Changing active schemas, we need to clear the cache.
                context.translatorCache.clear();
                if (enable) {
                    // REVIEW: is switch current translator the right behavior?
                    context.lastActionSchemaName = subAgentName;
                } else if (context.lastActionSchemaName === subAgentName) {
                    context.lastActionSchemaName = name;
                }
            });
        },
        addDynamicAgent,
        removeDynamicAgent,
        forceCleanupDynamicAgent,
        getSharedLocalHostPort: async (agentName: string) => {
            const localHostPort = await context.agents.getSharedLocalHostPort(
                name,
                agentName,
            );
            if (localHostPort === undefined) {
                throw new Error(
                    `Agent '${agentName}' does not have a shared local host port.`,
                );
            }
            return localHostPort;
        },
        setLocalHostPort(port: number) {
            context.agents.setLocalHostPort(name, port);
        },
        registerPort(role: string, port: number) {
            const regId = context.portRegistrar.register(
                name,
                role,
                port,
                sessionContextId,
            );
            return {
                release: () =>
                    context.portRegistrar.release(regId, sessionContextId),
            };
        },
        indexes(type: string): Promise<any[]> {
            return new Promise<IndexData[]>((resolve, reject) => {
                const iidx: IndexData[] =
                    IndexManager.getInstance().indexes.filter((value) => {
                        return type === "all" || value.source === type;
                    });

                resolve(iidx);
            });
        },
        popupQuestion(
            message: string,
            choices: string[] = ["Yes", "No"], // default choices
            defaultId?: number,
        ): Promise<number> {
            return context.clientIO.question(
                undefined,
                message,
                choices,
                defaultId,
                name,
            );
        },
        async reloadAgentSchema(): Promise<void> {
            await context.agents.reloadAgentSchema(name, context);
        },
        async notifyReadinessChanged(): Promise<void> {
            // Best-effort: swallow errors so an event-driven trigger
            // (e.g. WebSocket onClientConnected) never throws into the
            // event-emitter path. The cache will reconcile on the next
            // refresh.
            try {
                await context.agents.refreshReadiness(name);
            } catch {
                // ignore
            }
        },
        async notifyClientCountChanged(
            role: string,
            count: number,
        ): Promise<void> {
            // Best-effort: counts are informational (surfaced only via
            // `@system ports`) so any failure is swallowed; logged
            // under typeagent:dispatcher:clientCount:warn so it isn't
            // entirely invisible.
            try {
                context.portRegistrar.setClientCount(
                    name,
                    role,
                    sessionContextId,
                    count,
                );
            } catch (e) {
                debugClientCountWarn(
                    `notifyClientCountChanged failed for ${name}/${role}: ${e instanceof Error ? e.message : String(e)}`,
                );
            }
        },
        async validateGrammarPatterns(
            request: GrammarValidationRequest,
        ): Promise<GrammarValidationResult> {
            debug(
                `[${name}] Validating ${request.patterns.length} patterns for "${request.actionName}"`,
            );

            // Use command lock for thread safety
            return context.commandLock(async () => {
                return await validateGrammarPatternsImpl(
                    name,
                    request,
                    context.agentGrammarRegistry,
                    context.agentCache,
                );
            });
        },
    };

    (sessionContext as any).conversationManager = context.conversationManager;
    // Expose CommandHandlerContext for agents that need to execute nested actions
    (sessionContext as any)._systemContext = context;
    return sessionContext;
}
