#!/usr/bin/env node

// Z-Image AI MCP Server
// Exposes all z-image.ai services as MCP tools for Copilot CLI and other MCP clients

import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { interactiveAuth, directAuth, getSessionToken } from "./auth.js";
import { registerTextToImageTools } from "./tools/text-to-image.js";
import { registerImageToImageTools } from "./tools/image-to-image.js";
import { registerImageUtilTools } from "./tools/image-utils.js";
import { registerTextToVideoTools } from "./tools/text-to-video.js";
import { registerImageToVideoTools } from "./tools/image-to-video.js";
import { registerVideoTools } from "./tools/video-tools.js";
import { registerUserTools } from "./tools/user.js";
import { registerUtilityTools } from "./tools/utility.js";

// ─── CLI: handle `auth` subcommand before starting MCP server ───
async function handleCli(): Promise<boolean> {
  const args = process.argv.slice(2);

  if (args[0] === "auth") {
    // Direct token mode: zimage-mcp auth --token <value>
    const tokenIdx = args.indexOf("--token");
    if (tokenIdx !== -1 && args[tokenIdx + 1]) {
      directAuth(args[tokenIdx + 1]);
      console.error("✅ Session token saved.");
      return true;
    }

    // Cookie string mode: zimage-mcp auth --cookie "<cookie string>"
    const cookieIdx = args.indexOf("--cookie");
    if (cookieIdx !== -1 && args[cookieIdx + 1]) {
      directAuth(args[cookieIdx + 1]);
      console.error("✅ Session token extracted and saved.");
      return true;
    }

    // Interactive browser-based auth
    try {
      await interactiveAuth();
      console.error("✅ Authentication successful! Session saved.");
      console.error("   You can now use zimage-mcp as an MCP server.");
    } catch (err) {
      console.error(`❌ Auth failed: ${(err as Error).message}`);
      process.exit(1);
    }
    return true;
  }

  if (args[0] === "status") {
    const token = getSessionToken();
    if (token) {
      console.error(`✅ Authenticated (token: ${token.slice(0, 8)}...${token.slice(-4)})`);
    } else {
      console.error("❌ Not authenticated. Run: zimage-mcp auth");
    }
    return true;
  }

  if (args[0] === "help" || args[0] === "--help" || args[0] === "-h") {
    console.error(`
zimage-mcp — Z-Image AI MCP Server

Usage:
  zimage-mcp              Start the MCP server (stdio transport)
  zimage-mcp auth         Interactive browser login (opens z-image.ai)
  zimage-mcp auth --token <session_token>   Set token directly
  zimage-mcp auth --cookie "<cookie_str>"   Extract token from cookie string
  zimage-mcp status       Check authentication status
  zimage-mcp help         Show this help

Environment:
  ZIMAGE_SESSION_TOKEN    Set session token via environment variable

MCP Configuration (add to Copilot CLI / MCP client config):
  {
    "mcpServers": {
      "zimage": {
        "command": "node",
        "args": ["${process.argv[1]}"]
      }
    }
  }

Tools (40 total):
  Text-to-Image (10): zimage_generate, zimage_turbo_t2i, nano_banana_t2i,
    nano_banana_pro_t2i, nano_banana_2_t2i, seedream_v45_t2i,
    seedream_50_lite_t2i, flux_2_klein_t2i, gpt_image_1_t2i, qwen_t2i

  Image-to-Image (10): image_to_image, flux_2_klein_i2i, soul_i2i,
    zimage_turbo_i2i, grok_imagine, seedream_50_lite_i2i, qwen_max_i2i,
    nano_banana_i2i, nano_banana_2_i2i, nano_banana_pro_edit

  Image Utils (4): image_eraser, background_remover, image_upscaler, image_layered

  Text-to-Video (5): wan_26_t2v, ltx2_fast_t2v, veo_31_t2v,
    seedance_15_pro_t2v, pixverse_56_t2v

  Image-to-Video (5): ltx2_fast_i2v, veo_31_i2v, seedance_15_pro_i2v,
    pixverse_56_i2v, vidu_q3_i2v

  Video Special (3): motion_control_kling, video_face_swap, longcat_avatar

  Account (3): get_user_info, get_credits, get_config

  Utility (5): prompt_enhance, estimate_credits, upload_image,
    upload_video, upload_audio
`);
    return true;
  }

  return false;
}

// ─── Main: start MCP server ───
async function main(): Promise<void> {
  // Handle CLI subcommands first
  if (await handleCli()) {
    process.exit(0);
  }

  // Create MCP server
  const server = new McpServer({
    name: "zimage-mcp",
    version: "0.1.0",
  });

  // Register all tool categories
  registerTextToImageTools(server);
  registerImageToImageTools(server);
  registerImageUtilTools(server);
  registerTextToVideoTools(server);
  registerImageToVideoTools(server);
  registerVideoTools(server);
  registerUserTools(server);
  registerUtilityTools(server);

  // Connect via stdio transport
  const transport = new StdioServerTransport();
  await server.connect(transport);

  console.error("🚀 Z-Image MCP server running (stdio transport)");
}

main().catch((err) => {
  console.error("Fatal error:", err);
  process.exit(1);
});