#!/usr/bin/env python3
from __future__ import annotations

import base64
import os
from typing import Any, Dict, List, Optional

import httpx
from mcp.server.fastmcp import FastMCP


server = FastMCP("yooztech_mcp_minimax_t2i")

API_BASE_URL = "https://api.minimaxi.com/v1/image_generation"


def _get_api_key() -> str:
    """从环境变量获取 API Key"""
    api_key = os.environ.get("MINIMAX_API_KEY", "").strip()
    if not api_key:
        raise ValueError(
            "未设置 MINIMAX_API_KEY 环境变量。\n"
            "请在 MCP Server 配置中设置环境变量：\n"
            '"env": {"MINIMAX_API_KEY": "your-api-key"}'
        )
    return api_key


@server.tool()
async def text_to_image(
    prompt: str,
    model: str = "image-01",
    aspect_ratio: str = "16:9",
    response_format: str = "base64",
) -> Dict[str, Any]:
    """根据文本描述生成图片（文生图）

    必填参数：
    - prompt: 图片描述文本，建议详细描述场景、主体、风格等

    可选参数：
    - model: 使用的模型，默认 "image-01"
    - aspect_ratio: 图片宽高比，默认 "16:9"，可选值：1:1, 16:9, 9:16, 4:3, 3:4
    - response_format: 返回格式，默认 "base64"，可选值：base64, url

    返回：
    - success: 是否成功
    - images: 生成的图片列表（base64 或 url）
    - format: 返回格式
    - model: 使用的模型
    - aspect_ratio: 宽高比

    使用示例：
    text_to_image(prompt="men Dressing in white t shirt, full-body stand front view image :25, outdoor, Venice beach sign")
    """
    if not prompt or not isinstance(prompt, str) or not prompt.strip():
        raise ValueError("prompt 是必填项，必须提供图片描述文本")

    api_key = _get_api_key()
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    payload = {
        "model": model,
        "prompt": prompt,
        "aspect_ratio": aspect_ratio,
        "response_format": response_format,
    }

    try:
        timeout = httpx.Timeout(120.0, connect=30.0, read=120.0, write=30.0)
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(API_BASE_URL, headers=headers, json=payload)
            resp.raise_for_status()

        data = resp.json()

        if "data" in data and "image_base64" in data["data"]:
            images = data["data"]["image_base64"]
        elif "data" in data and "image_url" in data["data"]:
            images = data["data"]["image_url"]
        else:
            images = []

        return {
            "success": True,
            "images": images,
            "format": response_format,
            "model": model,
            "aspect_ratio": aspect_ratio,
            "response": data,
        }

    except httpx.HTTPStatusError as e:
        return {
            "success": False,
            "error": f"HTTP 错误: {e.response.status_code}",
            "message": e.response.text,
        }
    except Exception as e:
        return {
            "success": False,
            "error": "请求失败",
            "message": str(e),
        }


@server.tool()
async def image_to_image(
    prompt: str,
    subject_reference: List[Dict[str, Any]],
    model: str = "image-01",
    aspect_ratio: str = "16:9",
    response_format: str = "base64",
) -> Dict[str, Any]:
    """结合参考图生成图片（图生图）

    通过提供参考图片（包含清晰主体的图片），结合文本描述生成保留主体特征的新图片。
    适用于保持人物形象一致性的场景，例如为同一个虚拟角色生成不同情境下的图片。

    必填参数：
    - prompt: 图片描述文本，描述希望生成的场景
    - subject_reference: 参考图列表，每项包含：
      - type: 参考图类型，目前仅支持 "character"（角色）
      - image_file: 参考图 URL（支持网络图片链接）

    可选参数：
    - model: 使用的模型，默认 "image-01"
    - aspect_ratio: 图片宽高比，默认 "16:9"，可选值：1:1, 16:9, 9:16, 4:3, 3:4
    - response_format: 返回格式，默认 "base64"，可选值：base64, url

    返回：
    - success: 是否成功
    - images: 生成的图片列表（base64 或 url）
    - format: 返回格式
    - model: 使用的模型
    - aspect_ratio: 宽高比

    使用示例：
    image_to_image(
        prompt="女孩在图书馆的窗户前，看向远方",
        subject_reference=[
            {"type": "character", "image_file": "https://example.com/ref.jpg"}
        ]
    )
    """
    if not prompt or not isinstance(prompt, str) or not prompt.strip():
        raise ValueError("prompt 是必填项，必须提供图片描述文本")

    if not subject_reference or not isinstance(subject_reference, list):
        raise ValueError("subject_reference 是必填项，必须提供参考图列表")

    for idx, ref in enumerate(subject_reference):
        if not isinstance(ref, dict):
            raise ValueError(f"subject_reference[{idx}] 必须是对象")
        if "type" not in ref or "image_file" not in ref:
            raise ValueError(f"subject_reference[{idx}] 必须包含 type 和 image_file 字段")

    api_key = _get_api_key()
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    payload = {
        "model": model,
        "prompt": prompt,
        "aspect_ratio": aspect_ratio,
        "subject_reference": subject_reference,
        "response_format": response_format,
    }

    try:
        timeout = httpx.Timeout(120.0, connect=30.0, read=120.0, write=30.0)
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(API_BASE_URL, headers=headers, json=payload)
            resp.raise_for_status()

        data = resp.json()

        if "data" in data and "image_base64" in data["data"]:
            images = data["data"]["image_base64"]
        elif "data" in data and "image_url" in data["data"]:
            images = data["data"]["image_url"]
        else:
            images = []

        return {
            "success": True,
            "images": images,
            "format": response_format,
            "model": model,
            "aspect_ratio": aspect_ratio,
            "response": data,
        }

    except httpx.HTTPStatusError as e:
        return {
            "success": False,
            "error": f"HTTP 错误: {e.response.status_code}",
            "message": e.response.text,
        }
    except Exception as e:
        return {
            "success": False,
            "error": "请求失败",
            "message": str(e),
        }


@server.tool()
async def save_base64_images(
    images: List[str],
    output_dir: str = ".",
    filename_prefix: str = "output",
) -> Dict[str, Any]:
    """将 Base64 编码的图片保存为文件

    必填参数：
    - images: Base64 编码的图片列表

    可选参数：
    - output_dir: 输出目录，默认当前目录
    - filename_prefix: 文件名前缀，默认 "output"

    返回：
    - success: 是否成功
    - saved_files: 保存的文件路径列表

    注意：此工具仅用于保存 base64 格式的图片数据
    """
    if not images or not isinstance(images, list):
        raise ValueError("images 是必填项，必须提供 Base64 编码的图片列表")

    import os

    output_path = os.path.expanduser(output_dir)
    os.makedirs(output_path, exist_ok=True)

    saved_files: List[str] = []
    errors: List[str] = []

    for idx, img_data in enumerate(images):
        try:
            filename = f"{filename_prefix}-{idx}.jpeg"
            filepath = os.path.join(output_path, filename)

            img_bytes = base64.b64decode(img_data)
            with open(filepath, "wb") as f:
                f.write(img_bytes)

            saved_files.append(filepath)
        except Exception as e:
            errors.append(f"第 {idx} 张图片保存失败: {str(e)}")

    return {
        "success": len(errors) == 0,
        "saved_files": saved_files,
        "errors": errors if errors else None,
    }


def main() -> None:
    server.run()


if __name__ == "__main__":
    main()
