#!/usr/bin/env python3
"""
图像生成 MCP 服务器
支持智谱AI、OpenAI、自定义API的图像生成
"""

import os
import json
import urllib.request
from pathlib import Path
from typing import Optional

# MCP SDK (如果使用 FastMCP)
try:
    from mcp.server.fastmcp import FastMCP
    mcp = FastMCP("image-generator")
    HAS_MCP = True
except ImportError:
    HAS_MCP = False
    print("警告: 未安装 mcp 包，请运行: pip install mcp")


class ImageGenerator:
    """图像生成器"""

    def __init__(self, api_key: str, model: str, api_url: Optional[str] = None):
        self.api_key = api_key
        self.model = model
        self.api_url = api_url or self._get_default_url()

    def _get_default_url(self) -> str:
        """根据模型获取默认 API URL"""
        if "cogview" in self.model.lower():
            return "https://open.bigmodel.cn/api/paas/v4/images/generations"
        elif "dall-e" in self.model.lower():
            return "https://api.openai.com/v1/images/generations"
        else:
            raise ValueError(f"未知模型: {self.model}，请提供 api_url")

    def generate(self, prompt: str, size: str = "1024x1024", output_path: Optional[str] = None) -> dict:
        """生成图片"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        data = {
            "model": self.model,
            "prompt": prompt,
            "size": size
        }

        req = urllib.request.Request(
            self.api_url,
            data=json.dumps(data).encode("utf-8"),
            headers=headers
        )

        try:
            with urllib.request.urlopen(req) as response:
                result = json.loads(response.read().decode("utf-8"))

                if result.get("data"):
                    image_url = result["data"][0]["url"]
                    local_path = self._download_image(image_url, output_path, prompt)
                    return {
                        "success": True,
                        "image_url": image_url,
                        "local_path": local_path,
                        "size": size
                    }
                else:
                    return {
                        "success": False,
                        "error": f"生成失败: {result}"
                    }
        except urllib.error.HTTPError as e:
            return {
                "success": False,
                "error": f"HTTP错误: {e.code} - {e.read().decode()}"
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }

    def _download_image(self, url: str, output_path: Optional[str], prompt: str) -> str:
        """下载图片"""
        if not output_path:
            # 生成安全的文件名
            safe_prompt = "".join(c for c in prompt[:20] if c.isalnum() or c in (' ', '-', '_'))
            safe_prompt = safe_prompt.replace(' ', '_')
            output_path = f"generated_{safe_prompt}.png"

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        urllib.request.urlretrieve(url, output_path)
        return str(Path(output_path).absolute())


# 全局配置
DEFAULT_API_KEY = os.environ.get("IMAGE_GEN_API_KEY", "")
DEFAULT_MODEL = os.environ.get("IMAGE_GEN_MODEL", "cogview-3-flash")
DEFAULT_API_URL = os.environ.get("IMAGE_GEN_API_URL", "")


if HAS_MCP:
    @mcp.tool()
    def generate_image(
        prompt: str,
        api_key: Optional[str] = None,
        model: Optional[str] = None,
        size: str = "1024x1024",
        output_path: Optional[str] = None,
        api_url: Optional[str] = None
    ) -> str:
        """
        生成 AI 图片

        Args:
            prompt: 图片描述提示词
            api_key: API 密钥 (默认使用环境变量 IMAGE_GEN_API_KEY)
            model: 模型名称 (默认: cogview-3-flash)
            size: 图片尺寸 (默认: 1024x1024)
            output_path: 输出文件路径 (默认: 自动生成)
            api_url: API 地址 (默认根据模型自动选择)

        Returns:
            JSON 格式的结果，包含 success、image_url、local_path 或 error

        Examples:
            >>> generate_image("一只橘猫在窗台上")
            >>> generate_image("赛博朋克城市", model="cogview-3", size="768x1344")
        """
        api_key = api_key or DEFAULT_API_KEY
        model = model or DEFAULT_MODEL
        api_url = api_url or DEFAULT_API_URL or None

        if not api_key:
            return json.dumps({
                "success": False,
                "error": "请提供 api_key 或设置环境变量 IMAGE_GEN_API_KEY"
            })

        gen = ImageGenerator(api_key, model, api_url)
        result = gen.generate(prompt, size, output_path)

        return json.dumps(result, ensure_ascii=False, indent=2)


    @mcp.tool()
    def list_models() -> str:
        """
        列出支持的图像生成模型

        Returns:
            支持的模型列表和说明
        """
        models = {
            "智谱AI": {
                "models": ["cogview-3-flash", "cogview-3"],
                "api_url": "https://open.bigmodel.cn/api/paas/v4/images/generations",
                "sizes": ["1024x1024", "768x1344", "864x1152", "1152x864", "1344x768"]
            },
            "OpenAI": {
                "models": ["dall-e-2", "dall-e-3"],
                "api_url": "https://api.openai.com/v1/images/generations",
                "sizes": ["1024x1024", "1792x1024", "1024x1792"]
            }
        }
        return json.dumps(models, ensure_ascii=False, indent=2)


    @mcp.resource("image://config")
    def get_config() -> str:
        """获取当前配置"""
        return json.dumps({
            "model": DEFAULT_MODEL,
            "has_api_key": bool(DEFAULT_API_KEY),
            "custom_api_url": bool(DEFAULT_API_URL)
        }, ensure_ascii=False, indent=2)


def main():
    """启动 MCP 服务器"""
    if not HAS_MCP:
        print("请先安装 MCP SDK:")
        print("  pip install mcp")
        return

    print("🚀 图像生成 MCP 服务器启动中...")
    print(f"📝 默认模型: {DEFAULT_MODEL}")
    print(f"🔑 API Key: {'已设置' if DEFAULT_API_KEY else '未设置'}")
    print()

    mcp.run()


if __name__ == "__main__":
    main()
