import { McpServer} from "@modelcontextprotocol/sdk/server/mcp.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { z } from "zod";
import axios from 'axios';
import * as fs from 'fs';
import * as path from 'path';

// 视频生成模型配置
const MODEL = 'wanx2.1-t2v-turbo';

// 设置API密钥和API端点
const API_KEY = process.env.DASHSCOPE_API_KEY || ''; // 从环境变量读取API密钥
const VIDEO_GENERATION_API = 'https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation/video-synthesis';
const TASK_STATUS_API = 'https://dashscope.aliyuncs.com/api/v1/tasks/';

// 检查API密钥是否存在
if (!API_KEY) {
  console.warn('警告: DASHSCOPE_API_KEY环境变量未设置，视频生成功能将无法使用');
}

// API响应类型定义
interface TaskResponse {
  task_id?: string;
  output?: {
    task_id?: string;
    task_status?: string;
    video_url?: string;
    result?: {
      video_url?: string;
    };
    actual_prompt?: string;
  };
  task_status?: string;
  status?: string;
  message?: string;
}

// 创建MCP服务器
const server = new McpServer({
  name: "文生视频MCP服务器",
  version: "1.0.0",
  description: "基于阿里云的文本生成视频工具，使用环境变量 DASHSCOPE_API_KEY 设置API密钥"
});

// 添加文生视频工具
server.tool("gen_video",
  { 
    prompt: z.string().min(1, "提示词不能为空"),
    size: z.string().regex(/^\d+\*\d+$/, "尺寸格式应为宽*高，如1280*720").optional().default("1280*720"),
    duration: z.number().min(1, "时长不能小于1秒").max(60, "时长不能超过60秒").optional().default(5) 
  },
  async ({ prompt, size, duration }: { prompt: string; size: string; duration: number }) => {
    try {
      // 验证API密钥
      if (!API_KEY || API_KEY.trim() === '') {
        return {
          content: [{ type: "text", text: "错误: 缺少API密钥，请设置DASHSCOPE_API_KEY环境变量" }]
        };
      }
      
      // 发送视频生成请求
      const response = await axios.post(
        VIDEO_GENERATION_API,
        {
          model: MODEL,
          input: {
            prompt,
          },
          parameters: {
            size,
            duration,
            prompt_extend: true, // 启用提示词扩展
          },
        },
        {
          headers: {
            'Authorization': `Bearer ${API_KEY}`,
            'Content-Type': 'application/json',
            'X-DashScope-Async': 'enable',
          },
        }
      );
      
      const responseData = response.data as TaskResponse;
      const taskId = responseData.task_id || responseData.output?.task_id;
      
      if (!taskId) {
        throw new Error('无法从响应中获取任务ID');
      }
      
      // 每5秒检查一次任务状态，最多检查72次(最大6分钟)
      const maxAttempts = 72;
      const intervalMs = 5000;
      let attempts = 0;
      let result: TaskResponse | null = null;
      let errorCount = 0;
      const startTime = Date.now();
      
      while (attempts < maxAttempts) {
        await new Promise(resolve => setTimeout(resolve, intervalMs));
        attempts++;
        
        try {
          // 查询任务状态
          const statusResponse = await axios.get(
            `${TASK_STATUS_API}${taskId}`,
            {
              headers: {
                'Authorization': `Bearer ${API_KEY}`,
                'Content-Type': 'application/json',
              },
            }
          );
          
          result = statusResponse.data;
          const status = result?.task_status || result?.output?.task_status;
          const waitedSeconds = Math.floor((Date.now() - startTime) / 1000);
          
          if (status === 'SUCCEEDED') {
            break;
          } else if (status === 'FAILED') {
            throw new Error(`任务失败: ${result?.message || '未知错误'}`);
          } else if (status === 'PENDING' || status === 'RUNNING') {
            // 任务还在运行中，继续等待
            // 重置错误计数，因为我们成功获取了状态
            errorCount = 0;
          } else if (!status) {
            // 无法获取状态，但继续尝试
          }
          
        } catch (error: any) {
          errorCount++;
          const waitedSeconds = Math.floor((Date.now() - startTime) / 1000);
          
          // 连续多次查询失败，抛出异常
          if (errorCount >= 3) {
            throw new Error(`查询任务状态连续失败${errorCount}次: ${error.message || '未知错误'}`);
          }
        }
      }
      
      if (!result) {
        throw new Error('无法获取任务结果');
      }
      
      // 检查是否达到最大尝试次数
      if (attempts >= maxAttempts) {
        const waitedSeconds = Math.floor((Date.now() - startTime) / 1000);
        return {
          content: [
            { 
              type: "text", 
              text: `视频生成任务仍在处理中 • 已等待 ${waitedSeconds} 秒\n` +
                    `您可以稍后使用任务ID查询结果\n` +
                    `任务ID: ${taskId}` 
            }
          ]
        };
      }
      
      // 获取视频URL
      let videoUrl = result?.output?.video_url || result?.output?.result?.video_url;
      if (!videoUrl) {
        if (result?.status === 'SUCCEEDED' || result?.output?.task_status === 'SUCCEEDED') {
          return {
            content: [
              { 
                type: "text", 
                text: `视频生成成功，但无法获取视频URL，请查看完整响应数据\n` +
                       `任务ID: ${taskId}\n` +
                       `完整响应数据:\n${JSON.stringify(result, null, 2)}` 
              }
            ]
          };
        }
        throw new Error('无法获取视频URL');
      }
      
      // 获取扩展提示词(如果有)
      const actualPrompt = result.output?.actual_prompt;
      const waitedSeconds = Math.floor((Date.now() - startTime) / 1000);
      
      // 构建响应内容
      const content: any[] = [
        { 
          type: "text", 
          text: `✅ 视频生成成功! (用时 ${waitedSeconds} 秒)\n` +
                 `任务ID: ${taskId}\n` +
                 `参数信息:\n` +
                 `- 提示词: ${prompt}\n` +
                 `- 尺寸: ${size}\n` +
                 `- 时长: ${duration}秒\n`
        },
        // 单独添加一个只包含视频URL的文本框，方便点击
        {
          type: "text",
          text: `${videoUrl}`
        }
      ];
      
      // 如果有扩展提示词，添加到响应
      if (actualPrompt) {
        content.push({ 
          type: "text", 
          text: `\n系统扩展生成的完整提示词:\n${actualPrompt}` 
        });
      }
      
      return { content };
    } catch (error: any) {
      console.error('视频生成请求失败:');
      let errorMessage = '';
      
      if (axios.isAxiosError(error) && error.response) {
        console.error('API错误响应:', error.response.data);
        errorMessage = error.response.data.message || '未知错误';
        
        // 展示完整错误数据
        return {
          content: [{ 
            type: "text", 
            text: `生成视频失败: ${errorMessage}\n\n` +
                  `API错误详情:\n${JSON.stringify(error.response.data, null, 2)}`
          }]
        };
      }
      
      errorMessage = error.message || '未知错误';
      return {
        content: [{ type: "text", text: `生成视频失败: ${errorMessage}` }]
      };
    }
  }
);

