import * as readline from 'readline';
import axios from 'axios';

// 定义 MCP 命令类型
type MCPCommand = {
  execute: (params: any) => Promise<any>;
  description: string;
  parameters: Record<string, any>;
};

// 定义 MCP 服务器类
class MCPServer {
  private commands: Record<string, MCPCommand> = {};
  private stdio: readline.Interface;
  private apiKey: string | undefined;

  constructor() {
    // 初始化 stdio 接口用于与 Claude 通信
    this.stdio = readline.createInterface({
      input: process.stdin,
      output: process.stdout,
    });
    
    // 从环境变量获取 API Key
    this.apiKey = process.env.API_KEY;

    // 注册生成图片命令
    this.registerCommand('generateImage', {
      execute: this.generateImage.bind(this),
      description: '使用快手API生成AI图片',
      parameters: {
        prompt: {
          type: 'string',
          description: '图片生成提示词',
          required: true
        },
        size: {
          type: 'string',
          description: '图片尺寸',
          enum: ['1024x1024', '1024x576', '576x1024', '720x1280', '1280x720'],
          default: '1024x1024'
        },
        count: {
          type: 'number',
          description: '生成图片数量',
          minimum: 1,
          maximum: 4,
          default: 1
        }
      }
    });

    // 开始监听命令
    this.start();
  }

  // 注册命令
  private registerCommand(name: string, command: MCPCommand) {
    this.commands[name] = command;
  }

  // 生成图片的核心功能
  private async generateImage(params: any): Promise<any> {
    if (!this.apiKey) {
      throw new Error('未配置API密钥，请在 MCP 配置中设置 API_KEY 环境变量');
    }

    try {
      console.error(`开始生成图片, 提示词: ${params.prompt}`);
      
      const response = await axios.post('https://open.kuaishou.com/openapi/v1/ai/image/generate', {
        prompt: params.prompt,
        size: params.size || '1024x1024',
        count: params.count || 1
      }, {
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${this.apiKey}`
        }
      });

      const result = {
        success: true,
        data: {
          images: response.data.images || [],
          seed: response.data.seed || Math.floor(Math.random() * 1000000),
          time_taken: response.data.time_taken || 'unknown'
        }
      };

      console.error(`图片生成成功，返回 ${result.data.images.length} 张图片`);
      return result;
    } catch (error: any) {
      console.error(`图片生成失败: ${error.message}`);
      
      if (error.response) {
        // 处理API错误响应
        switch (error.response.status) {
          case 400:
            throw new Error(`请求错误: ${error.response.data.message}`);
          case 401:
            throw new Error('无效的API密钥');
          case 404:
            throw new Error('API端点不存在');
          case 429:
            throw new Error('请求频率超限');
          case 503:
            throw new Error('服务暂时不可用');
          case 504:
            throw new Error('请求超时');
          default:
            throw new Error(`未知错误: ${error.response.status}`);
        }
      } else if (error.request) {
        throw new Error('请求发送失败，网络问题');
      } else {
        throw new Error(`请求配置错误: ${error.message}`);
      }
    }
  }

  // 处理接收到的 JSON-RPC 2.0 消息
  private async handleMessage(message: any) {
    try {
      // 检查是否是有效的 JSON-RPC 2.0 请求
      if (message.jsonrpc !== '2.0' || !message.method) {
        return this.sendError(message.id, -32600, '无效的请求');
      }

      // 处理不同类型的请求
      if (message.method === 'initialize') {
        // 初始化请求，返回支持的命令
        const commandsInfo = Object.keys(this.commands).map(name => ({
          name,
          description: this.commands[name].description,
          parameters: this.commands[name].parameters
        }));

        return this.sendResult(message.id, {
          name: 'kuaishou-image-generator',
          version: '1.0.0',
          commands: commandsInfo
        });
      } else if (message.method === 'executeCommand') {
        // 执行命令请求
        const { command, params } = message.params || {};

        if (!command) {
          return this.sendError(message.id, -32602, '缺少命令名称');
        }

        if (!this.commands[command]) {
          return this.sendError(message.id, -32601, `未找到命令: ${command}`);
        }

        try {
          const result = await this.commands[command].execute(params || {});
          return this.sendResult(message.id, result);
        } catch (error: any) {
          return this.sendError(message.id, -32603, `执行命令出错: ${error.message}`);
        }
      } else {
        // 未知方法
        return this.sendError(message.id, -32601, `未知方法: ${message.method}`);
      }
    } catch (error: any) {
      return this.sendError(message.id, -32700, `解析错误: ${error.message}`);
    }
  }

  // 发送成功结果
  private sendResult(id: any, result: any) {
    const response = {
      jsonrpc: '2.0',
      id,
      result
    };
    console.log(JSON.stringify(response));
  }

  // 发送错误响应
  private sendError(id: any, code: number, message: string) {
    const response = {
      jsonrpc: '2.0',
      id,
      error: {
        code,
        message
      }
    };
    console.log(JSON.stringify(response));
  }

  // 启动服务器，监听消息
  private start() {
    this.stdio.on('line', async (line) => {
      try {
        const message = JSON.parse(line);
        await this.handleMessage(message);
      } catch (error: any) {
        console.error(`处理消息出错: ${error.message}`);
        this.sendError(null, -32700, `无效的JSON: ${error.message}`);
      }
    });

    // 在启动时打印一条调试信息，但不发送给客户端
    console.error('快手文生图 MCP 服务已启动');
    console.error('支持的命令:', Object.keys(this.commands).join(', '));
    
    if (!this.apiKey) {
      console.error('警告: 未设置 API_KEY 环境变量，无法调用快手 API');
    }
  }
}

// 创建并启动 MCP 服务器
new MCPServer();