import { fromJsonSchema } from "@modelcontextprotocol/server";
import { zodToJsonSchema } from "zod-to-json-schema";
import type { z } from "zod";
import crypto from "node:crypto";
import { bridgeToolErrorResult, resolveScratchpadFolder, startStdioMcpServer } from "./core-bundle.mjs";
import { GenerateToolParamsSchemaMinimalStrict, UpscaleToolSchemaStrict } from "../core-bundle.mjs";
import { applyMcpConfig, initCustomConfigs, readMcpConfig } from "./config.js";
import { bindChatContextToScratchpad } from "./chatContextBridge.js";
import { assertNoAttachmentNotation } from "./sourceNotation.js";
import { logger } from "./mcpLogger.js";
import { extractSummary, materializeNewImages, snapshotImageIndices, writeHtmlReport } from "./resultMaterializer.js";
import { generateImageToolResult, generateSingleImageToolResult } from "./toolResults.js";

type ToolArgs = Record<string, unknown> & { scratchpadFolder: string };

/**
 * Adds the MCP-only scratchpadFolder field on top of the plugin's real Zod schema — no field is
 * hand-duplicated. Also hides `quality`, mirroring the plugin's own tools (generate_image.ts /
 * upscale.ts), which always run at the backend's "auto" default and never expose it to the agent.
 */
function toMcpInputSchema(zodSchema: z.ZodTypeAny): Record<string, unknown> {
  const { $schema: _$schema, ...schema } = zodToJsonSchema(zodSchema) as Record<string, any>;
  const { quality: _quality, ...visibleProperties } = (schema.properties ?? {}) as Record<string, unknown>;
  return {
    ...schema,
    properties: {
      scratchpadFolder: {
        type: "string",
        description: "Required scratchpad session folder. Obtain this value with get_scratchpad_folder and pass it back unchanged.",
      },
      ...visibleProperties,
    },
    required: ["scratchpadFolder", ...(((schema.required as string[] | undefined) ?? []).filter((key) => key !== "quality"))],
  };
}

const SOURCE_NOTATION_GUIDANCE =
  "Sources (canvas/moodboard): pN (picture), iN (image, including a file this tool generated earlier), a filename in the scratchpad folder, or an absolute accessible path. " +
  "aN (LM-Studio attachment notation) is not supported here — if a user attachment is needed and not yet in the scratchpad, call load_file_attachment first and use the resulting scratchpad filename instead.";

const GENERATE_IMAGE_DESCRIPTION = `Generate an image or video using Draw Things.

Before calling this tool: call get_scratchpad_folder and pass its return value unchanged as scratchpadFolder.

${SOURCE_NOTATION_GUIDANCE}`;

const UPSCALE_DESCRIPTION = `Re-render the canvas at a higher resolution via Draw Things image2image. No crop — uses the full canvas as-is.

Before calling this tool: call get_scratchpad_folder and pass its return value unchanged as scratchpadFolder.

${SOURCE_NOTATION_GUIDANCE}`;

// The MCP SDK's CallToolResult content-block union isn't re-exported for reuse here,
// so the handler result stays structurally typed (any) rather than hand-duplicating it.
function toToolResult(result: unknown): any {
  const record = result as { content?: unknown; isError?: unknown } | undefined;
  if (record && Array.isArray(record.content)) {
    return record.isError === true ? { content: record.content, isError: true } : { content: record.content };
  }
  return { content: [{ type: "text", text: typeof result === "string" ? result : JSON.stringify(result) }] };
}

/**
 * Forwards handleGenerateImage/handleUpscale's ProgressCallback to MCP's
 * notifications/progress, tied to the current tool call via ctx.mcpReq.notify
 * (which auto-attaches relatedRequestId). Only sends anything when the client
 * actually opted in by sending a progressToken with the request (per MCP spec) —
 * returns undefined otherwise, so no progress machinery runs for clients that
 * never asked for it. The status text is built with the exact same formula as
 * src/tools/generate_image.ts / src/tools/upscale.ts's onProgress (ctx.status(...)) so the
 * MCP path shows identical wording to the LM Studio plugin instead of a separately invented
 * format — no "total" is sent, since that's what made a client render its own, duplicate
 * "x/y (z%)" next to our already-formatted text.
 */
