import { readFile, stat } from "node:fs/promises";
import path from "node:path";
import readline from "node:readline";

const SERVER_NAME = "Sites Design Picker";
const SERVER_VERSION = JSON.parse(
  await readFile(new URL("../.codex-plugin/plugin.json", import.meta.url), "utf8"),
).version;
const TOOL_NAME = "choose_site_design";
const MAX_IMAGE_BYTES = 8 * 1024 * 1024;
const JsonRpcError = {
  METHOD_NOT_FOUND: -32601,
  INVALID_PARAMS: -32602,
};

let clientSupportsOpenAIForm = false;
let nextRequestId = 1;
const pendingRequests = new Map();

function send(message) {
  process.stdout.write(`${JSON.stringify(message)}\n`);
}

function sendResult(id, result) {
  send({ jsonrpc: "2.0", id, result });
}

function sendError(id, code, message) {
  send({ jsonrpc: "2.0", id, error: { code, message } });
}

function request(method, params) {
  const id = `server-${nextRequestId++}`;
  const response = new Promise((resolve, reject) => {
    pendingRequests.set(id, { resolve, reject });
  });
  send({ jsonrpc: "2.0", id, method, params });
  return response;
}

function requireNonEmptyString(value, name) {
  if (typeof value !== "string" || value.trim().length === 0) {
    throw new Error(`${name} must be a non-empty string.`);
  }
  return value.trim();
}

function imageMimeType(bytes) {
  if (
    bytes
      .subarray(0, 8)
      .equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))
  ) {
    return "image/png";
  }
  if (bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff) {
    return "image/jpeg";
  }
  if (
    bytes.subarray(0, 4).toString("ascii") === "RIFF" &&
    bytes.subarray(8, 12).toString("ascii") === "WEBP"
  ) {
    return "image/webp";
  }
  throw new Error("imagePath must contain a PNG, JPEG, or WebP image.");
}

async function imagePickerOption(option, index) {
  if (typeof option !== "object" || option === null || Array.isArray(option)) {
    throw new Error(`options[${index}] must be an object.`);
  }

  const id = requireNonEmptyString(option.id, `options[${index}].id`);
  const title = requireNonEmptyString(option.title, `options[${index}].title`);
  const designBrief = requireNonEmptyString(
    option.designBrief,
    `options[${index}].designBrief`,
  );
  const imagePath = requireNonEmptyString(
    option.imagePath,
    `options[${index}].imagePath`,
  );
  if (!path.isAbsolute(imagePath)) {
    throw new Error(`options[${index}].imagePath must be absolute.`);
  }

  let imageStat;
  try {
    imageStat = await stat(imagePath);
  } catch {
    throw new Error(`options[${index}].imagePath must be a regular file.`);
  }
  if (!imageStat.isFile()) {
    throw new Error(`options[${index}].imagePath must be a regular file.`);
  }
  if (imageStat.size > MAX_IMAGE_BYTES) {
    throw new Error(`options[${index}].imagePath must be 8 MB or smaller.`);
  }

  const bytes = await readFile(imagePath);
  const mimeType = imageMimeType(bytes);
  return {
    id,
    title,
    designBrief,
    image: `data:${mimeType};base64,${bytes.toString("base64")}`,
  };
}

async function getImagePickerOptions(value) {
  if (!Array.isArray(value) || value.length !== 3) {
    throw new Error("options must contain exactly three design options.");
  }
  const options = await Promise.all(value.map(imagePickerOption));
  if (new Set(options.map(({ id }) => id)).size !== options.length) {
    throw new Error("options must use unique ids.");
  }
  return options;
}

async function handleToolCall(id, params) {
  if (params?.name !== TOOL_NAME) {
    sendError(
      id,
      JsonRpcError.INVALID_PARAMS,
      `Unknown tool: ${params?.name ?? ""}`,
    );
    return;
  }
  if (!clientSupportsOpenAIForm) {
    throw new Error("The MCP client does not support openai/form requests.");
  }

  const options = await getImagePickerOptions(params.arguments?.options);
  const question =
    params.arguments?.question === undefined
      ? "Choose a visual direction for your site"
      : requireNonEmptyString(params.arguments.question, "question");
  const elicitation = await request("openai/form", {
    message: question,
    requestedSchema: {
      type: "object",
      properties: {
        selection: {
          type: "openai/imagePicker",
          title: "Visual directions",
          items: options.map(({ id: optionId, title, image }) => ({
            id: optionId,
            title,
            image,
          })),
        },
      },
      required: ["selection"],
    },
  });

  if (elicitation?.action !== "accept") {
    sendResult(id, {
      content: [
        {
          type: "text",
          text: "No design option was selected. Do not begin the site implementation until the user chooses how to proceed.",
        },
      ],
      structuredContent: {
        status: "not_selected",
        action: elicitation?.action ?? "cancel",
      },
    });
    return;
  }

  const selected = options.find(
    ({ id: optionId }) => optionId === elicitation.content?.selection,
  );
  if (selected === undefined) {
    throw new Error(
      "The selected design option is not one of the supplied options.",
    );
  }

  sendResult(id, {
    content: [
      {
        type: "text",
        text: `Selected design option: ${selected.title}\n\n${selected.designBrief}`,
      },
    ],
    structuredContent: {
      status: "selected",
      id: selected.id,
      title: selected.title,
      designBrief: selected.designBrief,
    },
  });
}

