import { compile } from "json-schema-to-typescript";
import { format } from "prettier";

export interface ToolDefinition {
  name: string;
  inputSchema: object;
  outputSchema?: object;
}

export interface GenerateOptions {
  mcpId: string;
  tools: ToolDefinition[];
}

export interface ToolSchemaFile {
  filename: string;
  content: string;
}

/**
 * Produce one JSON Schema file per tool: `<TOOL>.json` holding
 * `{ name, description?, inputSchema, outputSchema? }`. Pure — the caller
 * writes them to disk. Filenames are sanitized to filesystem-safe chars while
 * the original tool name is preserved inside the file.
 */
export function toolSchemaFiles(
  tools: (ToolDefinition & { description?: string })[],
): ToolSchemaFile[] {
  return tools.map((tool) => {
    const body: Record<string, unknown> = { name: tool.name };
    if (tool.description) body.description = tool.description;
    body.inputSchema = tool.inputSchema;
    if (tool.outputSchema) body.outputSchema = tool.outputSchema;
    return {
      filename: `${tool.name.replace(/[^A-Za-z0-9_.-]/g, "_")}.json`,
      content: `${JSON.stringify(body, null, 2)}\n`,
    };
  });
}

const BANNER = `// This file was auto-generated by @decocms/typegen. Do not edit manually.
// Regenerate with: bunx @decocms/typegen --mcp <id> --key <key> --output <file>
`;

const PRETTIER_CONFIG = {
  parser: "typescript" as const,
  printWidth: 80,
  singleQuote: false,
  trailingComma: "all" as const,
  semi: true,
};

/**
 * Parse compile() output into individual declarations.
 * Each declaration may be preceded by a JSDoc comment block.
 */
function parseDeclarations(raw: string): {
  name: string;
  kind: "type" | "interface";
  body: string;
  full: string;
}[] {
  const declarations: {
    name: string;
    kind: "type" | "interface";
    body: string;
    full: string;
  }[] = [];

  const lines = raw.split("\n");
  let i = 0;

  while (i < lines.length) {
    // Skip blank lines
    if (lines[i].trim() === "") {
      i++;
      continue;
    }

    // Consume optional JSDoc comment
    if (lines[i].trim().startsWith("/**")) {
      while (i < lines.length && !lines[i].includes("*/")) i++;
      i++; // skip the closing */ line
      continue;
    }

    // Match `export type Name = ...` (may span multiple lines for union/intersection types)
    const typeMatch = lines[i].match(/^export\s+type\s+(\w+)\s*=\s*(.*)/);
    if (typeMatch) {
      const name = typeMatch[1];
      const bodyLines: string[] = [typeMatch[2]];
      let braceDepth = 0;
      for (const ch of typeMatch[2]) {
        if (ch === "{") braceDepth++;
        if (ch === "}") braceDepth--;
      }
      // Accumulate lines until all braces are closed and we hit a trailing semicolon
      while (i + 1 < lines.length) {
        const trimmed = bodyLines.join("\n").trimEnd();
        if (braceDepth === 0 && trimmed.endsWith(";")) break;
        i++;
        bodyLines.push(lines[i]);
        for (const ch of lines[i]) {
          if (ch === "{") braceDepth++;
          if (ch === "}") braceDepth--;
        }
      }
      const body = bodyLines.join("\n").replace(/;\s*$/, "").trim();
      declarations.push({
        name,
        kind: "type",
        body,
        full: bodyLines.join("\n"),
      });
      i++;
      continue;
    }

    // Match `export interface Name {`
    const ifaceMatch = lines[i].match(/^export\s+interface\s+(\w+)\s*\{/);
    if (ifaceMatch) {
      const name = ifaceMatch[1];
      let braceDepth = 0;
      const blockLines: string[] = [];
      for (; i < lines.length; i++) {
        blockLines.push(lines[i]);
        for (const ch of lines[i]) {
          if (ch === "{") braceDepth++;
          if (ch === "}") braceDepth--;
        }
        if (braceDepth === 0) break;
      }
      i++;
      const full = blockLines.join("\n");
      // Body = everything between the outer braces
      const body = full
        .replace(/^export\s+interface\s+\w+\s*/, "")
        .replace(/;\s*$/, "")
        .trim();
      declarations.push({ name, kind: "interface", body, full });
      continue;
    }

    i++;
  }

  return declarations;
}

async function schemaToTs(schema: object, typeName: string): Promise<string> {
  const raw = await compile(schema as never, typeName, {
    bannerComment: "",
    additionalProperties: false,
  });

  const decls = parseDeclarations(raw);

  // compile() may use the schema's `title` instead of our typeName for the primary
  // declaration. Find by name first, then fall back to the "root" declaration — the
  // one whose name isn't referenced in any other declaration's body.
  let primary = decls.find((d) => d.name === typeName);
  if (!primary && decls.length > 0) {
    const allBodies = decls.map((d) => d.body).join("\n");
    primary = decls.find((d) => {
      const others = allBodies.replace(d.body, "");
      return !new RegExp(`\\b${d.name}\\b`).test(others);
    });
    primary ??= decls[decls.length - 1];
  }

  if (!primary) {
    return raw.replace(/;\s*$/, "").trim();
  }

  let result = primary.body;

  const helpers = decls.filter((d) => d !== primary);
  if (helpers.length > 0) {
    const replacements = new Map<string, string>();
    for (const helper of helpers) {
      replacements.set(
        helper.name,
        helper.kind === "type" ? `(${helper.body})` : helper.body,
      );
    }

    // Multi-pass inlining: helpers may reference other helpers (e.g. Filter uses Operator).
    // Keep replacing until stable or a safety cap is reached.
    let prev = "";
    for (let pass = 0; pass < 10 && result !== prev; pass++) {
      prev = result;
      for (const [name, replacement] of replacements) {
        result = result.replace(new RegExp(`\\b${name}\\b`, "g"), replacement);
      }
    }
  }

  return result;
}

export async function generateClientCode(
  opts: GenerateOptions,
): Promise<string> {
  const { mcpId, tools } = opts;

  const toolEntries: string[] = [];

  for (const tool of tools) {
    const inputType = await schemaToTs(tool.inputSchema, `${tool.name}Input`);
    const outputType = tool.outputSchema
      ? await schemaToTs(tool.outputSchema, `${tool.name}Output`)
      : "unknown";

    toolEntries.push(
      `  ${tool.name}: {\n    input: ${inputType};\n    output: ${outputType};\n  };`,
    );
  }

  const toolsInterface =
    tools.length === 0
      ? "export interface Tools {}"
      : `export interface Tools {\n${toolEntries.join("\n")}\n}`;

  const code = `${BANNER}
import { createStudioClient } from "@decocms/typegen";

${toolsInterface}

export const client = createStudioClient<Tools>({
  mcpId: "${mcpId}",
  apiKey: process.env.STUDIO_API_KEY ?? process.env.MESH_API_KEY,
  baseUrl: process.env.STUDIO_BASE_URL ?? process.env.MESH_BASE_URL,
});
`;

  return format(code, PRETTIER_CONFIG);
}