function createProgressForwarder(ctx: unknown, tool: string): ((step: number, totalSteps?: number, message?: string) => void) | undefined {
  const mcpReq = (ctx as { mcpReq?: { _meta?: { progressToken?: unknown }; notify?: (n: unknown) => Promise<void> } } | undefined)?.mcpReq;
  const progressToken = mcpReq?._meta?.progressToken;
  if (progressToken === undefined || typeof mcpReq?.notify !== "function") return undefined;
  let lastProgress = 0;
  return (step, totalSteps, message) => {
    let statusText: string;
    if (step === -1 && message) {
      statusText = message;
    } else if (totalSteps && totalSteps > 0) {
      const pct = Math.round((step / (totalSteps + 1)) * 100);
      statusText = `Step ${step}/${totalSteps} (${pct}%)`;
    } else {
      statusText = `Step ${step}...`;
    }
    if (step >= 0) {
      lastProgress = Math.max(lastProgress, step);
    } else if (lastProgress > 0 && typeof totalSteps === "number") {
      // A label-only event (step=-1) after real steps already ran is the wrap-up ("Finishing...",
      // or the tail of a 2nd pass) — count it as the (totalSteps + 1)th step, matching the "+1"
      // denominator the percentages above already use, so it lands on "done" instead of staying
      // stuck at the last sampling step's number.
      lastProgress = totalSteps + 1;
    }
    mcpReq.notify!({
      method: "notifications/progress",
      params: {
        progressToken,
        progress: lastProgress,
        message: statusText,
      },
    }).catch((error) => {
      logger.logError("progress-notify-failed", error, { tool }).catch(() => {});
    });
  };
}

/**
 * Runs a render call, then diffs chat_media_state.json to find the iN entries it just
 * appended (see resultMaterializer.ts) and turns them into Bionic-guided result text. For a
 * single result, no HTML report is written — the browser opens the file's own preview and
 * metadata is inlined from the tool call's summary. For 2+ results, an HTML report is still
 * written and referenced. Falls back to the renderer's raw content when nothing new was
 * recorded (e.g. an early validation error inside handleGenerateImage/handleUpscale itself).
 */
async function runAndMaterialize(
  tool: "generate_image" | "upscale",
  scratchpadPath: string,
  renderArgs: Record<string, unknown>,
  render: () => Promise<unknown>
): Promise<any> {
  const before = await snapshotImageIndices(scratchpadPath);
  const result = await render();
  const record = result as { content?: unknown; isError?: unknown } | undefined;
  if (record?.isError === true) return toToolResult(result);

  const materialized = await materializeNewImages(scratchpadPath, before);
  if (materialized.length === 0) return toToolResult(result);

  // Single-result case: no HTML report — the in-app browser opens the file's own preview, and
  // metadata is inlined directly from the tool call's own summary (see toolResults.ts).
  if (materialized.length === 1) {
    const summary = extractSummary(result);
    await logger.logEvent("tool-materialized", { tool, scratchpadPath, notations: materialized[0]!.notation });
    return { content: [{ type: "text", text: generateSingleImageToolResult(tool, materialized[0]!, summary) }] };
  }

  const requestId = crypto.randomUUID();
  const reportPath = await writeHtmlReport(
    scratchpadPath,
    requestId,
    {
      tool,
      prompt: typeof renderArgs.prompt === "string" ? renderArgs.prompt : undefined,
      canvas: typeof renderArgs.canvas === "string" ? renderArgs.canvas : undefined,
      summary: extractSummary(result),
    },
    materialized
  );
  await logger.logEvent("tool-materialized", { tool, scratchpadPath, notations: materialized.map((m) => m.notation).join(","), reportPath });
  return { content: [{ type: "text", text: generateImageToolResult(tool, reportPath, materialized) }] };
}

