import { spawn } from "child_process";
import { existsSync, readdirSync } from "fs";
import { resolve } from "path";

/**
 * Minimal MCP-over-stdio client that talks to the upstream
 * `qemu-mcp-server` Rust binary.
 *
 * The Rust server communicates via JSON-RPC 2.0 over stdio. This client
 * performs the MCP initialization handshake once and then sends `tools/call`
 * requests for each tool invocation.
 */

const MAX_OUTPUT_LENGTH = 2000;

export interface QemuMcpClientOptions {
  /** Override the QEMU QMP unix socket path. */
  qmpSocketPath?: string;
  /** Override the guest serial unix socket path. */
  serialSocketPath?: string;
  /** Abort signal to cancel the in-flight tool call. */
  abortSignal?: AbortSignal;
}

/**
 * Invokes a single tool exposed by the upstream qemu-mcp-server.
 *
 * @returns The text content of the tool response, or an error message string.
 */
export async function invokeTool(
  name: string,
  args: Record<string, unknown>,
  options: QemuMcpClientOptions = {}
): Promise<string> {
  const { qmpSocketPath, serialSocketPath, abortSignal } = options;

  const binary = resolveBinary();

  // NOTE: resolveBinary() already throws a helpful error when the binary is
  // missing, so this check is a defensive fallback.
  if (!existsSync(binary)) {
    throw new Error(
      "qemu-mcp-server binary not found at:\n  " +
        binary +
        "\n\n" +
        "The pre-built binary (shipped in the published package) may be " +
        "corrupted or missing. Please try one of the following:\n" +
        "  1. Reinstall the plugin: `lms uninstall qemu-mcp` then `lms install qemu-mcp`\n" +
        "  2. Build from source (requires Rust/cargo):\n" +
        "     cd <plugin-directory>\n" +
        "     node scripts/setup.cjs"
    );
  }

  const env: NodeJS.ProcessEnv = { ...process.env };
  if (qmpSocketPath !== undefined) env.QMP_SOCKET_PATH = qmpSocketPath;
  if (serialSocketPath !== undefined) env.SERIAL_SOCKET_PATH = serialSocketPath;

  return new Promise<string>((outerResolve, outerReject) => {
    const child = spawn(binary, [], { env, stdio: ["pipe", "pipe", "pipe"] });

    // Map of pending response ids -> resolver pairs.
    const pending = new Map<
      number,
      { resolve: (r: unknown) => void; reject: (e: Error) => void }
    >();
    let buffer = "";

    const onStdout = (chunk: Buffer) => {
      buffer += chunk.toString();
      let idx: number;
      while ((idx = buffer.indexOf("\n")) !== -1) {
        const line = buffer.slice(0, idx).trim();
        buffer = buffer.slice(idx + 1);
        if (!line) continue;

        let parsed:
          | { id: number; result?: { content?: unknown[] }; error?: { message: string } }
          | undefined;
        try {
          const maybeParsed = JSON.parse(line);
          if (maybeParsed === null || typeof maybeParsed !== "object") {
            // Non-object JSON (e.g. a bare string/number) is not a valid MCP
            // message; ignore it and keep waiting for more data.
            continue;
          }
          parsed = maybeParsed as typeof parsed;
        } catch {
          // Incomplete message arriving mid-chunk; wait for more data.
          continue;
        }

        if (parsed === undefined || parsed.id === undefined) {
          // Not a complete JSON-RPC message (missing id); ignore.
          continue;
        }

        const waiter = pending.get(parsed.id);
        if (waiter) {
          pending.delete(parsed.id);
          if (parsed.error) {
            waiter.reject(new Error(parsed.error.message));
          } else {
            waiter.resolve(parsed.result ?? {});
          }
        }
        // Messages without a matching waiter (e.g. notifications) are ignored.
      }
    };

    child.stdout.on("data", onStdout);
    child.stderr.on("data", (chunk: Buffer) => {
      // Upstream server logs go to stderr per MCP convention. Surface them
      // without breaking the JSON-RPC flow.
      process.stderr.write("[qemu-mcp] " + chunk.toString());
    });

    const cleanup = () => {
      child.stdout.off("data", onStdout);
      child.stderr.off("data", onStdout);
      child.stdin.end();
      child.kill();
    };

    const onAbort = () => {
      cleanup();
      pending.forEach((v) => v.reject(new Error("Operation aborted by user.")));
      outerReject(new Error("Operation aborted by user."));
    };

    if (abortSignal) {
      if (abortSignal.aborted) {
        onAbort();
        return;
      }
      abortSignal.addEventListener("abort", onAbort, { once: true });
    }

    child.on("error", (err) => {
      abortSignal?.removeEventListener("abort", onAbort);
      outerReject(
        new Error("Failed to start qemu-mcp-server: " + err.message)
      );
    });

    // 1. MCP initialize handshake (id = 1).
    const initPromise = new Promise<unknown>((res, rej) => {
      pending.set(1, { resolve: res, reject: rej });
      child.stdin.write(
        JSON.stringify({
          jsonrpc: "2.0",
          id: 1,
          method: "initialize",
          params: {
            protocolVersion: "2024-11-05",
            capabilities: { tools: {} },
            clientInfo: { name: "qemu-mcp-lmstudio", version: "0.1.0" },
          },
        }) + "\n"
      );
    });

    // 2. After initialize, send `notifications/initialized` then `tools/call`.
    (async () => {
      try {
        await initPromise;
        child.stdin.write(
          JSON.stringify({ jsonrpc: "2.0", method: "notifications/initialized" }) +
            "\n"
        );

        const callId = 2;
        const callPromise = new Promise<unknown>((res, rej) => {
          pending.set(callId, { resolve: res, reject: rej });
          child.stdin.write(
            JSON.stringify({
              jsonrpc: "2.0",
              id: callId,
              method: "tools/call",
              params: { name, arguments: args },
            }) + "\n"
          );
        });

        const result = (await callPromise) as { content?: unknown[] } | undefined;
        abortSignal?.removeEventListener("abort", onAbort);
        cleanup();
        outerResolve(extractText(result?.content ?? []));
      } catch (err) {
        abortSignal?.removeEventListener("abort", onAbort);
        cleanup();
        outerReject(err);
      }
    })();
  });
}

