import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { spawn, type ChildProcess } from "node:child_process";
import readline from "node:readline";

/**
 * pi-mcp-bridge — a minimal generic MCP client extension for pi.
 *
 * Spawns an MCP server as a subprocess, discovers its tools via the standard
 * MCP handshake, and registers them with pi. Maintains a persistent connection
 * to the server for efficient tool calls (eliminates subprocess startup per call).
 *
 * All configuration is via environment variables:
 *   MCP_SERVER_CMD  — Command to run the MCP server (required)
 *   MCP_SERVER_DIR  — Working directory for the server (optional)
 *   MCP_SERVER_ARGS — JSON array of extra CLI arguments (optional)
 */

// ── Configuration ───────────────────────────────────────────────────────────

interface Config {
  serverCmd: string;
  serverDir: string;
  serverArgs: string[];
}

function loadConfig(): Config | null {
  const serverCmd = process.env.MCP_SERVER_CMD?.trim();
  if (!serverCmd) return null;

  let serverArgs: string[] = [];
  try {
    const parsed = JSON.parse(process.env.MCP_SERVER_ARGS || "[]");
    if (Array.isArray(parsed)) serverArgs = parsed;
  } catch {
    // ignore malformed JSON
  }

  return {
    serverCmd,
    serverDir: process.env.MCP_SERVER_DIR || "",
    serverArgs,
  };
}

// ── JSON Schema → TypeBox converter ────────────────────────────────────────

/**
 * Convert an MCP inputSchema (JSON Schema subset) to a TypeBox schema.
 * This lets the LLM see proper parameter types for each MCP tool.
 */
function inputSchemaToTypeBox(inputSchema: Record<string, unknown>): any {
  const convert = (schema: Record<string, unknown>): any => {
    switch (schema.type) {
      case "string":
        return Type.String({ description: schema.description as string | undefined });
      case "number":
        return Type.Number({ description: schema.description as string | undefined });
      case "integer":
        return Type.Integer({ description: schema.description as string | undefined });
      case "boolean":
        return Type.Boolean({ description: schema.description as string | undefined });
      case "array": {
        const items = schema.items as Record<string, unknown> | undefined;
        return Type.Array(items ? convert(items) : Type.Unknown());
      }
      case "object": {
        const properties =
          (schema.properties as Record<string, Record<string, unknown>>) || {};
        const required = new Set<string>((schema.required as string[]) || []);
        const converted: Record<string, any> = {};

        for (const [key, prop] of Object.entries(properties)) {
          const ts = convert(prop);
          converted[key] = required.has(key) ? ts : Type.Optional(ts);
        }

        return Type.Object(converted);
      }
      default:
        return Type.Unknown();
    }
  };

  return convert(inputSchema);
}

// ── MCP Client ──────────────────────────────────────────────────────────────

interface PendingRequest {
  resolve: (value: unknown) => void;
  reject: (reason: Error) => void;
  timer: NodeJS.Timeout;
}

class McpClient {
  private proc: ChildProcess | null = null;
  private rl: readline.Interface | null = null;
  private pending = new Map<number, PendingRequest>();
  private nextId = 1;
  private _closed = false;
  private jsonBuffer = "";

  constructor(
    private cmd: string,
    private args: string[],
    private cwd: string | undefined,
  ) {}

  get closed(): boolean {
    return this._closed;
  }