async function main(): Promise<void> {
  const config = readMcpConfig();
  applyMcpConfig(config);
  await initCustomConfigs(config.customConfigsPath);

  await logger.logEvent("MCP server starting", {
    drawThingsHost: config.drawThingsHost,
    drawThingsHttpPort: config.drawThingsHttpPort,
    drawThingsGrpcPort: config.drawThingsGrpcPort,
    previewInChat: config.previewInChat,
    embedPngMetadata: config.embedPngMetadata,
    chatWorkingDirectories: config.chatWorkingDirectories,
    nodeVersion: process.version,
  });

  startStdioMcpServer({
    name: "generate-image-mcp",
    version: "0.1.0",
    buildServer: (server) => {
      server.registerTool(
        "generate_image",
        { description: GENERATE_IMAGE_DESCRIPTION, inputSchema: fromJsonSchema(toMcpInputSchema(GenerateToolParamsSchemaMinimalStrict) as any) },
        async (args: unknown, ctx: unknown): Promise<any> => {
          const toolArgs = args as ToolArgs;
          try {
            const scratchpadPath = await resolveScratchpadFolder(config.chatWorkingDirectories, toolArgs.scratchpadFolder, logger);
            await logger.logEvent("tool-request", { tool: "generate_image", scratchpadPath });
            // quality is hidden from the schema above; also strip it here so it can never reach the
            // backend even if the caller sends it anyway — always resolves to the "auto" default.
            const { scratchpadFolder: _scratchpadFolder, quality: _quality, ...renderArgs } = toolArgs;
            assertNoAttachmentNotation(renderArgs);
            bindChatContextToScratchpad(scratchpadPath);
            const onProgress = createProgressForwarder(ctx, "generate_image");
            return await runAndMaterialize("generate_image", scratchpadPath, renderArgs, async () => {
              const { handleGenerateImage } = await import("../core/tools.js");
              return handleGenerateImage(renderArgs, onProgress);
            });
          } catch (error) {
            return await bridgeToolErrorResult("generate_image", error, logger);
          }
        }
      );

      server.registerTool(
        "upscale",
        { description: UPSCALE_DESCRIPTION, inputSchema: fromJsonSchema(toMcpInputSchema(UpscaleToolSchemaStrict) as any) },
        async (args: unknown, ctx: unknown): Promise<any> => {
          const toolArgs = args as ToolArgs;
          try {
            const scratchpadPath = await resolveScratchpadFolder(config.chatWorkingDirectories, toolArgs.scratchpadFolder, logger);
            await logger.logEvent("tool-request", { tool: "upscale", scratchpadPath });
            // quality is hidden from the schema above; also strip it here so it can never reach the
            // backend even if the caller sends it anyway — always resolves to the "auto" default.
            const { scratchpadFolder: _scratchpadFolder, quality: _quality, ...renderArgs } = toolArgs;
            assertNoAttachmentNotation(renderArgs);
            bindChatContextToScratchpad(scratchpadPath);
            const onProgress = createProgressForwarder(ctx, "upscale");
            return await runAndMaterialize("upscale", scratchpadPath, renderArgs, async () => {
              const { handleUpscale } = await import("../core/tools.js");
              return handleUpscale(renderArgs, onProgress);
            });
          } catch (error) {
            return await bridgeToolErrorResult("upscale", error, logger);
          }
        }
      );
    },
  });
}

main();

import { fromJsonSchema } from "@modelcontextprotocol/server";
import { zodToJsonSchema } from "zod-to-json-schema";
import type { z } from "zod";
import crypto from "node:crypto";
import { bridgeToolErrorResult, resolveScratchpadFolder, startStdioMcpServer } from "./core-bundle.mjs";
import { GenerateToolParamsSchemaMinimalStrict, UpscaleToolSchemaStrict } from "../core-bundle.mjs";
import { applyMcpConfig, initCustomConfigs, readMcpConfig } from "./config.js";
import { bindChatContextToScratchpad } from "./chatContextBridge.js";
import { assertNoAttachmentNotation } from "./sourceNotation.js";
import { logger } from "./mcpLogger.js";
import { extractSummary, materializeNewImages, snapshotImageIndices, writeHtmlReport } from "./resultMaterializer.js";
import { generateImageToolResult, generateSingleImageToolResult } from "./toolResults.js";

type ToolArgs = Record<string, unknown> & { scratchpadFolder: string };

/**
 * Adds the MCP-only scratchpadFolder field on top of the plugin's real Zod schema — no field is
 * hand-duplicated. Also hides `quality`, mirroring the plugin's own tools (generate_image.ts /
 * upscale.ts), which always run at the backend's "auto" default and never expose it to the agent.
 */
function toMcpInputSchema(zodSchema: z.ZodTypeAny): Record<string, unknown> {
  const { $schema: _$schema, ...schema } = zodToJsonSchema(zodSchema) as Record<string, any>;
  const { quality: _quality, ...visibleProperties } = (schema.properties ?? {}) as Record<string, unknown>;
  return {
    ...schema,
    properties: {
      scratchpadFolder: {
        type: "string",
        description: "Required scratchpad session folder. Obtain this value with get_scratchpad_folder and pass it back unchanged.",
      },
      ...visibleProperties,
    },
    required: ["scratchpadFolder", ...(((schema.required as string[] | undefined) ?? []).filter((key) => key !== "quality"))],
  };
}

const SOURCE_NOTATION_GUIDANCE =
  "Sources (canvas/moodboard): pN (picture), iN (image, including a file this tool generated earlier), a filename in the scratchpad folder, or an absolute accessible path. " +
  "aN (LM-Studio attachment notation) is not supported here — if a user attachment is needed and not yet in the scratchpad, call load_file_attachment first and use the resulting scratchpad filename instead.";

