#!/usr/bin/env node
import "dotenv/config";
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
  ErrorCode,
  McpError,
} from "@modelcontextprotocol/sdk/types.js";
import axios, { AxiosError } from "axios";
import { z } from "zod";
import FormData from "form-data";

type SeedanceAspectRatio = "1:1" | "21:9" | "4:3" | "3:4" | "16:9" | "9:16";
type SeedanceResolution = "480p" | "720p" | "1080p";

const FIXED_DURATION = "8";
const FIXED_RESOLUTION: SeedanceResolution = "480p";
const FIXED_GENERATE_AUDIO = true;
const DEFAULT_SIZE = "854x480";

const API_KEY =
  process.env.VIDEO_API_KEY || "";
const BASE_URL = process.env.VIDEO_BASE_URL || "https://api.kie.ai";
const UPLOAD_URL =
  "https://www.mcpcn.cc/api/fileUploadAndDownload/uploadMcpFile";
const KIE_CALLBACK_URL = process.env.KIE_CALLBACK_URL || "";
const VIDEO_MODEL = process.env.VIDEO_MODEL || "bytedance/seedance-2-fast";
const VIDEO_TEXT_MODEL = process.env.VIDEO_TEXT_MODEL || VIDEO_MODEL;
const VIDEO_IMAGE_MODEL = process.env.VIDEO_IMAGE_MODEL || VIDEO_MODEL;

interface VideoCreateResponse {
  id: string;
  object: string;
  model: string;
  status: string;
  progress: number;
  created_at: number;
  seconds: string;
  size: string;
  status_update_time?: number;
}

interface VideoQueryResponse {
  id: string;
  object?: string;
  model?: string;
  status: string;
  progress?: number;
  created_at?: number;
  completed_at?: number;
  seconds?: string;
  size?: string;
  video_url?: string | null;
  enhanced_prompt?: string;
  status_update_time?: number;
  detail?: any;
  error?: {
    message: string;
  };
  amount?: number;
}

interface KieCreateTaskRequest {
  model: string;
  callBackUrl?: string;
  input: {
    prompt: string;
    aspect_ratio: SeedanceAspectRatio;
    resolution: SeedanceResolution;
    duration: string;
    fixed_lens: boolean;
    generate_audio: boolean;
    input_urls?: string[];
  };
}

interface KieCreateTaskResponse {
  code: number;
  message?: string;
  msg?: string;
  data?: {
    taskId?: string;
  };
}

interface KieTaskRecord {
  taskId: string;
  model?: string;
  state?: string;
  param?: string;
  resultJson?: string;
  failCode?: string;
  failMsg?: string;
  costTime?: number;
  completeTime?: number;
  createTime?: number;
}

interface KieQueryTaskResponse {
  code: number;
  message?: string;
  msg?: string;
  data?: KieTaskRecord;
}

const createVideoSchema = z.object({
  prompt: z.string().describe("Text prompt describing the video content to generate"),
  size: z
    .string()
    .default(DEFAULT_SIZE)
    .describe('Video dimensions in "WxH" format, e.g. 1280x720 or 720x1280'),
  resolution: z
    .enum(["480p", "720p"])
    .default("480p")
    .describe('Video resolution, either "480p" or "720p"'),
});

const createVideoWithImagesSchema = z.object({
  prompt: z.string().describe("Text prompt describing how to animate the reference image"),
  input_reference: z.string().url().describe("URL of the reference image for image-to-video generation"),
  size: z
    .string()
    .default(DEFAULT_SIZE)
    .describe('Video dimensions in "WxH" format, e.g. 1280x720 or 720x1280'),
  resolution: z
    .enum(["480p", "720p"])
    .default("480p")
    .describe('Video resolution, either "480p" or "720p"'),
});

const queryTaskSchema = z.object({
  id: z
    .string()
    .describe("Video task ID, e.g. task_12345678"),
});

const getVideoContentSchema = z.object({
  id: z
    .string()
    .describe("Video task ID, e.g. task_12345678"),
});