/** Extract plain text from MCP `content` array. */
function extractText(content: unknown[] | undefined): string {
  if (!Array.isArray(content)) return JSON.stringify(content);
  const parts: string[] = [];
  for (const part of content) {
    if (
      part &&
      typeof part === "object" &&
      "type" in part &&
      (part as { type: string }).type === "text" &&
      "text" in part
    ) {
      parts.push((part as { text: string }).text);
    } else if (part && typeof part === "object") {
      parts.push(JSON.stringify(part));
    }
  }
  const joined = parts.join("\n");
  return joined.length > MAX_OUTPUT_LENGTH
    ? joined.substring(0, MAX_OUTPUT_LENGTH) + "... [truncated]"
    : joined;
}

/** Resolve the path to the compiled Rust binary. */
function resolveBinary(): string {
  // Walk up from __dirname (CJS) to find the plugin root (where package.json
  // lives). This is symlink-safe and works in the LM Studio Node environment.
  let pluginRoot: string | undefined;
  let dir = __dirname;
  for (let i = 0; i < 5; i++) {
    if (existsSync(resolve(dir, "package.json"))) {
      pluginRoot = dir;
      break;
    }
    const parent = resolve(dir, "..");
    if (parent === dir) break; // reached filesystem root
    dir = parent;
  }
  if (!pluginRoot) {
    throw new Error(
      "Could not locate plugin root (package.json not found near " + __dirname + ")"
    );
  }

  // 1. Prefer the binary shipped in the published package (bin/). This is the
  // fast path that works without any build step or cargo installation.
  const shippedBinary = resolve(pluginRoot, "bin", "qemu-mcp-server");
  if (existsSync(shippedBinary)) {
    return shippedBinary;
  }

  // 2. Fall back to the locally built binary from the source tree.
  const localBinary = resolve(
    pluginRoot,
    "qemu-mcp-server",
    "target",
    "release",
    "qemu-mcp-server"
  );
  if (existsSync(localBinary)) {
    return localBinary;
  }

  // 3. Binary not found — throw a helpful error.
  throw new Error(
    "qemu-mcp-server binary not found. " +
      "The pre-built binary (shipped in the published package) is missing, " +
      "and no local build was found.\n\n" +
      "Please try one of the following:\n" +
      "  1. Reinstall the plugin: `lms uninstall qemu-mcp` then `lms install qemu-mcp`\n" +
      "  2. Build from source (requires Rust/cargo):\n" +
      "     cd " + pluginRoot + "\n" +
      "     node scripts/setup.cjs\n" +
      "  3. Manually place a compiled qemu-mcp-server binary at: " +
      shippedBinary
  );
}