const GENERATE_IMAGE_DESCRIPTION = `Generate an image or video using Draw Things.

Before calling this tool: call get_scratchpad_folder and pass its return value unchanged as scratchpadFolder.

${SOURCE_NOTATION_GUIDANCE}`;

const UPSCALE_DESCRIPTION = `Re-render the canvas at a higher resolution via Draw Things image2image. No crop — uses the full canvas as-is.

Before calling this tool: call get_scratchpad_folder and pass its return value unchanged as scratchpadFolder.

${SOURCE_NOTATION_GUIDANCE}`;

// The MCP SDK's CallToolResult content-block union isn't re-exported for reuse here,
// so the handler result stays structurally typed (any) rather than hand-duplicating it.
function toToolResult(result: unknown): any {
  const record = result as { content?: unknown; isError?: unknown } | undefined;
  if (record && Array.isArray(record.content)) {
    return record.isError === true ? { content: record.content, isError: true } : { content: record.content };
  }
  return { content: [{ type: "text", text: typeof result === "string" ? result : JSON.stringify(result) }] };
}

/**
 * Forwards handleGenerateImage/handleUpscale's ProgressCallback to MCP's
 * notifications/progress, tied to the current tool call via ctx.mcpReq.notify
 * (which auto-attaches relatedRequestId). Only sends anything when the client
 * actually opted in by sending a progressToken with the request (per MCP spec) —
 * returns undefined otherwise, so no progress machinery runs for clients that
 * never asked for it. The status text is built with the exact same formula as
 * src/tools/generate_image.ts / src/tools/upscale.ts's onProgress (ctx.status(...)) so the
 * MCP path shows identical wording to the LM Studio plugin instead of a separately invented
 * format — no "total" is sent, since that's what made a client render its own, duplicate
 * "x/y (z%)" next to our already-formatted text.
 */
function createProgressForwarder(ctx: unknown, tool: string): ((step: number, totalSteps?: number, message?: string) => void) | undefined {
  const mcpReq = (ctx as { mcpReq?: { _meta?: { progressToken?: unknown }; notify?: (n: unknown) => Promise<void> } } | undefined)?.mcpReq;
  const progressToken = mcpReq?._meta?.progressToken;
  if (progressToken === undefined || typeof mcpReq?.notify !== "function") return undefined;
  let lastProgress = 0;
  return (step, totalSteps, message) => {
    let statusText: string;
    if (step === -1 && message) {
      statusText = message;
    } else if (totalSteps && totalSteps > 0) {
      const pct = Math.round((step / (totalSteps + 1)) * 100);
      statusText = `Step ${step}/${totalSteps} (${pct}%)`;
    } else {
      statusText = `Step ${step}...`;
    }
    if (step >= 0) {
      lastProgress = Math.max(lastProgress, step);
    } else if (lastProgress > 0 && typeof totalSteps === "number") {
      // A label-only event (step=-1) after real steps already ran is the wrap-up ("Finishing...",
      // or the tail of a 2nd pass) — count it as the (totalSteps + 1)th step, matching the "+1"
      // denominator the percentages above already use, so it lands on "done" instead of staying
      // stuck at the last sampling step's number.
      lastProgress = totalSteps + 1;
    }
    mcpReq.notify!({
      method: "notifications/progress",
      params: {
        progressToken,
        progress: lastProgress,
        message: statusText,
      },
    }).catch((error) => {
      logger.logError("progress-notify-failed", error, { tool }).catch(() => {});
    });
  };
}

/**
 * Runs a render call, then diffs chat_media_state.json to find the iN entries it just
 * appended (see resultMaterializer.ts) and turns them into Bionic-guided result text. For a
 * single result, no HTML report is written — the browser opens the file's own preview and
 * metadata is inlined from the tool call's summary. For 2+ results, an HTML report is still
 * written and referenced. Falls back to the renderer's raw content when nothing new was
 * recorded (e.g. an early validation error inside handleGenerateImage/handleUpscale itself).
 */
