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

import os
from typing import Any, Dict, List

import httpx
from mcp.server.fastmcp import FastMCP


server = FastMCP("yooztech_mcp_minimax_t2i")

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

# 文档：response_format=url 时返回 data.image_urls；下载图片需带 Bearer 鉴权


async def _download_images_with_auth(
    image_urls: List[str],
    output_dir: str,
    filename_prefix: str,
    api_key: str,
) -> tuple[List[str], List[str]]:
    """使用 Bearer 鉴权下载图片到本地，返回 (saved_files, errors)。"""
    out_path = os.path.expanduser(output_dir)
    os.makedirs(out_path, exist_ok=True)
    headers = {"Authorization": f"Bearer {api_key}"}
    saved_files: List[str] = []
    errors: List[str] = []
    timeout = httpx.Timeout(60.0, connect=15.0, read=60.0)
    async with httpx.AsyncClient(timeout=timeout) as client:
        for idx, url in enumerate(image_urls):
            try:
                r = await client.get(url, headers=headers)
                r.raise_for_status()
                path = os.path.join(out_path, f"{filename_prefix}-{idx}.jpeg")
                with open(path, "wb") as f:
                    f.write(r.content)
                saved_files.append(path)
            except Exception as e:
                errors.append(f"第 {idx} 张: {str(e)}")
    return saved_files, errors


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",
    output_dir: str = ".",
    filename_prefix: str = "output",
) -> Dict[str, Any]:
    """根据文本描述生成图片（文生图），并在同一次调用内用 API Key 鉴权下载并保存到本地。

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

    可选参数：
    - model: 使用的模型，默认 "image-01"
    - aspect_ratio: 图片宽高比，默认 "16:9"，可选值：1:1, 16:9, 4:3, 3:2, 2:3, 3:4, 9:16, 21:9（21:9 仅 image-01）
    - output_dir: 保存目录，默认当前目录
    - filename_prefix: 本地文件名前缀，默认 "output"，保存为 {prefix}-0.jpeg ...

    返回：
    - success: 生成与下载是否全部成功
    - saved_files: 已保存的本地文件路径列表
    - images: 接口返回的图片 URL（MiniMax URL 需鉴权，请使用 saved_files）
    - model / aspect_ratio / status_code / status_msg / metadata 等同原接口

    使用示例：
    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": "url",
    }

    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()

        # 文档：url 时返回 data.image_urls
        images: List[str] = []
        if "data" in data and "image_urls" in data["data"]:
            images = data["data"]["image_urls"]

        base_resp = data.get("base_resp") or {}
        status_code = base_resp.get("status_code", 0)
        status_msg = base_resp.get("status_msg", "")

        ok_gen = status_code == 0 and bool(images)
        saved_files: List[str] = []
        download_errors: List[str] = []
        if ok_gen:
            saved_files, download_errors = await _download_images_with_auth(
                images, output_dir, filename_prefix, api_key
            )

        return {
            "success": ok_gen and len(download_errors) == 0,
            "saved_files": saved_files,
            "download_errors": download_errors if download_errors else None,
            "images": images,
            "model": model,
            "aspect_ratio": aspect_ratio,
            "status_code": status_code,
            "status_msg": status_msg,
            "metadata": data.get("metadata"),
            "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",
    output_dir: str = ".",
    filename_prefix: str = "output",
) -> Dict[str, Any]:
    """结合参考图生成图片（图生图），并在同一次调用内用 API Key 鉴权下载并保存到本地。

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

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

    可选参数：
    - model: 使用的模型，默认 "image-01"
    - aspect_ratio: 图片宽高比，默认 "16:9"，可选值：1:1, 16:9, 4:3, 3:2, 2:3, 3:4, 9:16, 21:9
    - output_dir / filename_prefix: 本地保存目录与文件名前缀（同 text_to_image）

    返回：
    - success: 生成与下载是否全部成功
    - saved_files: 已保存的本地文件路径列表
    - images: 接口返回的图片 URL

    使用示例：
    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": "url",
    }

    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()

        # 文档：url 时返回 data.image_urls
        images = []
        if "data" in data and "image_urls" in data["data"]:
            images = data["data"]["image_urls"]

        base_resp = data.get("base_resp") or {}
        status_code = base_resp.get("status_code", 0)
        status_msg = base_resp.get("status_msg", "")

        ok_gen = status_code == 0 and bool(images)
        saved_files: List[str] = []
        download_errors: List[str] = []
        if ok_gen:
            saved_files, download_errors = await _download_images_with_auth(
                images, output_dir, filename_prefix, api_key
            )

        return {
            "success": ok_gen and len(download_errors) == 0,
            "saved_files": saved_files,
            "download_errors": download_errors if download_errors else None,
            "images": images,
            "model": model,
            "aspect_ratio": aspect_ratio,
            "status_code": status_code,
            "status_msg": status_msg,
            "metadata": data.get("metadata"),
            "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),
        }


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


if __name__ == "__main__":
    main()
