import type { AgentToolResult, AgentToolUpdateCallback, ExtensionContext } from "@earendil-works/pi-coding-agent";
import type { McpExtensionState } from "./state.ts";
import type { DirectToolSpec, McpConfig, McpContent, McpResource, McpTool, ServerEntry } from "./types.ts";
import type { MetadataCache, ServerCacheEntry } from "./metadata-cache.ts";
import { lazyConnect, getFailureAgeSeconds } from "./init.ts";
import { computeServerHash, isServerCacheValid, saveMetadataCache, serializeResources, serializeTools } from "./metadata-cache.ts";
import { formatSchema } from "./tool-metadata.ts";
import { transformMcpContent } from "./tool-registrar.ts";
import { maybeStartUiSession, type UiSessionRuntime } from "./ui-session.ts";
import { formatToolName, isToolExcluded } from "./types.ts";
import { resourceNameToToolName } from "./resource-tools.ts";
import { authenticate, supportsOAuth } from "./mcp-auth-flow.ts";
import { formatAuthRequiredMessage } from "./utils.ts";
import { McpServerManager } from "./server-manager.ts";
import { logger } from "./logger.ts";

const BUILTIN_NAMES = new Set(["read", "bash", "edit", "write", "grep", "find", "ls", "mcp"]);

type DirectAutoAuthResult =
  | { status: "skipped" }
  | { status: "success" }
  | { status: "failed"; message: string };
type DirectToolFilter = true | string[] | false;
type ParsedDirectToolEnvOverride = {
  servers: Set<string>;
  toolsByServer: Map<string, Set<string>>;
};
type DirectToolMetadataManager = {
  connect: (serverName: string, definition: ServerEntry) => Promise<{
    status: "connected" | "closed" | "needs-auth";
    tools: McpTool[];
    resources: McpResource[];
  }>;
  closeAll: () => Promise<void>;
};

function getDirectAuthRequiredMessage(
  state: McpExtensionState,
  serverName: string,
  defaultMessage = `MCP server "${serverName}" requires OAuth authentication. Run /mcp-auth ${serverName} first.`,
): string {
  return formatAuthRequiredMessage(state.config, serverName, defaultMessage);
}

function getDirectAuthFailedMessage(state: McpExtensionState, serverName: string, message: string): string {
  const customGuidance = state.config.settings?.authRequiredMessage;
  if (customGuidance) {
    return `OAuth authentication failed for "${serverName}": ${message}. ${getDirectAuthRequiredMessage(state, serverName)}`;
  }
  return `OAuth authentication failed for "${serverName}": ${message}. Run /mcp-auth ${serverName} first.`;
}

function parseDirectToolEnvOverride(envOverride?: string[]): ParsedDirectToolEnvOverride | undefined {
  if (!envOverride) return undefined;

  const servers = new Set<string>();
  const toolsByServer = new Map<string, Set<string>>();

  for (let item of envOverride) {
    item = item.replace(/\/+$/, "");
    if (item.includes("/")) {
      const [server, tool] = item.split("/", 2);
      if (server && tool) {
        if (!toolsByServer.has(server)) toolsByServer.set(server, new Set());
        toolsByServer.get(server)!.add(tool);
      } else if (server) {
        servers.add(server);
      }
    } else if (item) {
      servers.add(item);
    }
  }

  return {
    servers,
    toolsByServer,
  };
}

function getDirectToolFilter(
  serverName: string,
  definition: ServerEntry,
  globalDirect: McpConfig["settings"]["directTools"],
  envOverride?: ParsedDirectToolEnvOverride,
): DirectToolFilter {
  if (envOverride) {
    if (envOverride.servers.has(serverName)) {
      return true;
    }

    const tools = envOverride.toolsByServer.get(serverName);
    return tools ? [...tools] : false;
  }

  if (definition.directTools !== undefined) {
    return definition.directTools;
  }

  return globalDirect || false;
}

async function attemptDirectAutoAuth(
  state: McpExtensionState,
  serverName: string,
): Promise<DirectAutoAuthResult> {
  if (state.config.settings?.autoAuth !== true) {
    return { status: "skipped" };
  }

  const definition = state.config.mcpServers[serverName];
  if (!definition || !supportsOAuth(definition) || !definition.url) {
    return { status: "skipped" };
  }

  const oauth = definition.oauth && typeof definition.oauth === "object" ? definition.oauth : undefined;
  const grantType = oauth?.grantType ?? "authorization_code";
  if (!state.ui && grantType !== "client_credentials") {
    return {
      status: "failed",
      message: getDirectAuthRequiredMessage(
        state,
        serverName,
        `MCP server "${serverName}" requires OAuth authentication. Run /mcp-auth ${serverName} in an interactive session.`,
      ),
    };
  }

  try {
    await authenticate(serverName, definition.url, definition);
    return { status: "success" };
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    return {
      status: "failed",
      message: getDirectAuthFailedMessage(state, serverName, message),
    };
  }
}

