#!/usr/bin/env node
/**
 * MCP Schema Sanitizer Proxy
 *
 * Wraps any MCP stdio server and sanitizes tool inputSchema objects
 * to remove top-level oneOf/allOf/anyOf that the Claude API rejects.
 *
 * Usage (in .mcp.json):
 *   "command": "node",
 *   "args": ["fluent-ai-skills/tools/mcp-schema-sanitizer.mjs", "--", "npx", "-y", "chrome-devtools-mcp@latest", ...]
 *
 * Everything after "--" is the real MCP server command + args.
 */

import { spawn } from "node:child_process";

// ── Parse args ──────────────────────────────────────────────────────────────
const separatorIdx = process.argv.indexOf("--");
if (separatorIdx === -1 || separatorIdx === process.argv.length - 1) {
  process.stderr.write(
    "mcp-schema-sanitizer: usage: node mcp-schema-sanitizer.mjs -- <command> [args...]\n"
  );
  process.exit(1);
}
const childCmd = process.argv[separatorIdx + 1];
const childArgs = process.argv.slice(separatorIdx + 2);

// ── Schema sanitization ─────────────────────────────────────────────────────

/** Merge an array of JSON Schema objects into one (best-effort). */
function mergeSchemas(schemas) {
  const merged = { type: "object", properties: {}, required: [] };
  for (const s of schemas) {
    if (!s || typeof s !== "object") continue;
    if (s.properties) Object.assign(merged.properties, s.properties);
    if (Array.isArray(s.required)) {
      for (const r of s.required) {
        if (!merged.required.includes(r)) merged.required.push(r);
      }
    }
    // Carry over additionalProperties if any schema sets it
    if (s.additionalProperties === false) merged.additionalProperties = false;
  }
  if (merged.required.length === 0) delete merged.required;
  if (Object.keys(merged.properties).length === 0) delete merged.properties;
  return merged;
}

/**
 * Recursively sanitize a JSON Schema node.
 * - Top-level anyOf/oneOf/allOf → merge into single object schema
 * - Nested anyOf/oneOf/allOf inside properties → keep (API allows nested)
 * - nullable anyOf [schema, {type:"null"}] → use the non-null branch
 */
function sanitizeSchema(schema, depth = 0) {
  if (!schema || typeof schema !== "object") return schema;

  // Only fix top-level (depth 0) — nested is fine
  if (depth === 0) {
    for (const keyword of ["anyOf", "oneOf", "allOf"]) {
      if (Array.isArray(schema[keyword])) {
        const branches = schema[keyword];

        // nullable pattern: anyOf: [realSchema, {type:"null"}]
        const nonNull = branches.filter(
          (b) => !(b && b.type === "null" && Object.keys(b).length === 1)
        );
        if (nonNull.length === 1) {
          const cleaned = { ...schema };
          delete cleaned[keyword];
          return sanitizeSchema({ ...nonNull[0], ...cleaned }, 0);
        }

        // General case: merge all branches
        const cleaned = { ...schema };
        delete cleaned[keyword];
        const merged = mergeSchemas(branches);
        return sanitizeSchema({ ...merged, ...cleaned }, 0);
      }
    }
  }

  // Recurse into properties
  if (schema.properties) {
    const newProps = {};
    for (const [key, val] of Object.entries(schema.properties)) {
      newProps[key] = sanitizeSchema(val, depth + 1);
    }
    schema = { ...schema, properties: newProps };
  }

  // Recurse into items (array schemas)
  if (schema.items) {
    schema = { ...schema, items: sanitizeSchema(schema.items, depth + 1) };
  }

  return schema;
}

/** Patch a tools/list response to sanitize all inputSchema objects. */
function patchToolsList(response) {
  if (!response?.result?.tools) return response;
  let patched = 0;
  for (const tool of response.result.tools) {
    if (!tool.inputSchema) continue;
    const original = JSON.stringify(tool.inputSchema);
    if (/oneOf|allOf|anyOf/.test(original)) {
      tool.inputSchema = sanitizeSchema(tool.inputSchema, 0);
      patched++;
    }
  }
  if (patched > 0) {
    process.stderr.write(
      `[mcp-schema-sanitizer] patched ${patched} tool schema(s)\n`
    );
  }
  return response;
}

// ── Spawn child MCP server ──────────────────────────────────────────────────
const child = spawn(childCmd, childArgs, {
  stdio: ["pipe", "pipe", "inherit"], // inherit stderr so child logs pass through
  env: process.env,
  shell: process.platform === "win32", // needed on Windows for npx etc.
});

child.on("error", (err) => {
  process.stderr.write(`[mcp-schema-sanitizer] spawn error: ${err.message}\n`);
  process.exit(1);
});

child.on("exit", (code) => {
  process.exit(code ?? 1);
});

// ── Pipe stdin → child (passthrough) ────────────────────────────────────────
process.stdin.pipe(child.stdin);

// ── Pipe child stdout → sanitize → our stdout ──────────────────────────────
let buffer = "";

child.stdout.on("data", (chunk) => {
  buffer += chunk.toString();

  // JSON-RPC messages are newline-delimited
  let newlineIdx;
  while ((newlineIdx = buffer.indexOf("\n")) !== -1) {
    const line = buffer.slice(0, newlineIdx).trim();
    buffer = buffer.slice(newlineIdx + 1);

    if (!line) {
      process.stdout.write("\n");
      continue;
    }

    try {
      const msg = JSON.parse(line);
      // Check if this is a tools/list response
      if (msg.result?.tools && Array.isArray(msg.result.tools)) {
        const patched = patchToolsList(msg);
        process.stdout.write(JSON.stringify(patched) + "\n");
      } else {
        process.stdout.write(line + "\n");
      }
    } catch {
      // Not valid JSON — pass through as-is
      process.stdout.write(line + "\n");
    }
  }
});

// Handle graceful shutdown
process.on("SIGTERM", () => child.kill("SIGTERM"));
process.on("SIGINT", () => child.kill("SIGINT"));