class SoraVideoMCPServer {
  private server: Server;
  private axiosInstance = axios.create({
    baseURL: BASE_URL,
    timeout: 30000,
    headers: {
      Authorization: `Bearer ${API_KEY}`,
      "Content-Type": "application/json",
      Accept: "application/json",
    },
  });

  constructor() {
    this.server = new Server(
      {
        name: "sora-video-mcp",
        version: "1.0.0",
      },
      {
        capabilities: {
          tools: {},
        },
      },
    );

    this.setupHandlers();
    this.setupErrorHandling();
  }

  private setupErrorHandling(): void {
    this.server.onerror = (error) => {
      console.error("[MCP Error]", error);
    };

    process.on("SIGINT", async () => {
      await this.server.close();
      process.exit(0);
    });
  }

  private setupHandlers(): void {
    this.server.setRequestHandler(ListToolsRequestSchema, async () => ({
      tools: [
        {
          name: "create_video",
          description:
            "Submit a text-to-video generation task. Returns a taskId immediately. Use query_video_task to check progress.",
          inputSchema: {
            type: "object",
            properties: {
              prompt: {
                type: "string",
                description:
                  'Text description of the video. Include scene, action, style, and camera details. Example: "A cute orange cat playing on a sunny meadow, slow-motion, cinematic quality"',
              },
              size: {
                type: "string",
                description:
                  'Video dimensions in "WxH" format, e.g. 1280x720 or 720x1280',
                default: DEFAULT_SIZE,
              },
              resolution: {
                type: "string",
                enum: ["480p", "720p"],
                description: 'Video resolution, either "480p" or "720p"',
                default: "480p",
              },
            },
            required: ["prompt"],
          },
        },
        {
          name: "create_video_with_images",
          description:
            "Submit an image-to-video generation task. Animates a reference image based on the text prompt. Returns a taskId immediately. Use query_video_task to check progress.",
          inputSchema: {
            type: "object",
            properties: {
              prompt: {
                type: "string",
                description:
                  'Description of how to animate the image. Include motion, transitions, and style. Example: "Animate the person with natural breathing and a subtle smile"',
              },
              input_reference: {
                type: "string",
                description:
                  "URL of the reference image. Must be a publicly accessible HTTP/HTTPS link (jpg, png, etc.).",
              },
              size: {
                type: "string",
                description:
                  'Video dimensions in "WxH" format, e.g. 1280x720 or 720x1280',
                default: DEFAULT_SIZE,
              },
              resolution: {
                type: "string",
                enum: ["480p", "720p"],
                description: 'Video resolution, either "480p" or "720p"',
                default: "480p",
              },
            },
            required: ["prompt", "input_reference"],
          },
        },
        {
          name: "query_video_task",
          description:
            "Query the status of a video generation task by taskId. Returns status (in_progress / completed / failed) and video_url when complete.",
          inputSchema: {
            type: "object",
            properties: {
              id: {
                type: "string",
                description:
                  'Unique task ID returned by create_video or create_video_with_images, e.g. "task_12345678".',
              },
            },
            required: ["id"],
          },
        },
        {
          name: "get_video_content",
          description:
            "Download the completed video and re-upload it to a custom storage server. Only needed when migrating video to a custom storage service.",
          inputSchema: {
            type: "object",
            properties: {
              id: {
                type: "string",
                description:
                  'Unique task ID, e.g. "task_12345678".',
              },
            },
            required: ["id"],
          },
        },
      ],
    }));

    this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
      try {
        if (!API_KEY) {
          throw new McpError(
            ErrorCode.InvalidRequest,
            "VIDEO_API_KEY environment variable is not set",
          );
        }

        switch (request.params.name) {
          case "create_video":
            return await this.handleCreateVideo(request.params.arguments);
          case "create_video_with_images":
            return await this.handleCreateVideoWithImages(
              request.params.arguments,
            );
          case "query_video_task":
            return await this.handleQueryTask(request.params.arguments);
          case "get_video_content":
            return await this.handleGetVideoContent(request.params.arguments);
          default:
            throw new McpError(
              ErrorCode.MethodNotFound,
              `Unknown tool: ${request.params.name}`,
            );
        }
      } catch (error) {
        if (error instanceof McpError) {
          throw error;
        }

        if (axios.isAxiosError(error)) {
          const axiosError = error as AxiosError;

          if (
            axiosError.code === "ECONNABORTED" ||
            axiosError.message.includes("timeout")
          ) {
            throw new McpError(
              ErrorCode.InternalError,
              "API request timed out. Check network connection and retry.",
            );
          }

          if (
            axiosError.code === "ECONNREFUSED" ||
            axiosError.code === "ENOTFOUND"
          ) {
            throw new McpError(
              ErrorCode.InternalError,
              "Cannot connect to API server. Check network or BASE_URL configuration.",
            );
          }

          const status = axiosError.response?.status;
          const data = axiosError.response?.data;

          if (status === 401) {
            throw new McpError(
              ErrorCode.InvalidRequest,
              "API authentication failed. Check VIDEO_API_KEY.",
            );
          }

          if (status === 429) {
            throw new McpError(
              ErrorCode.InternalError,
              "API rate limit exceeded. Please retry later.",
            );
          }

          throw new McpError(
            ErrorCode.InternalError,
            `API request failed (${status}): ${JSON.stringify(data) || axiosError.message}`,
          );
        }

        throw new McpError(
          ErrorCode.InternalError,
          `Error: ${error instanceof Error ? error.message : String(error)}`,
        );
      }
    });
  }

  private mapSizeToAspectRatio(size: string): SeedanceAspectRatio {
    const [width, height] = size.split("x").map(Number);
    if (!Number.isFinite(width) || !Number.isFinite(height)) {
      return "16:9";
    }

    const ratio = width / height;
    const supportedRatios: Array<{
      aspectRatio: SeedanceAspectRatio;
      value: number;
    }> = [
        { aspectRatio: "1:1", value: 1 },
        { aspectRatio: "21:9", value: 21 / 9 },
        { aspectRatio: "4:3", value: 4 / 3 },
        { aspectRatio: "3:4", value: 3 / 4 },
        { aspectRatio: "16:9", value: 16 / 9 },
        { aspectRatio: "9:16", value: 9 / 16 },
      ];

    return supportedRatios.reduce((closest, current) =>
      Math.abs(current.value - ratio) < Math.abs(closest.value - ratio)
        ? current
        : closest
    ).aspectRatio;
  }

  private mapAspectRatioToSize(
    aspectRatio?: string,
    resolution: SeedanceResolution = "720p",
  ): string | undefined {
    const sizeMap: Record<
      SeedanceAspectRatio,
      Record<SeedanceResolution, string>
    > = {
      "1:1": {
        "480p": "480x480",
        "720p": "720x720",
        "1080p": "1080x1080",
      },
      "21:9": {
        "480p": "1120x480",
        "720p": "1680x720",
        "1080p": "2520x1080",
      },
      "4:3": {
        "480p": "640x480",
        "720p": "960x720",
        "1080p": "1440x1080",
      },
      "3:4": {
        "480p": "480x640",
        "720p": "720x960",
        "1080p": "1080x1440",
      },
      "16:9": {
        "480p": "854x480",
        "720p": "1280x720",
        "1080p": "1920x1080",
      },
      "9:16": {
        "480p": "480x854",
        "720p": "720x1280",
        "1080p": "1080x1920",
      },
    };

    if (!aspectRatio || !(aspectRatio in sizeMap)) {
      return undefined;
    }

    return sizeMap[aspectRatio as SeedanceAspectRatio][resolution];
  }

  private normalizeVideoModel(model?: string): string {
    if (!model) {
      return VIDEO_MODEL;
    }

    return model.replace(/-(text|image)-to-video$/, "");
  }

  private mapKieStateToStatus(
    state?: string,
  ): "in_progress" | "completed" | "failed" {
    const normalizedState = (state || "").toLowerCase();

    if (
      normalizedState === "success" ||
      normalizedState === "succeeded" ||
      normalizedState === "completed"
    ) {
      return "completed";
    }

    if (
      normalizedState === "waiting" ||
      normalizedState === "pending" ||
      normalizedState === "processing"
    ) {
      return "in_progress";
    }

    if (
      normalizedState === "failed" ||
      normalizedState === "fail" ||
      normalizedState === "error" ||
      normalizedState === "cancelled"
    ) {
      return "failed";
    }

    return "in_progress";
  }

  private parseKieResultJson(resultJson?: string): { videoUrl: string | null } {
    if (!resultJson) {
      return { videoUrl: null };
    }

    try {
      const parsed = JSON.parse(resultJson) as {
        resultUrls?: unknown;
      };
      const resultUrls = parsed.resultUrls;

      if (Array.isArray(resultUrls)) {
        const firstUrl = resultUrls.find((item) => typeof item === "string");
        if (typeof firstUrl === "string") {
          return { videoUrl: firstUrl };
        }
      }
    } catch (error) {
      console.error("Failed to parse Kie resultJson:", error);
    }

    return { videoUrl: null };
  }

  private extractTaskMetaFromParam(param?: string): {
    seconds?: string;
    size?: string;
    resolution?: string;
    isWithVideo?: boolean;
  } {
    if (!param) {
      return {};
    }

    try {
      const parsed = JSON.parse(param) as {
        input?: {
          duration?: string | number;
          n_frames?: string | number;
          aspect_ratio?: string;
          resolution?: string;
          input_urls?: string[];
        };
      };

      const seconds =
        typeof parsed.input?.duration === "number"
          ? String(parsed.input.duration)
          : typeof parsed.input?.duration === "string"
            ? parsed.input.duration
            : typeof parsed.input?.n_frames === "number"
              ? String(parsed.input.n_frames)
              : typeof parsed.input?.n_frames === "string"
                ? parsed.input.n_frames
                : undefined;

      const resolution = parsed.input?.resolution;
      const size = this.mapAspectRatioToSize(
        parsed.input?.aspect_ratio,
        resolution === "480p" || resolution === "720p" || resolution === "1080p"
          ? resolution
          : "720p",
      );
      const isWithVideo = Array.isArray(parsed.input?.input_urls) && parsed.input.input_urls.length > 0;

      return {
        seconds,
        size,
        resolution,
        isWithVideo,
      };
    } catch (error) {
      console.error("Failed to parse Kie param:", error);
      return {};
    }
  }

  private getKieResponseMessage(response: {
    message?: string;
    msg?: string;
  }): string {
    return response.msg || response.message || "Unknown error";
  }

  private async handleCreateVideo(args: unknown) {
    const params = createVideoSchema.parse(args);
    const aspectRatio = this.mapSizeToAspectRatio(params.size);
    const resolution = params.resolution as SeedanceResolution;
    const outputSize =
      this.mapAspectRatioToSize(aspectRatio, resolution) || DEFAULT_SIZE;

    const requestBody: KieCreateTaskRequest = {
      model: VIDEO_TEXT_MODEL,
      input: {
        prompt: params.prompt,
        aspect_ratio: aspectRatio,
        resolution: resolution,
        duration: FIXED_DURATION,
        fixed_lens: false,
        generate_audio: FIXED_GENERATE_AUDIO,
      },
    };

    if (KIE_CALLBACK_URL) {
      requestBody.callBackUrl = KIE_CALLBACK_URL;
    }

    const response = await axios.post<KieCreateTaskResponse>(
      `${BASE_URL}/api/v1/jobs/createTask`,
      requestBody,
      {
        headers: {
          "Content-Type": "application/json",
          Authorization: `Bearer ${API_KEY}`,
        },
        timeout: 30000,
      },
    );

    const taskId = response.data.data?.taskId;
    if (response.data.code !== 200 || !taskId) {
      throw new Error(
        `Failed to create video task: ${this.getKieResponseMessage(response.data)}`,
      );
    }
    console.error(`Video task submitted. Task ID: ${taskId}`);

    return {
      content: [
        {
          type: "text",
          text: JSON.stringify(
            {
              success: true,
              message: "Video generation task submitted.",
              task_id: taskId,
              note: "Use query_video_task with this task_id to check progress (processing takes ~1-5 minutes).",
              data: {
                id: taskId,
                object: "video",
                model: VIDEO_MODEL,
                status: "in_progress",
                progress: 0,
                created_at: Date.now(),
                seconds: FIXED_DURATION,
                size: outputSize,
              } satisfies VideoCreateResponse,
            },
            null,
            2,
          ),
        },
      ],
    };
  }

  private async handleCreateVideoWithImages(args: unknown) {
    const params = createVideoWithImagesSchema.parse(args);
    const aspectRatio = this.mapSizeToAspectRatio(params.size);
    const resolution = params.resolution as SeedanceResolution;
    const outputSize =
      this.mapAspectRatioToSize(aspectRatio, resolution) || DEFAULT_SIZE;

    const requestBody: KieCreateTaskRequest = {
      model: VIDEO_IMAGE_MODEL,
      input: {
        prompt: params.prompt,
        input_urls: [params.input_reference],
        aspect_ratio: aspectRatio,
        resolution: resolution,
        duration: FIXED_DURATION,
        fixed_lens: false,
        generate_audio: FIXED_GENERATE_AUDIO,
      },
    };

    if (KIE_CALLBACK_URL) {
      requestBody.callBackUrl = KIE_CALLBACK_URL;
    }

    const response = await axios.post<KieCreateTaskResponse>(
      `${BASE_URL}/api/v1/jobs/createTask`,
      requestBody,
      {
        headers: {
          "Content-Type": "application/json",
          Authorization: `Bearer ${API_KEY}`,
        },
        timeout: 30000,
      },
    );

    const taskId = response.data.data?.taskId;
    if (response.data.code !== 200 || !taskId) {
      throw new Error(
        `Failed to create image-to-video task: ${this.getKieResponseMessage(response.data)}`,
      );
    }
    console.error(`Image-to-video task submitted. Task ID: ${taskId}`);

    return {
      content: [
        {
          type: "text",
          text: JSON.stringify(
            {
              success: true,
              message: "Image-to-video task submitted.",
              task_id: taskId,
              note: "Use query_video_task with this task_id to check progress (processing takes ~1-5 minutes).",
              data: {
                id: taskId,
                object: "video",
                model: VIDEO_MODEL,
                status: "in_progress",
                progress: 0,
                created_at: Date.now(),
                seconds: FIXED_DURATION,
                size: outputSize,
              } satisfies VideoCreateResponse,
            },
            null,
            2,
          ),
        },
      ],
    };
  }

  private async handleQueryTask(args: unknown) {
    const params = queryTaskSchema.parse(args);

    const response = await axios.get<KieQueryTaskResponse>(
      `${BASE_URL}/api/v1/jobs/recordInfo`,
      {
        params: {
          taskId: params.id,
        },
        headers: {
          Authorization: `Bearer ${API_KEY}`,
        },
        timeout: 30000,
      },
    );

    const kieResponse = response.data;
    if (kieResponse.code !== 200 || !kieResponse.data) {
      throw new Error(
        `Failed to query task: ${this.getKieResponseMessage(kieResponse)}`,
      );
    }

    const record = kieResponse.data;
    const meta = this.extractTaskMetaFromParam(record.param);
    const { videoUrl } = this.parseKieResultJson(record.resultJson);

    const data: VideoQueryResponse = {
      id: record.taskId,
      object: "video",
      model: this.normalizeVideoModel(record.model),
      status: this.mapKieStateToStatus(record.state),
      progress:
        this.mapKieStateToStatus(record.state) === "completed" ? 100 : 0,
      created_at: record.createTime,
      completed_at: record.completeTime,
      seconds: meta.seconds,
      size: meta.size,
      video_url: videoUrl,
    };

    if (data.status === "failed") {
      data.error = {
        message:
          record.failMsg ||
          (record.failCode ? `Task failed (${record.failCode})` : "Task failed"),
      };
    }

    let statusMessage = "";

    if (data.status === "completed" || data.status === "COMPLETED") {
      const isWithVideo = meta.isWithVideo || false;
      const res = meta.resolution || "480p";
      const seconds = Number(meta.seconds) || 8;
      
      let unitPricePerSecUSD = 0;
      if (res === "720p") {
        unitPricePerSecUSD = isWithVideo ? 0.100 : 0.165;
      } else {
        unitPricePerSecUSD = isWithVideo ? 0.045 : 0.0775;
      }
      
      const unitPricePerSecCNY = unitPricePerSecUSD * 7.2;
      data.amount = Number((unitPricePerSecCNY * seconds).toFixed(3));

      return {
        content: [
          {
            type: "text",
            text: JSON.stringify(
              {
                success: true,
                message: "Task completed",
                data: data,
              },
              null,
              2,
            ),
          },
        ],
      };
    }

    switch (data.status) {
      case "in_progress":
      case "IN_PROGRESS":
        statusMessage = `Processing... progress: ${data.progress || 0}%`;
        break;
      case "completed":
      case "COMPLETED":
        statusMessage = "Task completed";
        break;
      case "failed":
      case "FAILED":
        if (data.error?.message) {
          statusMessage = `Task failed: ${data.error.message}`;
        } else {
          statusMessage = "Task failed";
        }
        break;
      default:
        statusMessage = `Task status: ${data.status}`;
    }

    const responseData: any = {
      id: data.id,
      object: data.object,
      model: data.model,
      status: data.status,
      progress: data.progress,
      created_at: data.created_at,
      completed_at: data.completed_at,
      seconds: data.seconds,
      size: data.size,
    };

    if ((data.status === "failed" || data.status === "FAILED") && data.error) {
      responseData.error = data.error;
    }

    const isFailed = data.status === "failed" || data.status === "FAILED";

    return {
      content: [
        {
          type: "text",
          text: JSON.stringify(
            {
              success: !isFailed,
              message: statusMessage,
              data: responseData,
            },
            null,
            2,
          ),
        },
      ],
      isError: isFailed,
    };
  }

  private async handleGetVideoContent(args: unknown) {
    const params = getVideoContentSchema.parse(args);

    const taskResult = await this.handleQueryTask({ id: params.id });
    const taskResultText = taskResult.content[0]?.text;
    const taskResultJson = taskResultText ? JSON.parse(taskResultText) : null;
    const taskStatus = taskResultJson?.data?.status;
    const videoUrl = taskResultJson?.data?.video_url as string | undefined;

    if (!videoUrl) {
      if (taskStatus === "in_progress" || taskStatus === "IN_PROGRESS") {
        throw new Error("Task is still in progress. No video available yet.");
      }
      throw new Error("No video download URL found for this task.");
    }

    const videoResponse = await axios.get(videoUrl, {
      responseType: "arraybuffer",
      timeout: 120000,
    });

    const formData = new FormData();
    formData.append("file", Buffer.from(videoResponse.data), {
      filename: `${params.id}.mp4`,
      contentType: "video/mp4",
    });

    const uploadResponse = await axios.post(UPLOAD_URL, formData, {
      headers: {
        ...formData.getHeaders(),
      },
      timeout: 120000,
    });

    if (uploadResponse.data.code !== 0) {
      throw new Error(`Video upload failed: ${uploadResponse.data.msg}`);
    }

    return {
      content: [
        {
          type: "text",
          text: JSON.stringify(
            {
              success: true,
              message: "Video successfully uploaded to server",
              data: {
                id: params.id,
                video_url: uploadResponse.data.data.url,
                original_video_id: params.id,
              },
            },
            null,
            2,
          ),
        },
      ],
    };
  }

  async run(): Promise<void> {
    const transport = new StdioServerTransport();
    await this.server.connect(transport);
    console.error("Video MCP Server running on stdio");
  }
}

const server = new SoraVideoMCPServer();
server.run().catch(console.error);