// Image-to-Image tools — 10 models
import { z } from "zod";
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { submitAndPoll } from "../poller.js";
import { TOOL_ENDPOINTS, SIZE_MAP } from "../types.js";
import { resolveImageInput } from "../upload.js";

export function registerImageToImageTools(server: McpServer): void {
  // ─── General Image-to-Image ───
  server.tool(
    "image_to_image",
    "Transform an image using a text prompt (6 credits). General-purpose image-to-image. Accepts local file path or URL.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Text prompt describing the desired transformation"),
      output_format: z.enum(["png", "jpeg", "webp"]).default("jpeg"),
      seed: z.number().int().default(-1),
    },
    async ({ image, prompt, output_format, seed }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.image_to_image;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        image: imageUrl,
        prompt: prompt.trim(),
        output_format,
        seed,
      });
      return { content: [{ type: "text", text: `✅ Image transformed: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );

  // ─── Flux 2 Klein I2I ───
  server.tool(
    "flux_2_klein_i2i",
    "Transform image using Flux 2 Klein (2 credits). Controllable strength & size.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Transformation prompt"),
      aspect_ratio: z.enum(["1:1", "16:9", "3:2", "2:3", "3:4", "4:3", "9:16"]).default("1:1"),
      strength: z.number().min(0).max(1).default(0.75).describe("Transformation strength (0=preserve original, 1=full change)"),
      output_format: z.enum(["png", "jpeg", "webp"]).default("jpeg"),
      seed: z.number().int().default(-1),
    },
    async ({ image, prompt, aspect_ratio, strength, output_format, seed }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.flux_2_klein_i2i;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        image: imageUrl,
        prompt: prompt.trim(),
        size: SIZE_MAP[aspect_ratio] || SIZE_MAP["1:1"],
        strength,
        output_format,
        seed,
        enable_base64_output: false,
        enable_sync_mode: false,
      });
      return { content: [{ type: "text", text: `✅ Image transformed: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );

  // ─── Soul I2I ───
  server.tool(
    "soul_i2i",
    "Transform image using Soul model (2 credits). Artistic style transfer.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Transformation prompt"),
    },
    async ({ image, prompt }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.soul_i2i;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        image: imageUrl,
        prompt: prompt.trim(),
      });
      return { content: [{ type: "text", text: `✅ Image transformed: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );

  // ─── Z-Image Turbo I2I ───
  server.tool(
    "zimage_turbo_i2i",
    "Transform image using Z-Image Turbo (fast). Image-to-image transformation.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Transformation prompt"),
    },
    async ({ image, prompt }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.zimage_turbo_i2i;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        image: imageUrl,
        prompt: prompt.trim(),
      });
      return { content: [{ type: "text", text: `✅ Image transformed: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );

  // ─── Grok Imagine ───
  server.tool(
    "grok_imagine",
    "Transform image using Grok Imagine (6 credits). xAI's image transformation with quality/style controls.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Transformation prompt"),
      quality: z.enum(["auto", "high", "medium", "low"]).default("auto"),
      aspect_ratio: z.enum(["1:1", "16:9", "3:2", "2:3", "3:4", "4:3", "9:16"]).default("1:1"),
      strength: z.number().min(0).max(1).default(0.75).describe("Transformation strength"),
      style: z.string().default("").describe("Style preset (optional)"),
    },
    async ({ image, prompt, quality, aspect_ratio, strength, style }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.grok_imagine;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        image: imageUrl,
        prompt: prompt.trim(),
        quality,
        size: SIZE_MAP[aspect_ratio] || SIZE_MAP["1:1"],
        strength: Math.round(100 * strength) / 100,
        style,
      });
      return { content: [{ type: "text", text: `✅ Image transformed: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );

  // ─── Seedream 5.0 Lite I2I ───
  server.tool(
    "seedream_50_lite_i2i",
    "Transform image using Seedream 5.0 Lite (2 credits). ByteDance I2I with strength control.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Transformation prompt"),
      aspect_ratio: z.enum(["1:1", "16:9", "3:2", "2:3", "3:4", "4:3", "9:16"]).default("1:1"),
      strength: z.number().min(0).max(1).default(0.75),
      output_format: z.enum(["png", "jpeg", "webp"]).default("jpeg"),
      seed: z.number().int().default(-1),
    },
    async ({ image, prompt, aspect_ratio, strength, output_format, seed }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.seedream_50_lite_i2i;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        image: imageUrl,
        prompt: prompt.trim(),
        size: SIZE_MAP[aspect_ratio] || SIZE_MAP["1:1"],
        strength,
        output_format,
        seed,
        enable_base64_output: false,
        enable_sync_mode: false,
      });
      return { content: [{ type: "text", text: `✅ Image transformed: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );

  // ─── Qwen Image Max I2I ───
  server.tool(
    "qwen_max_i2i",
    "Transform image using Qwen Image Max (4 credits). Alibaba's powerful I2I model.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Transformation prompt"),
    },
    async ({ image, prompt }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.qwen_max_i2i;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        image: imageUrl,
        prompt: prompt.trim(),
      });
      return { content: [{ type: "text", text: `✅ Image transformed: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );

  // ─── Nano Banana I2I ───
  server.tool(
    "nano_banana_i2i",
    "Transform image using Nano Banana (2 credits). Image-to-image transformation.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Transformation prompt"),
      seed: z.number().int().default(-1),
    },
    async ({ image, prompt, seed }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.nano_banana_i2i;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        images: [imageUrl],
        prompt: prompt.trim(),
        enable_base64_output: false,
        enable_sync_mode: false,
        seed,
      });
      return { content: [{ type: "text", text: `✅ Image transformed: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );

  // ─── Nano Banana 2 I2I ───
  server.tool(
    "nano_banana_2_i2i",
    "Transform image using Nano Banana 2 (2 credits). Latest Nano Banana I2I.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Transformation prompt"),
      seed: z.number().int().default(-1),
    },
    async ({ image, prompt, seed }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.nano_banana_2_i2i;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        images: [imageUrl],
        prompt: prompt.trim(),
        enable_base64_output: false,
        enable_sync_mode: false,
        seed,
      });
      return { content: [{ type: "text", text: `✅ Image transformed: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );

  // ─── Nano Banana Pro Edit Ultra ───
  server.tool(
    "nano_banana_pro_edit",
    "Edit image using Nano Banana Pro Edit Ultra (2 credits). Advanced image editing.",
    {
      image: z.string().describe("Image URL or local file path"),
      prompt: z.string().describe("Edit instruction prompt"),
    },
    async ({ image, prompt }) => {
      const imageUrl = await resolveImageInput(image);
      const ep = TOOL_ENDPOINTS.nano_banana_pro_edit;
      const result = await submitAndPoll(ep.endpoint, ep.scenario, {
        images: [imageUrl],
        prompt: prompt.trim(),
      });
      return { content: [{ type: "text", text: `✅ Image edited: ${result.output}\nTask ID: ${result.taskId}` }] };
    },
  );
}