import { spawn } from "child_process";
import { existsSync, readdirSync } from "fs";
import { resolve } from "path";

/**
 * Minimal MCP-over-stdio client that talks to the upstream
 * `qemu-mcp-server` Rust binary.
 *
 * The Rust server communicates via JSON-RPC 2.0 over stdio. This client
 * performs the MCP initialization handshake once and then sends `tools/call`
 * requests for each tool invocation.
 */

const MAX_OUTPUT_LENGTH = 2000;

export interface QemuMcpClientOptions {
  /** Override the QEMU QMP unix socket path. */
  qmpSocketPath?: string;
  /** Override the guest serial unix socket path. */
  serialSocketPath?: string;
  /** Abort signal to cancel the in-flight tool call. */
  abortSignal?: AbortSignal;
}

/**
 * Invokes a single tool exposed by the upstream qemu-mcp-server.
 *
 * @returns The text content of the tool response, or an error message string.
 */
export async function invokeTool(
  name: string,
  args: Record<string, unknown>,
  options: QemuMcpClientOptions = {}
): Promise<string> {
  const { qmpSocketPath, serialSocketPath, abortSignal } = options;

  const binary = resolveBinary();

  // NOTE: resolveBinary() already throws a helpful error when the binary is
  // missing, so this check is a defensive fallback.
  if (!existsSync(binary)) {
    throw new Error(
      "qemu-mcp-server binary not found at:\n  " +
        binary +
        "\n\n" +
        "The pre-built binary (shipped in the published package) may be " +
        "corrupted or missing. Please try one of the following:\n" +
        "  1. Reinstall the plugin: `lms uninstall qemu-mcp` then `lms install qemu-mcp`\n" +
        "  2. Build from source (requires Rust/cargo):\n" +
        "     cd <plugin-directory>\n" +
        "     node scripts/setup.cjs"
    );
  }

  const env: NodeJS.ProcessEnv = { ...process.env };
  if (qmpSocketPath !== undefined) env.QMP_SOCKET_PATH = qmpSocketPath;
  if (serialSocketPath !== undefined) env.SERIAL_SOCKET_PATH = serialSocketPath;

  return new Promise<string>((outerResolve, outerReject) => {
    const child = spawn(binary, [], { env, stdio: ["pipe", "pipe", "pipe"] });

    // Map of pending response ids -> resolver pairs.
    const pending = new Map<
      number,
      { resolve: (r: unknown) => void; reject: (e: Error) => void }
    >();
    let buffer = "";

    const onStdout = (chunk: Buffer) => {
      buffer += chunk.toString();
      let idx: number;
      while ((idx = buffer.indexOf("\n")) !== -1) {
        const line = buffer.slice(0, idx).trim();
        buffer = buffer.slice(idx + 1);
        if (!line) continue;

        let parsed:
          | { id: number; result?: { content?: unknown[] }; error?: { message: string } }
          | undefined;
        try {
          const maybeParsed = JSON.parse(line);
          if (maybeParsed === null || typeof maybeParsed !== "object") {
            // Non-object JSON (e.g. a bare string/number) is not a valid MCP
            // message; ignore it and keep waiting for more data.
            continue;
          }
          parsed = maybeParsed as typeof parsed;
        } catch {
          // Incomplete message arriving mid-chunk; wait for more data.
          continue;
        }

        if (parsed === undefined || parsed.id === undefined) {
          // Not a complete JSON-RPC message (missing id); ignore.
          continue;
        }

        const waiter = pending.get(parsed.id);
        if (waiter) {
          pending.delete(parsed.id);
          if (parsed.error) {
            waiter.reject(new Error(parsed.error.message));
          } else {
            waiter.resolve(parsed.result ?? {});
          }
        }
        // Messages without a matching waiter (e.g. notifications) are ignored.
      }
    };

    child.stdout.on("data", onStdout);
    child.stderr.on("data", (chunk: Buffer) => {
      // Upstream server logs go to stderr per MCP convention. Surface them
      // without breaking the JSON-RPC flow.
      process.stderr.write("[qemu-mcp] " + chunk.toString());
    });

    const cleanup = () => {
      child.stdout.off("data", onStdout);
      child.stderr.off("data", onStdout);
      child.stdin.end();
      child.kill();
    };

    const onAbort = () => {
      cleanup();
      pending.forEach((v) => v.reject(new Error("Operation aborted by user.")));
      outerReject(new Error("Operation aborted by user."));
    };

    if (abortSignal) {
      if (abortSignal.aborted) {
        onAbort();
        return;
      }
      abortSignal.addEventListener("abort", onAbort, { once: true });
    }

    child.on("error", (err) => {
      abortSignal?.removeEventListener("abort", onAbort);
      outerReject(
        new Error("Failed to start qemu-mcp-server: " + err.message)
      );
    });

    // 1. MCP initialize handshake (id = 1).
    const initPromise = new Promise<unknown>((res, rej) => {
      pending.set(1, { resolve: res, reject: rej });
      child.stdin.write(
        JSON.stringify({
          jsonrpc: "2.0",
          id: 1,
          method: "initialize",
          params: {
            protocolVersion: "2024-11-05",
            capabilities: { tools: {} },
            clientInfo: { name: "qemu-mcp-lmstudio", version: "0.1.0" },
          },
        }) + "\n"
      );
    });

    // 2. After initialize, send `notifications/initialized` then `tools/call`.
    (async () => {
      try {
        await initPromise;
        child.stdin.write(
          JSON.stringify({ jsonrpc: "2.0", method: "notifications/initialized" }) +
            "\n"
        );

        const callId = 2;
        const callPromise = new Promise<unknown>((res, rej) => {
          pending.set(callId, { resolve: res, reject: rej });
          child.stdin.write(
            JSON.stringify({
              jsonrpc: "2.0",
              id: callId,
              method: "tools/call",
              params: { name, arguments: args },
            }) + "\n"
          );
        });

        const result = (await callPromise) as { content?: unknown[] } | undefined;
        abortSignal?.removeEventListener("abort", onAbort);
        cleanup();
        outerResolve(extractText(result?.content ?? []));
      } catch (err) {
        abortSignal?.removeEventListener("abort", onAbort);
        cleanup();
        outerReject(err);
      }
    })();
  });
}