  /** Start the MCP server subprocess and set up line-based JSON-RPC handling. */
  start(): void {
    if (this.proc) return;

    try {
      this.proc = spawn(this.cmd, this.args, {
        cwd: this.cwd,
        stdio: ["pipe", "pipe", "pipe"],
        env: process.env as Record<string, string>,
      });
    } catch (err) {
      const msg = err instanceof Error ? err.message : String(err);
      throw new Error(`Failed to spawn MCP server "${this.cmd}": ${msg}`);
    }

    const stderrChunks: Buffer[] = [];
    let stderrTimer: ReturnType<typeof setTimeout> | null = null;

    // Forward server stderr for debugging; buffer to avoid interleaving
    this.proc.stderr!.on("data", (data: Buffer) => {
      stderrChunks.push(data);
      // Flush stderr periodically rather than per-chunk
      if (!stderrTimer) {
        stderrTimer = setTimeout(() => {
          stderrTimer = null;
          const text = Buffer.concat(stderrChunks)
            .toString("utf-8")
            .replace(/\n$/, "");
          stderrChunks.length = 0;
          if (text) process.stderr.write(`[mcp-bridge:stderr] ${text}\n`);
        }, 50).unref();
      }
    });

    // Handle unexpected exit
    this.proc.on("exit", (code, signal) => {
      if (!this._closed) {
        const reason = signal
          ? `killed with signal ${signal}`
          : `exited with code ${code}`;
        this.rejectAll(new Error(`MCP server ${reason}`));
      }
      this.proc = null;
      this.rl = null;
    });

    this.proc.on("error", (err) => {
      this.rejectAll(new Error(`MCP server process error: ${err.message}`));
    });

    // Set up line-based reading of stdout
    this.rl = readline.createInterface({
      input: this.proc.stdout!,
      terminal: false,
    });
    this.rl.on("line", (line: string) => this.handleLine(line));
  }

  /**
   * Send a JSON-RPC request and wait for a response.
   * Default timeout is 30 seconds.
   */
  async request(
    method: string,
    params?: unknown,
    timeoutMs = 30_000,
  ): Promise<unknown> {
    if (this._closed) {
      throw new Error("MCP client is closed");
    }
    if (!this.proc || !this.proc.stdin?.writable) {
      throw new Error("MCP server is not running");
    }

    const id = this.nextId++;

    return new Promise<unknown>((resolve, reject) => {
      const timer = setTimeout(() => {
        this.pending.delete(id);
        reject(
          new Error(
            `MCP request "${method}" timed out after ${timeoutMs}ms`,
          ),
        );
      }, timeoutMs);

      this.pending.set(id, { resolve, reject, timer });

      const message = JSON.stringify({
        jsonrpc: "2.0",
        id,
        method,
        params: params !== undefined ? params : undefined,
      });

      const written = this.proc.stdin.write(message + "\n");

      // If the internal buffer is full, wait for drain before continuing
      if (!written) {
        this.proc.stdin.once("drain", () => {});
      }
    });
  }

  /** Send a JSON-RPC notification (no response expected). */
  notify(method: string, params?: unknown): void {
    if (this._closed || !this.proc?.stdin?.writable) return;

    const message = JSON.stringify({
      jsonrpc: "2.0",
      method,
      params: params !== undefined ? params : undefined,
    });

    this.proc.stdin.write(message + "\n");
  }

  /** Gracefully stop the MCP server and clean up resources. */
  async stop(): Promise<void> {
    if (this._closed) return;
    this._closed = true;

    // Try graceful shutdown via MCP protocol
    if (this.proc && !this.proc.killed) {
      try {
        await this.request("shutdown", undefined, 5_000);
      } catch {
        // Server may not support shutdown — that's fine
      }
    }

    this.cleanup();
    this.rejectAll(new Error("MCP client stopped"));
  }

  // ── Private ──

  /**
   * Try to parse a line as JSON. If it fails, accumulate in a buffer
   * (handles multi-line JSON or interleaved debug output).
   */
  private handleLine(line: string): void {
    const trimmed = line.trim();
    if (!trimmed) return;

    // Try parsing as standalone JSON
    let msg: any;
    let raw = trimmed;

    try {
      msg = JSON.parse(raw);
      this.jsonBuffer = ""; // clear buffer on success
    } catch {
      // Not valid JSON on its own — try buffering
      this.jsonBuffer += raw;
      try {
        msg = JSON.parse(this.jsonBuffer);
        this.jsonBuffer = ""; // parsed successfully from buffer
      } catch {
        // Still not valid JSON — if buffer is large, it's probably
        // debug output, so clear it to avoid memory growth
        if (this.jsonBuffer.length > 1024 * 1024 /* 1MiB */) {
          this.jsonBuffer = "";
        }
        return; // skip line
      }
    }

    this.dispatchMessage(msg);
  }