export function resolveDirectTools(
  config: McpConfig,
  cache: MetadataCache | null,
  prefix: "server" | "none" | "short",
  envOverride?: string[],
): DirectToolSpec[] {
  const specs: DirectToolSpec[] = [];
  if (!cache) return specs;

  const seenNames = new Set<string>();
  const globalDirect = config.settings?.directTools;
  const parsedEnvOverride = parseDirectToolEnvOverride(envOverride);

  for (const [serverName, definition] of Object.entries(config.mcpServers)) {
    const serverCache = cache.servers[serverName];
    if (!serverCache || !isServerCacheValid(serverCache, definition)) continue;

    const toolFilter = getDirectToolFilter(serverName, definition, globalDirect, parsedEnvOverride);

    if (!toolFilter) continue;

    for (const tool of serverCache.tools ?? []) {
      if (toolFilter !== true && !toolFilter.includes(tool.name)) continue;
      if (isToolExcluded(tool.name, serverName, prefix, definition.excludeTools)) continue;
      const prefixedName = formatToolName(tool.name, serverName, prefix);
      if (BUILTIN_NAMES.has(prefixedName)) {
        console.warn(`MCP: skipping direct tool "${prefixedName}" (collides with builtin)`);
        continue;
      }
      if (seenNames.has(prefixedName)) {
        console.warn(`MCP: skipping duplicate direct tool "${prefixedName}" from "${serverName}"`);
        continue;
      }
      seenNames.add(prefixedName);
      specs.push({
        serverName,
        originalName: tool.name,
        prefixedName,
        description: tool.description ?? "",
        inputSchema: tool.inputSchema,
        uiResourceUri: tool.uiResourceUri,
        uiStreamMode: tool.uiStreamMode,
      });
    }

    if (definition.exposeResources !== false) {
      for (const resource of serverCache.resources ?? []) {
        const baseName = `get_${resourceNameToToolName(resource.name)}`;
        if (toolFilter !== true && !toolFilter.includes(baseName)) continue;
        if (isToolExcluded(baseName, serverName, prefix, definition.excludeTools)) continue;
        const prefixedName = formatToolName(baseName, serverName, prefix);
        if (BUILTIN_NAMES.has(prefixedName)) {
          console.warn(`MCP: skipping direct resource tool "${prefixedName}" (collides with builtin)`);
          continue;
        }
        if (seenNames.has(prefixedName)) {
          console.warn(`MCP: skipping duplicate direct resource tool "${prefixedName}" from "${serverName}"`);
          continue;
        }
        seenNames.add(prefixedName);
        specs.push({
          serverName,
          originalName: baseName,
          prefixedName,
          description: resource.description ?? `Read resource: ${resource.uri}`,
          resourceUri: resource.uri,
        });
      }
    }
  }

  return specs;
}

export function getMissingConfiguredDirectToolServers(
  config: McpConfig,
  cache: MetadataCache | null,
  envOverride?: string[],
): string[] {
  const missing: string[] = [];
  const globalDirect = config.settings?.directTools;
  const parsedEnvOverride = parseDirectToolEnvOverride(envOverride);

  for (const [serverName, definition] of Object.entries(config.mcpServers)) {
    const hasDirectTools = !!getDirectToolFilter(serverName, definition, globalDirect, parsedEnvOverride);

    if (!hasDirectTools) continue;

    const serverCache = cache?.servers?.[serverName];
    if (!serverCache || !isServerCacheValid(serverCache, definition)) {
      missing.push(serverName);
    }
  }

  return missing;
}

function buildLiveServerCacheEntry(definition: ServerEntry, tools: McpTool[], resources: McpResource[]): ServerCacheEntry {
  return {
    configHash: computeServerHash(definition),
    tools: serializeTools(tools),
    resources: definition.exposeResources === false ? [] : serializeResources(resources),
    cachedAt: Date.now(),
  };
}

