from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, List
import httpx
import os
import json
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# Agnes API 配置
AGNES_API_KEY = os.environ.get("AGNES_API_KEY")
AGNES_API_URL = "https://apihub.agnes-ai.com/v1/images/generations"  # 文生图端点[citation:2][citation:4]


class GenerateRequest(BaseModel):
    """图片生成请求 - 所有参数都可选，由 n8n 传入"""
    prompt: str                                      # 必填：图片描述
    model: Optional[str] = "agnes-image-2.1-flash"  # 模型：2.0-flash(图生图)/2.1-flash(文生图)[citation:3]
    size: Optional[str] = "1024x1024"               # 图片尺寸
    negative_prompt: Optional[str] = None           # 负向提示词
    num_inference_steps: Optional[int] = 20         # 推理步数
    guidance_scale: Optional[float] = 7.5           # 引导比例
    seed: Optional[int] = None                      # 随机种子
    image_url: Optional[str] = None                 # 图生图/图编图时传入[citation:4]
    response_format: Optional[str] = "url"          # url 或 b64_json
    # n8n 工作流传过来的元数据
    row_number: Optional[int] = None
    job_id: Optional[str] = None
    topic: Optional[str] = None


class GenerateResponse(BaseModel):
    """图片生成响应"""
    status: str
    image_url: Optional[str] = None
    image_base64: Optional[str] = None
    prompt: str
    model: str
    size: str
    row_number: Optional[int] = None
    job_id: Optional[str] = None
    topic: Optional[str] = None
    error: Optional[str] = None


@app.post("/generate", response_model=GenerateResponse)
async def generate_image(request: GenerateRequest):
    """
    通用图片生成 API
    - 文生图：只传 prompt
    - 图生图/图片编辑：传 prompt + image_url
    """
    logger.info(f"Image generation request - prompt: {request.prompt[:100]}..., model: {request.model}")

    if not AGNES_API_KEY:
        raise HTTPException(status_code=500, detail="AGNES_API_KEY not configured")

    # 确定使用的模型
    # 如果提供了 image_url，自动切换到图生图模型
    model = request.model
    if request.image_url and model == "agnes-image-2.1-flash":
        model = "agnes-image-2.0-flash"  # 图生图专用模型[citation:2][citation:3]
        logger.info(f"Auto-switched to {model} for img2img")

    # 构建请求体（完全动态，由 n8n 决定传哪些参数）
    payload = {
        "model": model,
        "prompt": request.prompt,
    }

    # 可选参数 - 只有 n8n 传了才添加
    if request.size:
        payload["size"] = request.size
    if request.seed is not None:
        payload["seed"] = request.seed
    if request.negative_prompt:
        payload["negative_prompt"] = request.negative_prompt

    # extra_body 参数 - 图生图/图片编辑需要[citation:2][citation:4]
    extra_body = {}
    if request.image_url:
        extra_body["image"] = [request.image_url]
        extra_body["tags"] = ["img2img"]
        extra_body["response_format"] = request.response_format or "url"
    if request.num_inference_steps:
        extra_body["num_inference_steps"] = request.num_inference_steps
    if request.guidance_scale:
        extra_body["guidance_scale"] = request.guidance_scale

    if extra_body:
        payload["extra_body"] = extra_body

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

    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            response = await client.post(AGNES_API_URL, headers=headers, json=payload)
            response.raise_for_status()
            result = response.json()

        logger.info(f"Agnes API response status: {response.status_code}")

        # 解析响应
        image_url = None
        image_base64 = None

        if "data" in result and len(result["data"]) > 0:
            if "url" in result["data"][0]:
                image_url = result["data"][0]["url"]
            elif "b64_json" in result["data"][0]:
                image_base64 = result["data"][0]["b64_json"]

        return GenerateResponse(
            status="success",
            image_url=image_url,
            image_base64=image_base64,
            prompt=request.prompt,
            model=model,
            size=request.size or "1024x1024",
            row_number=request.row_number,
            job_id=request.job_id,
            topic=request.topic
        )

    except httpx.HTTPStatusError as e:
        logger.error(f"Agnes API error: {e.response.status_code} - {e.response.text}")
        return GenerateResponse(
            status="error",
            prompt=request.prompt,
            model=model,
            size=request.size or "1024x1024",
            row_number=request.row_number,
            job_id=request.job_id,
            topic=request.topic,
            error=f"API error: {e.response.status_code}"
        )
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        return GenerateResponse(
            status="error",
            prompt=request.prompt,
            model=model,
            size=request.size or "1024x1024",
            row_number=request.row_number,
            job_id=request.job_id,
            topic=request.topic,
            error=str(e)
        )


@app.get("/health")
async def health():
    return {"status": "ok", "service": "Agnes Image Generator"}


@app.get("/")
async def root():
    return {
        "name": "Agnes Image Generator API",
        "version": "1.0.0",
        "description": "Generic image generation API for n8n integration",
        "endpoint": "POST /generate",
        "params": {
            "prompt": "required - image description",
            "model": "optional - agnes-image-2.1-flash (T2I) or agnes-image-2.0-flash (I2I)",
            "size": "optional - e.g., 1024x1024, 1024x768",
            "image_url": "optional - for img2img/edit",
            "seed": "optional - for reproducibility",
            "row_number": "optional - for n8n workflow",
            "job_id": "optional - for n8n workflow",
            "topic": "optional - for n8n workflow"
        }
    }