  /** Route a parsed JSON-RPC message to its pending promise. */
  private dispatchMessage(msg: any): void {
    // Notification (no id) — not expected after init, but safe to ignore
    if (msg.id === undefined) return;

    // Error response
    if (msg.error) {
      const pending = this.pending.get(msg.id);
      if (pending) {
        clearTimeout(pending.timer);
        this.pending.delete(msg.id);
        pending.reject(
          new Error(msg.error.message || "Unknown MCP error"),
        );
      }
      return;
    }

    // Successful response
    const pending = this.pending.get(msg.id);
    if (pending) {
      clearTimeout(pending.timer);
      this.pending.delete(msg.id);
      pending.resolve(msg.result);
    }
  }

  private cleanup(): void {
    this.rl?.close();
    this.rl = null;

    if (this.proc) {
      this.proc.stdin?.end();
      this.proc.kill("SIGTERM");

      // Force kill after 2s if still alive
      setTimeout(() => {
        if (this.proc && !this.proc.killed) {
          this.proc.kill("SIGKILL");
        }
      }, 2_000).unref();

      this.proc = null;
    }
  }

  private rejectAll(reason: Error): void {
    for (const [, pending] of this.pending) {
      clearTimeout(pending.timer);
      pending.reject(reason);
    }
    this.pending.clear();
  }
}

// ── Main Extension ──────────────────────────────────────────────────────────

export default async function (pi: ExtensionAPI) {
  const config = loadConfig();
  if (!config) {
    console.error(
      "[mcp-bridge] MCP_SERVER_CMD is not set. " +
        "Set this environment variable to the command that starts your MCP server.",
    );
    return;
  }

  const parts = config.serverCmd.split(/\s+/);
  const client = new McpClient(
    parts[0],
    [...parts.slice(1), ...config.serverArgs],
    config.serverDir || undefined,
  );

  try {
    // Start the MCP server subprocess
    client.start();

    // Phase 1: Initialize
    const initResult = (await client.request("initialize", {
      protocolVersion: "2024-11-05",
      capabilities: {},
      clientInfo: { name: "pi-mcp-bridge", version: "0.1.0" },
    })) as Record<string, unknown>;

    const serverVersion =
      ((initResult?.serverInfo as Record<string, unknown>)?.name) ?? "unknown";
    console.error(`[mcp-bridge] connected to ${String(serverVersion)}`);

    // Phase 2: Send initialized notification (fire-and-forget, no response)
    client.notify("notifications/initialized");

    // Phase 3: Discover and register tools
    const listResult = (await client.request("tools/list")) as Record<
      string,
      unknown
    >;
    const tools = (listResult?.tools as Array<{
      name: string;
      description?: string;
      inputSchema?: Record<string, unknown>;
    }>) ?? [];

    console.error(
      `[mcp-bridge] registered ${tools.length} tool(s) from ${config.serverCmd}`,
    );

    for (const tool of tools) {
      const schema = tool.inputSchema
        ? inputSchemaToTypeBox(tool.inputSchema)
        : Type.Object({});

      // Capture tool name in closure for the execute handler
      const toolName = tool.name;

      pi.registerTool({
        name: toolName,
        label: tool.name.replace(/_/g, " "),
        description: tool.description ?? `MCP tool: ${toolName}`,
        parameters: schema,
        async execute(_toolCallId, params) {
          const result = (await client.request("tools/call", {
            name: toolName,
            arguments: params,
          })) as Record<string, unknown>;

          const content = (result?.content ?? []) as Array<{
            type: string;
            text?: string;
            data?: unknown;
          }>;
          const text = content
            .filter((c) => c.type === "text")
            .map((c) => c.text ?? "")
            .join("\n");

          if (text) {
            return { content: [{ type: "text" as const, text }] };
          }

          return {
            content: [{ type: "text" as const, text: JSON.stringify(result) }],
          };
        },
      });
    }

    // Cleanup when pi session ends
    pi.on("session_shutdown", async () => {
      await client.stop();
    });
  } catch (err) {
    const msg = err instanceof Error ? err.message : String(err);
    console.error(`[mcp-bridge] initialization failed: ${msg}`);
    await client.stop();
  }
}

// ── Exports for testing ────────────────────────────────────────────────────

/** @internal Exported for unit testing only. */
export { loadConfig, inputSchemaToTypeBox, McpClient };
export type { Config };