/**
 * Bridge the mabl MCP server into pi.dev as custom tools.
 *
 * This extension does NOT run the mabl OAuth check at startup. Instead it
 * registers a `/mabl-auth` command that the user can run explicitly. The
 * command checks authentication, runs `mabl auth login --auto` if needed, and
 * then starts `mabl mcp start` and registers the exposed tools.
 */
import { createRequire } from "node:module";
import path from "node:path";
import { spawn } from "node:child_process";
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import type {
  ExtensionAPI,
  ExtensionCommandContext,
} from "@earendil-works/pi-coding-agent";
import { Type, type TSchema } from "typebox";

const require = createRequire(import.meta.url);

function jsonSchemaToTypeBox(schema: unknown): TSchema {
  if (!schema || typeof schema !== "object") {
    return Type.Any();
  }
  const s = schema as Record<string, any>;

  switch (s.type) {
    case "object": {
      const properties: Record<string, TSchema> = {};
      const required = new Set<string>(s.required ?? []);
      for (const [key, value] of Object.entries(s.properties ?? {})) {
        const child = jsonSchemaToTypeBox(value);
        properties[key] = required.has(key) ? child : Type.Optional(child);
      }
      return Type.Object(properties);
    }
    case "array":
      return Type.Array(jsonSchemaToTypeBox(s.items));
    case "string":
      return Type.String({ description: s.description });
    case "number":
    case "integer":
      return Type.Number({ description: s.description });
    case "boolean":
      return Type.Boolean({ description: s.description });
    case "null":
      return Type.Null();
    default:
      break;
  }

  if (Array.isArray(s.enum)) {
    return Type.Union(s.enum.map((value: unknown) => Type.Literal(value as any)));
  }

  return Type.Any({ description: s.description });
}

function runCommand(
  command: string,
  args: string[],
  options?: { stdio?: "inherit" | "pipe"; timeoutMs?: number }
): Promise<{ exitCode: number; stdout: string; stderr: string }> {
  return new Promise((resolve) => {
    const stdio = options?.stdio === "inherit" ? "inherit" : "pipe";
    const child = spawn(command, args, {
      stdio: ["inherit", stdio, stdio],
      shell: false,
      windowsHide: true,
    });

    const stdoutChunks: Buffer[] = [];
    const stderrChunks: Buffer[] = [];
    let killedByTimeout = false;

    if (child.stdout && stdio === "pipe") {
      child.stdout.on("data", (chunk) => stdoutChunks.push(chunk));
    }
    if (child.stderr && stdio === "pipe") {
      child.stderr.on("data", (chunk) => stderrChunks.push(chunk));
    }

    const timeout = options?.timeoutMs
      ? setTimeout(() => {
          killedByTimeout = true;
          child.kill("SIGTERM");
        }, options.timeoutMs)
      : undefined;

    let settled = false;

    child.on("error", (err) => {
      if (settled) return;
      settled = true;
      if (timeout) clearTimeout(timeout);
      resolve({
        exitCode: -1,
        stdout: Buffer.concat(stdoutChunks).toString("utf-8"),
        stderr: Buffer.concat(stderrChunks).toString("utf-8") || String(err),
      });
    });

    child.on("exit", (code) => {
      if (settled) return;
      settled = true;
      if (timeout) clearTimeout(timeout);
      resolve({
        exitCode: killedByTimeout ? -1 : code ?? 0,
        stdout: Buffer.concat(stdoutChunks).toString("utf-8"),
        stderr: Buffer.concat(stderrChunks).toString("utf-8"),
      });
    });
  });
}

async function resolveMablBin(): Promise<string | undefined> {
  try {
    const pkgPath = require.resolve("@mablhq/mabl-cli/package.json");
    return path.join(path.dirname(pkgPath), "cli.js");
  } catch {
    return undefined;
  }
}

async function checkAuthentication(
  nodePath: string,
  mablBin: string
): Promise<boolean> {
  const infoResult = await runCommand(nodePath, [mablBin, "auth", "info"], {
    timeoutMs: 30000,
  });
  return (
    infoResult.exitCode === 0 &&
    !infoResult.stdout.toLowerCase().includes("not logged in") &&
    !infoResult.stderr.toLowerCase().includes("not logged in")
  );
}