async function runAndMaterialize(
  tool: "generate_image" | "upscale",
  scratchpadPath: string,
  renderArgs: Record<string, unknown>,
  render: () => Promise<unknown>
): Promise<any> {
  const before = await snapshotImageIndices(scratchpadPath);
  const result = await render();
  const record = result as { content?: unknown; isError?: unknown } | undefined;
  if (record?.isError === true) return toToolResult(result);

  const materialized = await materializeNewImages(scratchpadPath, before);
  if (materialized.length === 0) return toToolResult(result);

  // Single-result case: no HTML report — the in-app browser opens the file's own preview, and
  // metadata is inlined directly from the tool call's own summary (see toolResults.ts).
  if (materialized.length === 1) {
    const summary = extractSummary(result);
    await logger.logEvent("tool-materialized", { tool, scratchpadPath, notations: materialized[0]!.notation });
    return { content: [{ type: "text", text: generateSingleImageToolResult(tool, materialized[0]!, summary) }] };
  }

  const requestId = crypto.randomUUID();
  const reportPath = await writeHtmlReport(
    scratchpadPath,
    requestId,
    {
      tool,
      prompt: typeof renderArgs.prompt === "string" ? renderArgs.prompt : undefined,
      canvas: typeof renderArgs.canvas === "string" ? renderArgs.canvas : undefined,
      summary: extractSummary(result),
    },
    materialized
  );
  await logger.logEvent("tool-materialized", { tool, scratchpadPath, notations: materialized.map((m) => m.notation).join(","), reportPath });
  return { content: [{ type: "text", text: generateImageToolResult(tool, reportPath, materialized) }] };
}

async function main(): Promise<void> {
  const config = readMcpConfig();
  applyMcpConfig(config);
  await initCustomConfigs(config.customConfigsPath);

  await logger.logEvent("MCP server starting", {
    drawThingsHost: config.drawThingsHost,
    drawThingsHttpPort: config.drawThingsHttpPort,
    drawThingsGrpcPort: config.drawThingsGrpcPort,
    previewInChat: config.previewInChat,
    embedPngMetadata: config.embedPngMetadata,
    chatWorkingDirectories: config.chatWorkingDirectories,
    nodeVersion: process.version,
  });

  startStdioMcpServer({
    name: "generate-image-mcp",
    version: "0.1.0",
    buildServer: (server) => {
      server.registerTool(
        "generate_image",
        { description: GENERATE_IMAGE_DESCRIPTION, inputSchema: fromJsonSchema(toMcpInputSchema(GenerateToolParamsSchemaMinimalStrict) as any) },
        async (args: unknown, ctx: unknown): Promise<any> => {
          const toolArgs = args as ToolArgs;
          try {
            const scratchpadPath = await resolveScratchpadFolder(config.chatWorkingDirectories, toolArgs.scratchpadFolder, logger);
            await logger.logEvent("tool-request", { tool: "generate_image", scratchpadPath });
            // quality is hidden from the schema above; also strip it here so it can never reach the
            // backend even if the caller sends it anyway — always resolves to the "auto" default.
            const { scratchpadFolder: _scratchpadFolder, quality: _quality, ...renderArgs } = toolArgs;
            assertNoAttachmentNotation(renderArgs);
            bindChatContextToScratchpad(scratchpadPath);
            const onProgress = createProgressForwarder(ctx, "generate_image");
            return await runAndMaterialize("generate_image", scratchpadPath, renderArgs, async () => {
              const { handleGenerateImage } = await import("../core/tools.js");
              return handleGenerateImage(renderArgs, onProgress);
            });
          } catch (error) {
            return await bridgeToolErrorResult("generate_image", error, logger);
          }
        }
      );

      server.registerTool(
        "upscale",
        { description: UPSCALE_DESCRIPTION, inputSchema: fromJsonSchema(toMcpInputSchema(UpscaleToolSchemaStrict) as any) },
        async (args: unknown, ctx: unknown): Promise<any> => {
          const toolArgs = args as ToolArgs;
          try {
            const scratchpadPath = await resolveScratchpadFolder(config.chatWorkingDirectories, toolArgs.scratchpadFolder, logger);
            await logger.logEvent("tool-request", { tool: "upscale", scratchpadPath });
            // quality is hidden from the schema above; also strip it here so it can never reach the
            // backend even if the caller sends it anyway — always resolves to the "auto" default.
            const { scratchpadFolder: _scratchpadFolder, quality: _quality, ...renderArgs } = toolArgs;
            assertNoAttachmentNotation(renderArgs);
            bindChatContextToScratchpad(scratchpadPath);
            const onProgress = createProgressForwarder(ctx, "upscale");
            return await runAndMaterialize("upscale", scratchpadPath, renderArgs, async () => {
              const { handleUpscale } = await import("../core/tools.js");
              return handleUpscale(renderArgs, onProgress);
            });
          } catch (error) {
            return await bridgeToolErrorResult("upscale", error, logger);
          }
        }
      );
    },
  });
}

main();