export async function resolveDirectToolsWithLiveMetadata(
  config: McpConfig,
  cache: MetadataCache | null,
  prefix: "server" | "none" | "short",
  envOverride?: string[],
  options: {
    createManager?: () => DirectToolMetadataManager;
    saveCache?: (cache: MetadataCache) => void;
  } = {},
): Promise<DirectToolSpec[]> {
  const liveCache: MetadataCache = {
    version: cache?.version ?? 1,
    servers: {
      ...(cache?.servers ?? {}),
    },
  };
  const missingServers = getMissingConfiguredDirectToolServers(config, liveCache, envOverride);

  if (missingServers.length === 0) {
    return resolveDirectTools(config, liveCache, prefix, envOverride);
  }

  const manager = options.createManager?.() ?? new McpServerManager();
  const writeCache = options.saveCache ?? saveMetadataCache;

  try {
    await Promise.all(
      missingServers.map(async (serverName) => {
        const definition = config.mcpServers[serverName];
        try {
          const connection = await manager.connect(serverName, definition);
          if (connection.status !== "connected") {
            logger.debug(`MCP: direct tool metadata unavailable for ${serverName}: ${connection.status}`);
            return;
          }

          const entry = buildLiveServerCacheEntry(definition, connection.tools, connection.resources);
          liveCache.servers[serverName] = entry;
          writeCache({ version: 1, servers: { [serverName]: entry } });
        } catch (error) {
          const message = error instanceof Error ? error.message : String(error);
          logger.debug(`MCP: direct tool metadata preload failed for ${serverName}: ${message}`);
        }
      })
    );
  } finally {
    await manager.closeAll();
  }

  return resolveDirectTools(config, liveCache, prefix, envOverride);
}

export function buildProxyDescription(
  config: McpConfig,
  cache: MetadataCache | null,
  directSpecs: DirectToolSpec[],
): string {
  const prefix = config.settings?.toolPrefix ?? "server";
  let desc = `MCP gateway - connect to MCP servers and call their tools. Non-MCP Pi tools should be called directly, not through mcp.\n`;

  const directByServer = new Map<string, number>();
  for (const spec of directSpecs) {
    directByServer.set(spec.serverName, (directByServer.get(spec.serverName) ?? 0) + 1);
  }
  if (directByServer.size > 0) {
    const parts = [...directByServer.entries()].map(
      ([server, count]) => `${server} (${count})`,
    );
    desc += `\nDirect tools available (call as normal tools): ${parts.join(", ")}\n`;
  }

  const serverSummaries: string[] = [];
  for (const serverName of Object.keys(config.mcpServers)) {
    const entry = cache?.servers?.[serverName];
    const definition = config.mcpServers[serverName];
    const toolCount = (entry?.tools ?? []).filter(
      (tool) => !isToolExcluded(tool.name, serverName, prefix, definition.excludeTools),
    ).length;
    const resourceCount = definition?.exposeResources !== false
      ? (entry?.resources ?? []).filter((resource) => {
          const baseName = `get_${resourceNameToToolName(resource.name)}`;
          return !isToolExcluded(baseName, serverName, prefix, definition.excludeTools);
        }).length
      : 0;
    const totalItems = toolCount + resourceCount;
    if (totalItems === 0) continue;
    const directCount = directByServer.get(serverName) ?? 0;
    const proxyCount = totalItems - directCount;
    if (proxyCount > 0) {
      serverSummaries.push(`${serverName} (${proxyCount} tools)`);
    }
  }

  if (serverSummaries.length > 0) {
    desc += `\nServers: ${serverSummaries.join(", ")}\n`;
  }

  desc += `\nUsage:\n`;
  desc += `  mcp({ })                              → Show server status\n`;
  desc += `  mcp({ server: "name" })               → List tools from server\n`;
  desc += `  mcp({ search: "query" })              → Search MCP tools by name/description\n`;
  desc += `  mcp({ describe: "tool_name" })        → Show tool details and parameters\n`;
  desc += `  mcp({ connect: "server-name" })       → Connect to a server and refresh metadata\n`;
  desc += `  mcp({ tool: "name", args: '{"key": "value"}' })    → Call a tool (args is JSON string)\n`;
  desc += `  mcp({ action: "ui-messages" })        → Retrieve accumulated messages from completed UI sessions\n`;
  desc += `\nMode: tool (call) > connect > describe > search > server (list) > action > nothing (status)`;

  return desc;
}

type DirectToolExecute = (
  toolCallId: string,
  params: Record<string, unknown>,
  signal: AbortSignal | undefined,
  onUpdate: AgentToolUpdateCallback<Record<string, unknown>> | undefined,
  ctx: ExtensionContext,
) => Promise<AgentToolResult<Record<string, unknown>>>;

