#!/usr/bin/env node

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ErrorCode,
  ListToolsRequestSchema,
  McpError,
} from '@modelcontextprotocol/sdk/types.js';
import { ImageGenerator } from './imageGenerator.js';

/**
 * AI image generation MCP server
 */
class AIImageGenerationMCPServer {
  private server: Server;
  private imageGenerator?: ImageGenerator;

  constructor() {
    this.server = new Server(
      {
        name: 'ai-image-generation-mcp',
        version: '2.0.0',
      },
      {
        capabilities: {
          tools: {},
        },
      }
    );

    this.setupTools();
  }

  private getImageGenerator(): ImageGenerator {
    if (!this.imageGenerator) {
      const apiKey = process.env.KIE_API_KEY;
      const model = process.env.KIE_MODEL;
      if (!apiKey) {
        throw new McpError(
          ErrorCode.InvalidRequest,
          'KIE_API_KEY was not detected. Set it in the environment and try again. Visit https://kie.ai/api-key to get an API key.'
        );
      }
      this.imageGenerator = new ImageGenerator(apiKey, model);
    }
    return this.imageGenerator;
  }

  private setupTools() {
    // Register the tool list handler.
    this.server.setRequestHandler(ListToolsRequestSchema, async () => {
      return {
        tools: [
          {
            name: 'generate-image',
            description: 'AI image generation tool (GPT Image-2). Creates an image generation task and returns a task ID. Use query-task to check progress.',
            inputSchema: {
              type: 'object',
              properties: {
                prompt: {
                  type: 'string',
                  description: 'Image generation prompt. Describe the desired image in detail (up to 20,000 characters).'
                },
                nsfw_checker: {
                  type: 'boolean',
                  description: 'Whether to disable platform-side content filtering. Defaults to false.'
                },
                callBackUrl: {
                  type: 'string',
                  description: 'Optional callback URL invoked after the task completes.'
                }
              },
              required: ['prompt']
            }
          },
          {
            name: 'query-task',
            description: 'Query the status and progress of an image generation task. Returns the current state, elapsed time, progress details, and related task data.',
            inputSchema: {
              type: 'object',
              properties: {
                taskId: {
                  type: 'string',
                  description: 'Task ID returned by generate-image.'
                }
              },
              required: ['taskId']
            }
          }
        ]
      };
    });

    // Register the tool call handler.
    this.server.setRequestHandler(CallToolRequestSchema, async (request: any) => {
      const { name, arguments: rawArgs } = request.params;

      // Normalize arguments.
      let args: any;
      if (typeof rawArgs === 'string') {
        try {
          args = JSON.parse(rawArgs);
        } catch (error) {
          throw new McpError(
            ErrorCode.InvalidParams,
            'Invalid argument format: unable to parse JSON arguments.'
          );
        }
      } else {
        args = rawArgs;
      }

      try {
        switch (name) {
          case 'generate-image':
            return await this.handleGenerateImage(args);

          case 'query-task':
            return await this.handleQueryTask(args);

          default:
            throw new McpError(
              ErrorCode.MethodNotFound,
              `Unknown tool: ${name}`
            );
        }
      } catch (error) {
        console.error(`Tool '${name}' execution error:`, error);
        throw new McpError(
          ErrorCode.InternalError,
          `Tool execution failed: ${error instanceof Error ? error.message : String(error)}`
        );
      }
    });
  }

  /**
   * Handle image generation tasks in hybrid mode: return the result synchronously within 2 seconds, otherwise return the taskId.
   */
  private async handleGenerateImage(args: any) {
    const { 
      prompt, 
      nsfw_checker = false,
      callBackUrl
    } = args;

    if (!prompt) {
      throw new McpError(
        ErrorCode.InvalidParams,
        'Missing required parameter: prompt'
      );
    }

    try {
      console.error(`Starting image generation task...`);
      
      // Call the timeout-aware generation method.
      const result = await this.getImageGenerator().generateImageWithTimeout(
        {
          prompt,
          nsfw_checker,
          callBackUrl
        },
        (message: string, elapsedSeconds: number) => {
          console.error(`[${elapsedSeconds}s] ${message}`);
        },
        2 // 2-second timeout.
      );

      if (result.success) {
        // Completed within 2 seconds, return the result directly.
        console.error(`Image generation completed; returning the result directly.`);
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                status: 'success',
                message: 'Image generation completed.',
                prompt: prompt,
                nsfw_checker: nsfw_checker,
                imageUrl: result.imageUrl
              }, null, 2)
            }
          ],
          isError: false
        };
      } else {
        // Timed out, return the taskId.
        console.error(`Image generation timed out; returning taskId: ${result.taskId}`);
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                status: 'async',
                message: 'Image generation is taking longer than expected and has switched to async mode.',
                taskId: result.taskId,
                prompt: prompt,
                nsfw_checker: nsfw_checker,
                tip: 'Use the query-task tool to check task progress and results.'
              }, null, 2)
            }
          ],
          isError: false
        };
      }
    } catch (error) {
      console.error(`Image generation failed:`, error);
      throw new McpError(
        ErrorCode.InternalError,
        `Image generation failed: ${error instanceof Error ? error.message : String(error)}`
      );
    }
  }

  /**
   * Query task status.
   */
  private async handleQueryTask(args: any) {
    const { taskId } = args;

    if (!taskId) {
      throw new McpError(
        ErrorCode.InvalidParams,
        'Missing required parameter: taskId'
      );
    }

    try {
      // Query task status.
      const taskData = await this.getImageGenerator().queryTask(taskId);

      // Parse the result.
      let imageUrl: string | undefined;
      if (taskData.state === 'success' && taskData.resultJson) {
        try {
          const resultJson = JSON.parse(taskData.resultJson);
          if (resultJson.resultUrls && resultJson.resultUrls.length > 0) {
            imageUrl = resultJson.resultUrls[0];
          }
        } catch (e) {
          console.error('Failed to parse resultJson:', e);
        }
      }

      // Calculate elapsed time.
      const elapsedSeconds = taskData.costTime 
        ? Math.floor(taskData.costTime / 1000) 
        : Math.floor((Date.now() - taskData.createTime) / 1000);

      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify({
              taskId: taskData.taskId,
              state: taskData.state,
              model: taskData.model,
              elapsedSeconds: elapsedSeconds,
              imageUrl: imageUrl,
              failCode: taskData.failCode,
              failMsg: taskData.failMsg,
              createdAt: new Date(taskData.createTime).toISOString(),
              completedAt: taskData.completeTime ? new Date(taskData.completeTime).toISOString() : undefined
            }, null, 2)
          }
        ],
        isError: false
      };
    } catch (error) {
      console.error(`Task query failed:`, error);
      throw new McpError(
        ErrorCode.InternalError,
        `Task query failed: ${error instanceof Error ? error.message : String(error)}`
      );
    }
  }

  /**
   * Start the server.
   */
  async start() {
    const transport = new StdioServerTransport();
    await this.server.connect(transport);
    console.error('AI image generation MCP server started (GPT Image-2 API)');
  }
}

// Start the server.
async function main() {
  try {
    const server = new AIImageGenerationMCPServer();
    await server.start();
  } catch (error) {
    console.error('Server startup failed:', error);
    process.exit(1);
  }
}

main().catch((error) => {
  console.error('Server runtime error:', error);
  process.exit(1);
});