async function handleRequest(message) {
  const { id, method, params } = message;

  if (method === "initialize") {
    clientSupportsOpenAIForm =
      typeof params?.capabilities?.extensions?.["openai/form"] === "object";
    sendResult(id, {
      protocolVersion: params?.protocolVersion ?? "2025-11-25",
      capabilities: { tools: {} },
      serverInfo: {
        name: SERVER_NAME,
        version: SERVER_VERSION,
      },
      instructions:
        "Use choose_site_design only after the Sites skill generates exactly three comparable one-shot design-option previews. It may be called sequentially for up to four distinct design decisions; wait for each selection before preparing the next picker.",
    });
    return;
  }

  if (method === "ping") {
    sendResult(id, {});
    return;
  }

  if (method === "tools/list") {
    sendResult(id, {
      tools: [
        {
          name: TOOL_NAME,
          title: "Choose a Sites Design Option",
          description:
            "Show exactly three generated site-design options and ask the user a focused choice question. Options may be full-page concepts or HTML-rendered palettes, layouts, typography pairs, and other visual systems. Use only for the Sites one-shot fast path after the previews exist as local PNG, JPEG, or WebP files. The tool may be called sequentially for up to four distinct decisions; wait for each result before calling it again.",
          inputSchema: {
            type: "object",
            properties: {
              question: {
                type: "string",
                minLength: 1,
                description:
                  'Focused user-facing question, such as "Pick a color palette" or "Pick a typography pair".',
              },
              options: {
                type: "array",
                description:
                  "Exactly three comparable generated site-design options.",
                minItems: 3,
                maxItems: 3,
                items: {
                  type: "object",
                  properties: {
                    id: {
                      type: "string",
                      minLength: 1,
                      description: "Stable unique id for the direction.",
                    },
                    title: {
                      type: "string",
                      minLength: 1,
                      description: "Short user-facing direction name.",
                    },
                    designBrief: {
                      type: "string",
                      minLength: 1,
                      description:
                        "Implementation-ready visual and interaction guidance for the generated preview.",
                    },
                    imagePath: {
                      type: "string",
                      minLength: 1,
                      description:
                        "Absolute path to a generated PNG, JPEG, or WebP preview no larger than 8 MB.",
                    },
                  },
                  required: ["id", "title", "designBrief", "imagePath"],
                  additionalProperties: false,
                },
              },
            },
            required: ["options"],
            additionalProperties: false,
          },
          annotations: {
            readOnlyHint: true,
            destructiveHint: false,
            idempotentHint: true,
            openWorldHint: false,
          },
        },
      ],
    });
    return;
  }

  if (method === "tools/call") {
    try {
      await handleToolCall(id, params);
    } catch (error) {
      sendError(
        id,
        JsonRpcError.INVALID_PARAMS,
        error instanceof Error ? error.message : String(error),
      );
    }
    return;
  }

  if (id !== undefined) {
    sendError(id, JsonRpcError.METHOD_NOT_FOUND, `Method not found: ${method}`);
  }
}

const lines = readline.createInterface({
  input: process.stdin,
  crlfDelay: Infinity,
});

lines.on("line", (line) => {
  if (line.trim().length === 0) {
    return;
  }

  let message;
  try {
    message = JSON.parse(line);
  } catch {
    return;
  }

  if (message.method === undefined && message.id !== undefined) {
    const pending = pendingRequests.get(message.id);
    if (pending !== undefined) {
      pendingRequests.delete(message.id);
      if (message.error !== undefined) {
        pending.reject(
          new Error(message.error.message ?? "MCP request failed."),
        );
      } else {
        pending.resolve(message.result);
      }
    }
    return;
  }

  void handleRequest(message);
});