export default async function (pi: ExtensionAPI) {
  let mcpState:
    | {
        transport: StdioClientTransport;
        client: Client;
        toolsRegistered: number;
      }
    | undefined;
  let cleanupRegistered = false;

  const cleanup = async () => {
    if (!mcpState) return;
    try {
      await mcpState.transport.close();
      await mcpState.client.close();
    } catch {
      // ignore cleanup errors
    }
    mcpState = undefined;
  };

  pi.registerCommand("mabl-auth", {
    description:
      "Check mabl OAuth status, log in if needed, and load mabl MCP tools.",
    handler: async (_args: string, ctx: ExtensionCommandContext) => {
      const mablBin = await resolveMablBin();
      if (!mablBin) {
        ctx.ui.notify(
          "[mabl-mcp] @mablhq/mabl-cli is not installed. Run `npm install` in the Dino package directory.",
          "error"
        );
        return;
      }

      const nodePath = process.execPath;

      if (mcpState) {
        const stillAuthenticated = await checkAuthentication(nodePath, mablBin);
        if (stillAuthenticated) {
          ctx.ui.notify(
            `[mabl-mcp] Already authenticated. ${mcpState.toolsRegistered} mabl MCP tool(s) loaded.`,
            "info"
          );
        } else {
          ctx.ui.notify(
            "[mabl-mcp] MCP server is running but authentication appears expired. Run `/mabl-auth` again after logging in.",
            "warning"
          );
        }
        return;
      }

      ctx.ui.notify("[mabl-mcp] Checking mabl authentication...", "info");
      let authenticated = await checkAuthentication(nodePath, mablBin);

      if (!authenticated) {
        ctx.ui.notify(
          "[mabl-mcp] Not authenticated. Starting OAuth login (auto-capture mode)...",
          "info"
        );
        const loginResult = await runCommand(
          nodePath,
          [mablBin, "auth", "login", "--auto"],
          { stdio: "inherit" }
        );
        if (loginResult.exitCode !== 0) {
          ctx.ui.notify(
            `[mabl-mcp] mabl OAuth login failed or was cancelled (exit ${loginResult.exitCode}). Run manually with: npx @mablhq/mabl-cli auth login --auto`,
            "error"
          );
          return;
        }
        authenticated = await checkAuthentication(nodePath, mablBin);
        if (!authenticated) {
          ctx.ui.notify(
            "[mabl-mcp] mabl OAuth login appeared to succeed but is still not authenticated.",
            "error"
          );
          return;
        }
        ctx.ui.notify("[mabl-mcp] mabl OAuth login succeeded.", "info");
      } else {
        ctx.ui.notify("[mabl-mcp] Already authenticated with mabl.", "info");
      }

      // Start the mabl MCP server and register tools.
      ctx.ui.notify("[mabl-mcp] Starting mabl MCP server...", "info");
      const transport = new StdioClientTransport({
        command: nodePath,
        args: [mablBin, "mcp", "start"],
        stderr: "inherit",
      });

      const client = new Client({
        name: "dino-mabl-mcp",
        version: "0.1.0",
      });

      try {
        await client.connect(transport);
      } catch (err) {
        ctx.ui.notify(
          `[mabl-mcp] Could not connect to mabl MCP server: ${err}`,
          "error"
        );
        return;
      }

      let tools;
      try {
        const result = await client.listTools();
        tools = result.tools;
      } catch (err) {
        ctx.ui.notify(
          `[mabl-mcp] Could not list mabl MCP tools: ${err}`,
          "error"
        );
        return;
      }

      for (const tool of tools) {
        pi.registerTool({
          name: tool.name,
          label: tool.name,
          description: tool.description ?? ``,
          promptSnippet: `mabl MCP tool: ${tool.description ?? tool.name}`,
          parameters: jsonSchemaToTypeBox(tool.inputSchema),
          async execute(_toolCallId, params) {
            const result = await client.callTool({
              name: tool.name,
              arguments: (params ?? {}) as Record<string, unknown>,
            });

            const content = (result.content as any[]).map((item: any) => {
              if (item.type === "text") {
                return { type: "text" as const, text: item.text };
              }
              if (item.type === "image") {
                return {
                  type: "image" as const,
                  data: item.data,
                  mimeType: item.mimeType,
                };
              }
              return { type: "text" as const, text: JSON.stringify(item) };
            });

            return { content, details: result };
          },
        });
      }

      mcpState = { transport, client, toolsRegistered: tools.length };
      ctx.ui.notify(
        `[mabl-mcp] Registered ${tools.length} mabl MCP tool(s).`,
        "info"
      );

      if (!cleanupRegistered) {
        cleanupRegistered = true;
        process.on("SIGINT", cleanup);
        process.on("SIGTERM", cleanup);
        pi.on("session_shutdown", cleanup);
      }
    },
  });
}