export function createDirectToolExecutor(
  getState: () => McpExtensionState | null,
  getInitPromise: () => Promise<McpExtensionState> | null,
  spec: DirectToolSpec
): DirectToolExecute {
  return async function execute(_toolCallId, params) {
    let state = getState();
    const initPromise = getInitPromise();

    if (!state && initPromise) {
      try {
        state = await initPromise;
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        return {
          content: [{ type: "text" as const, text: `MCP initialization failed: ${message}` }],
          details: { error: "init_failed", message },
        };
      }
    }
    if (!state) {
      return {
        content: [{ type: "text" as const, text: "MCP not initialized" }],
        details: { error: "not_initialized" },
      };
    }

    let connected = await lazyConnect(state, spec.serverName);
    let autoAuthAttempted = false;

    if (!connected && state.manager.getConnection(spec.serverName)?.status === "needs-auth") {
      autoAuthAttempted = true;
      const autoAuth = await attemptDirectAutoAuth(state, spec.serverName);
      if (autoAuth.status === "failed") {
        return {
          content: [{ type: "text" as const, text: autoAuth.message }],
          details: { error: "auth_required", server: spec.serverName, message: autoAuth.message },
        };
      }
      if (autoAuth.status === "success") {
        await state.manager.close(spec.serverName);
        state.failureTracker.delete(spec.serverName);
        connected = await lazyConnect(state, spec.serverName);
      }
    }

    if (!connected) {
      const authConnection = state.manager.getConnection(spec.serverName);
      if (authConnection?.status === "needs-auth") {
        const message = getDirectAuthRequiredMessage(state, spec.serverName);
        return {
          content: [{ type: "text" as const, text: message }],
          details: { error: "auth_required", server: spec.serverName, message, autoAuthAttempted },
        };
      }
      const failedAgo = getFailureAgeSeconds(state, spec.serverName);
      return {
        content: [{ type: "text" as const, text: `MCP server "${spec.serverName}" not available${failedAgo !== null ? ` (failed ${failedAgo}s ago)` : ""}` }],
        details: { error: "server_unavailable", server: spec.serverName },
      };
    }

    const connection = state.manager.getConnection(spec.serverName);
    if (!connection || connection.status !== "connected") {
      return {
        content: [{ type: "text" as const, text: `MCP server "${spec.serverName}" not connected` }],
        details: { error: "not_connected", server: spec.serverName },
      };
    }

    let uiSession: UiSessionRuntime | null = null;

    try {
      state.manager.touch(spec.serverName);
      state.manager.incrementInFlight(spec.serverName);

      if (spec.resourceUri) {
        const result = await connection.client.readResource({ uri: spec.resourceUri });
        const content = (result.contents ?? []).map(c => ({
          type: "text" as const,
          text: "text" in c ? c.text : ("blob" in c ? `[Binary data: ${(c as { mimeType?: string }).mimeType ?? "unknown"}]` : JSON.stringify(c)),
        }));
        return {
          content: content.length > 0 ? content : [{ type: "text" as const, text: "(empty resource)" }],
          details: { server: spec.serverName, resourceUri: spec.resourceUri },
        };
      }

      const hasUi = !!spec.uiResourceUri;
      uiSession = hasUi
        ? await maybeStartUiSession(state, {
            serverName: spec.serverName,
            toolName: spec.originalName,
            toolArgs: params ?? {},
            uiResourceUri: spec.uiResourceUri!,
            streamMode: spec.uiStreamMode,
          })
        : null;

      const resultPromise = connection.client.callTool({
        name: spec.originalName,
        arguments: params ?? {},
        _meta: uiSession?.requestMeta,
      });

      const result = await resultPromise;
      uiSession?.sendToolResult(result as unknown as import("@modelcontextprotocol/sdk/types.js").CallToolResult);

      const mcpContent = (result.content ?? []) as McpContent[];
      const content = transformMcpContent(mcpContent);

      if (result.isError) {
        let errorText = content.filter(c => c.type === "text").map(c => (c as { text: string }).text).join("\n") || "Tool execution failed";
        if (spec.inputSchema) {
          errorText += `\n\nExpected parameters:\n${formatSchema(spec.inputSchema)}`;
        }
        return {
          content: [{ type: "text" as const, text: `Error: ${errorText}` }],
          details: { error: "tool_error", server: spec.serverName },
        };
      }

      const resultText = content.filter(c => c.type === "text").map(c => (c as { text: string }).text).join("\n") || "(empty result)";
      if (hasUi) {
        const uiMessage = uiSession?.reused
          ? "Updated the open UI."
          : "📺 Interactive UI is now open in your browser. I'll respond to your prompts and intents as you interact with it.";
        return {
          content: [{ type: "text" as const, text: `${resultText}\n\n${uiMessage}` }],
          details: { server: spec.serverName, tool: spec.originalName, uiOpen: true },
        };
      }

      return {
        content: content.length > 0 ? content : [{ type: "text" as const, text: "(empty result)" }],
        details: { server: spec.serverName, tool: spec.originalName },
      };
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      uiSession?.sendToolCancelled(message);
      let errorText = `Failed to call tool: ${message}`;
      if (spec.inputSchema) {
        errorText += `\n\nExpected parameters:\n${formatSchema(spec.inputSchema)}`;
      }
      return {
        content: [{ type: "text" as const, text: errorText }],
        details: { error: "call_failed", server: spec.serverName },
      };
    } finally {
      if (uiSession?.reused) {
        uiSession.close();
      }
      state.manager.decrementInFlight(spec.serverName);
      state.manager.touch(spec.serverName);
    }
  };
}