/** Extract plain text from MCP `content` array. */
function extractText(content: unknown[] | undefined): string {
  if (!Array.isArray(content)) return JSON.stringify(content);
  const parts: string[] = [];
  for (const part of content) {
    if (
      part &&
      typeof part === "object" &&
      "type" in part &&
      (part as { type: string }).type === "text" &&
      "text" in part
    ) {
      parts.push((part as { text: string }).text);
    } else if (part && typeof part === "object") {
      parts.push(JSON.stringify(part));
    }
  }
  const joined = parts.join("\n");
  return joined.length > MAX_OUTPUT_LENGTH
    ? joined.substring(0, MAX_OUTPUT_LENGTH) + "... [truncated]"
    : joined;
}

/** Resolve the path to the compiled Rust binary. */
function resolveBinary(): string {
  // Walk up from __dirname (CJS) to find the plugin root (where package.json
  // lives). This is symlink-safe and works in the LM Studio Node environment.
  let pluginRoot: string | undefined;
  let dir = __dirname;
  for (let i = 0; i < 5; i++) {
    if (existsSync(resolve(dir, "package.json"))) {
      pluginRoot = dir;
      break;
    }
    const parent = resolve(dir, "..");
    if (parent === dir) break; // reached filesystem root
    dir = parent;
  }
  if (!pluginRoot) {
    throw new Error(
      "Could not locate plugin root (package.json not found near " + __dirname + ")"
    );
  }

  // 1. Prefer the binary shipped in the published package (bin/). This is the
  // fast path that works without any build step or cargo installation.
  const shippedBinary = resolve(pluginRoot, "bin", "qemu-mcp-server");
  if (existsSync(shippedBinary)) {
    return shippedBinary;
  }

  // 2. Fall back to the locally built binary from the source tree.
  const localBinary = resolve(
    pluginRoot,
    "qemu-mcp-server",
    "target",
    "release",
    "qemu-mcp-server"
  );
  if (existsSync(localBinary)) {
    return localBinary;
  }

  // 3. Binary not found — throw a helpful error.
  throw new Error(
    "qemu-mcp-server binary not found. " +
      "The pre-built binary (shipped in the published package) is missing, " +
      "and no local build was found.\n\n" +
      "Please try one of the following:\n" +
      "  1. Reinstall the plugin: `lms uninstall qemu-mcp` then `lms install qemu-mcp`\n" +
      "  2. Build from source (requires Rust/cargo):\n" +
      "     cd " + pluginRoot + "\n" +
      "     node scripts/setup.cjs\n" +
      "  3. Manually place a compiled qemu-mcp-server binary at: " +
      shippedBinary
  );
}