#!/usr/bin/env python3
"""通过 FastAPI 提供 GPT Image 和 Grok Image 文生图服务。"""

import base64
import os
import time
import urllib.error
import urllib.request
from pathlib import Path
from uuid import uuid4

from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from openai import APIConnectionError, APIStatusError, APITimeoutError, OpenAI
from pydantic import BaseModel, Field, model_validator

load_dotenv()

API_BASE_URL = os.getenv("IMAGE_API_BASE_URL", "https://api.slomerex.xyz/v1")
OUTPUT_DIR = Path("output")
MAX_RETRIES = 3
RETRYABLE_STATUS_CODES = {408, 429, 500, 502, 503, 504}
PROVIDERS = {
    "grok": {
        "api_key_env": "GROK_API_KEY",
        "model": "grok-imagine-image-lite",
    },
    "gpt": {
        "api_key_env": "GPT_API_KEY",
        "model": "gpt-image-2-1K",
    },
}

app = FastAPI(title="AI Image T2I API", version="1.0.0")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")


class ImageGenerationRequest(BaseModel):
    prompt: str = Field(min_length=1, description="图片提示词")
    provider: str = Field(default="grok", pattern="^(grok|gpt)$", description="模型提供方")
    model: str | None = Field(default=None, description="覆盖提供方默认模型")
    size: str = Field(default="1024x1024", description="图片尺寸")
    n: int = Field(default=1, ge=1, le=10, description="生成数量")
    profile_name: str = Field(pattern=r"^[\w-]+$", description="调用方标识（必填），图片会保存到 output/<profile_name>/ 下")

    @model_validator(mode="before")
    @classmethod
    def check_profile_name(cls, data):
        if isinstance(data, dict) and not data.get("profile_name"):
            raise ValueError("缺少 profile_name 参数，请在请求中带上你的 agent/profile 名称")
        return data


class GeneratedImage(BaseModel):
    filename: str
    url: str
    bytes: int


class ImageGenerationResponse(BaseModel):
    provider: str
    model: str
    images: list[GeneratedImage]


def generate_with_retry(client: OpenAI, *, model: str, prompt: str, size: str, n: int):
    """调用图片接口，并对可恢复的上游错误进行指数退避重试。"""
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            return client.images.generate(model=model, prompt=prompt, size=size, n=n)
        except APIStatusError as e:
            if e.status_code not in RETRYABLE_STATUS_CODES or attempt == MAX_RETRIES:
                raise HTTPException(status_code=502, detail=f"上游图片服务错误: HTTP {e.status_code}") from e
            time.sleep(2 ** (attempt - 1))
        except (APIConnectionError, APITimeoutError) as e:
            if attempt == MAX_RETRIES:
                raise HTTPException(status_code=504, detail="上游图片服务连接或请求超时") from e
            time.sleep(2 ** (attempt - 1))


def download_image(url: str) -> bytes:
    if url.startswith("data:image/"):
        try:
            _, encoded = url.split(",", 1)
            return base64.b64decode(encoded, validate=True)
        except (ValueError, TypeError) as e:
            raise HTTPException(status_code=502, detail="图片 data URL 格式无效") from e

    request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            with urllib.request.urlopen(request, timeout=120) as response:
                return response.read()
        except urllib.error.HTTPError as e:
            if e.code not in RETRYABLE_STATUS_CODES or attempt == MAX_RETRIES:
                raise HTTPException(status_code=502, detail=f"图片 URL 下载失败: HTTP {e.code}") from e
        except (urllib.error.URLError, TimeoutError) as e:
            if attempt == MAX_RETRIES:
                raise HTTPException(status_code=504, detail="图片 URL 下载超时或网络连接失败") from e
        time.sleep(2 ** (attempt - 1))


@app.get("/health")
def health_check():
    return {"status": "ok"}


@app.post("/images/generations", response_model=ImageGenerationResponse)
def generate_images(request: ImageGenerationRequest):
    config = PROVIDERS[request.provider]
    api_key = os.getenv(config["api_key_env"])
    if not api_key:
        raise HTTPException(status_code=500, detail=f"未配置环境变量 {config['api_key_env']}")

    model = request.model or config["model"]
    profile_dir = OUTPUT_DIR / request.profile_name
    profile_dir.mkdir(parents=True, exist_ok=True)
    client = OpenAI(base_url=API_BASE_URL, api_key=api_key, timeout=120.0, max_retries=0)
    response = generate_with_retry(
        client,
        model=model,
        prompt=request.prompt,
        size=request.size,
        n=request.n,
    )

    images = []
    for item in response.data:
        try:
            if item.b64_json:
                raw = base64.b64decode(item.b64_json, validate=True)
            elif item.url:
                raw = download_image(item.url)
            else:
                raise HTTPException(status_code=502, detail="上游图片服务未返回 b64_json 或 url")
        except ValueError as e:
            raise HTTPException(status_code=502, detail="上游 Base64 图片数据格式无效") from e

        filename = f"{model}-{uuid4().hex}.png"
        (profile_dir / filename).write_bytes(raw)
        images.append(GeneratedImage(filename=filename, url=f"/output/{request.profile_name}/{filename}", bytes=len(raw)))

    return ImageGenerationResponse(provider=request.provider, model=model, images=images)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)