// 添加查询任务状态工具
server.tool("check_task",
  { 
    taskId: z.string().min(1, "任务ID不能为空")
  },
  async ({ taskId }: { taskId: string }) => {
    try {
      // 验证API密钥
      if (!API_KEY || API_KEY.trim() === '') {
        return {
          content: [{ type: "text", text: "错误: 缺少API密钥，请设置DASHSCOPE_API_KEY环境变量" }]
        };
      }
      
      // 查询任务状态
      const response = await axios.get(
        `${TASK_STATUS_API}${taskId}`,
        {
          headers: {
            'Authorization': `Bearer ${API_KEY}`,
            'Content-Type': 'application/json',
          },
        }
      );
      
      const result = response.data as TaskResponse;
      const status = result.task_status || result.output?.task_status;
      
      // 获取视频URL
      let videoUrl = result.output?.video_url || result.output?.result?.video_url;
      
      // 获取扩展提示词(如果有)
      const actualPrompt = result.output?.actual_prompt;
      
      let content: any[] = [];
      
      // 基于状态构建不同的响应
      if (status === 'SUCCEEDED' && videoUrl) {
        content.push({ 
          type: "text", 
          text: `✅ 视频生成成功!\n` +
                 `任务ID: ${taskId}\n`
        });
        
        // 单独添加一个只包含视频URL的文本框，方便点击
        content.push({
          type: "text",
          text: `${videoUrl}`
        });
        
        if (actualPrompt) {
          content.push({ type: "text", text: `\n系统扩展生成的完整提示词:\n${actualPrompt}` });
        }
      } else if (status === 'SUCCEEDED' && !videoUrl) {
        content.push({
          type: "text",
          text: `任务已完成，但无法获取视频URL\n完整响应数据:\n${JSON.stringify(result, null, 2)}`
        });
      } else if (status === 'FAILED') {
        content.push({
          type: "text",
          text: `❌ 任务失败\n失败原因: ${result.message || '未知错误'}`
        });
      } else if (status === 'PENDING' || status === 'RUNNING') {
        content.push({
          type: "text",
          text: `⏳ 任务正在运行中\n状态: ${status}\n任务ID: ${taskId}`
        });
      } else {
        content.push({
          type: "text",
          text: `⚠️ 无法识别任务状态\n状态值: ${status || '无状态'}\n完整响应数据:\n${JSON.stringify(result, null, 2)}`
        });
      }
      
      return { content };
    } catch (error: any) {
      // 处理错误
      const errorMessage = error.response?.data?.message || error.message || '未知错误';
      
      if (error.response?.status === 404) {
        return {
          content: [{ type: "text", text: `错误: 找不到指定的任务ID (${taskId})` }]
        };
      }
      
      return {
        content: [{ type: "text", text: `查询任务状态失败: ${errorMessage}` }]
      };
    }
  }
);

// 初始化stdin/stdout传输
const transport = new StdioServerTransport();

// 启动服务器的异步函数
export async function startServer() {
  console.log("文生视频MCP服务器正在启动...");
  await server.connect(transport);
  console.log("文生视频MCP服务器已启动，等待连接...");
}

// 如果直接运行此文件，则启动服务器
if (require.main === module) {
  startServer().catch(error => {
    console.error("服务器启动错误:", error);
    process.exit(1);
  });
}