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

import base64
import os
import time
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

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="生成数量")


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


@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"]
    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:
        if not item.b64_json:
            raise HTTPException(status_code=502, detail="上游图片服务未返回 Base64 图片数据")

        try:
            raw = base64.b64decode(item.b64_json)
        except ValueError as e:
            raise HTTPException(status_code=502, detail="上游图片数据格式无效") from e

        filename = f"{model}-{uuid4().hex}.png"
        (OUTPUT_DIR / filename).write_bytes(raw)
        images.append(GeneratedImage(filename=filename, url=f"/output/{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)