#!/usr/bin/env python3
"""
独立图像生成工具
支持多种AI图像生成API
"""

import os
import sys
import json
import argparse
import urllib.request
from pathlib import Path


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

    def __init__(self, api_key: str, model: str):
        self.api_key = api_key
        self.model = model

    def generate(self, prompt: str, **kwargs) -> str:
        """生成图片，返回图片路径"""
        raise NotImplementedError

    def _download_image(self, url: str, output_path: str = None) -> str:
        """下载图片到本地"""
        if not output_path:
            output_path = f"generated_{self.model}_{int(os.times()[4])}.png"

        # 确保目录存在
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)

        urllib.request.urlretrieve(url, output_path)
        return output_path


class ZhipuGenerator(ImageGenerator):
    """智谱AI CogView 生成器"""

    API_URL = "https://open.bigmodel.cn/api/paas/v4/images/generations"

    def generate(self, prompt: str, size: str = "1024x1024", **kwargs) -> str:
        import urllib.error

        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"]
                    return self._download_image(image_url, kwargs.get("output_path"))
                else:
                    raise Exception(f"生成失败: {result}")
        except urllib.error.HTTPError as e:
            raise Exception(f"HTTP错误: {e.code} - {e.read().decode()}")


class OpenAIGenerator(ImageGenerator):
    """OpenAI DALL-E 生成器"""

    API_URL = "https://api.openai.com/v1/images/generations"

    def generate(self, prompt: str, size: str = "1024x1024", **kwargs) -> str:
        import urllib.error

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        data = {
            "model": self.model,
            "prompt": prompt,
            "n": 1,
            "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"]
                    return self._download_image(image_url, kwargs.get("output_path"))
                else:
                    raise Exception(f"生成失败: {result}")
        except urllib.error.HTTPError as e:
            raise Exception(f"HTTP错误: {e.code} - {e.read().decode()}")


class CustomAPIGenerator(ImageGenerator):
    """自定义API生成器"""

    def __init__(self, api_key: str, model: str, api_url: str):
        super().__init__(api_key, model)
        self.api_url = api_url

    def generate(self, prompt: str, size: str = "1024x1024", **kwargs) -> str:
        import urllib.error

        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 isinstance(result, dict):
                    data = result.get("data", [])
                else:
                    data = result

                if data:
                    # 尝试获取URL
                    if isinstance(data[0], dict):
                        image_url = data[0].get("url") or data[0].get("image_url")
                    else:
                        image_url = data[0]

                    if image_url:
                        return self._download_image(str(image_url), kwargs.get("output_path"))

                raise Exception(f"无法解析响应: {result}")
        except urllib.error.HTTPError as e:
            raise Exception(f"HTTP错误: {e.code} - {e.read().decode()}")


def get_generator(provider: str, api_key: str, model: str, api_url: str = None):
    """获取对应的生成器"""
    generators = {
        "zhipu": ZhipuGenerator,
        "openai": OpenAIGenerator,
        "custom": CustomAPIGenerator,
    }

    generator_class = generators.get(provider.lower())
    if not generator_class:
        raise ValueError(f"不支持的提供商: {provider}，可选: {list(generators.keys())}")

    if provider.lower() == "custom" and not api_url:
        raise ValueError("自定义API需要提供 --api-url 参数")

    if provider.lower() == "custom":
        return generator_class(api_key, model, api_url)
    return generator_class(api_key, model)


def main():
    parser = argparse.ArgumentParser(
        description="独立图像生成工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  %(prog)s --provider zhipu --api-key YOUR_KEY --prompt "一只猫"
  %(prog)s --provider openai --api-key YOUR_KEY --model dall-e-3 --output ./image.png
  %(prog)s --provider custom --api-url https://api.example.com/v1/images --api-key KEY --prompt "风景"

支持的提供商:
  zhipu   - 智谱AI (cogview-3, cogview-3-flash)
  openai  - OpenAI (dall-e-2, dall-e-3)
  custom  - 自定义API (需提供 --api-url)
        """
    )

    parser.add_argument("--provider", "-p", default="zhipu",
                        choices=["zhipu", "openai", "custom"],
                        help="API提供商 (默认: zhipu)")
    parser.add_argument("--api-key", "-k", required=True,
                        help="API密钥")
    parser.add_argument("--model", "-m", default="cogview-3-flash",
                        help="模型名称 (默认: cogview-3-flash)")
    parser.add_argument("--api-url", "-u",
                        help="自定义API地址 (仅provider=custom时需要)")
    parser.add_argument("--prompt", "-t", required=True,
                        help="图片描述提示词")
    parser.add_argument("--size", "-s", default="1024x1024",
                        help="图片尺寸 (默认: 1024x1024)")
    parser.add_argument("--output", "-o",
                        help="输出文件路径 (默认: auto-generated)")
    parser.add_argument("--verbose", "-v", action="store_true",
                        help="显示详细信息")

    args = parser.parse_args()

    try:
        # 获取生成器
        generator = get_generator(args.provider, args.api_key, args.model, args.api_url)

        if args.verbose:
            print(f"[配置] 提供商: {args.provider}")
            print(f"[配置] 模型: {args.model}")
            print(f"[配置] 尺寸: {args.size}")
            print(f"[配置] 提示词: {args.prompt}")

        # 生成图片
        print(f"🎨 正在生成图片...")
        output_path = generator.generate(
            prompt=args.prompt,
            size=args.size,
            output_path=args.output
        )

        print(f"✅ 图片已保存: {output_path}")
        return 0

    except Exception as e:
        print(f"❌ 错误: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())
