import os
import base64
import json
import uuid
import threading
import traceback
import time
from datetime import datetime
import requests
from io import BytesIO

from PIL import Image
from fastapi import FastAPI, HTTPException, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from pydantic import BaseModel

def log(msg: str):
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
    print(f"[{ts}] {msg}")

app = FastAPI()

# =========================
# CORS
# =========================

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# =========================
# ENV
# =========================

OPENROUTER_KEY = os.getenv("OPENROUTER_API_KEY")

# 给插件调用的密钥
API_KEY = os.getenv("API_KEY")

# 未来接 Modal 时使用
MODAL_API_URL = os.getenv("MODAL_API_URL")

# =========================
# 异步任务存储 & 图片存储（分离图片数据，避免 JSON 响应过大）
# =========================

tasks = {}
image_store = {}  # task_id -> (bytes, content_type)

def extract_image_bytes(data: dict) -> tuple:
    """
    从 Modal 响应中提取图片
    支持：
    - image_base64
    - base64
    - image_data
    - image_url
    - url
    """

    # ===== 你的 Modal 返回格式 =====
    if "image_base64" in data and data["image_base64"]:
        return (
            base64.b64decode(data["image_base64"]),
            "image/png"
        )

    # ===== 常见 Base64 字段 =====
    for field in [
        "base64",
        "image_data",
        "image"
    ]:
        value = data.get(field)

        if value and isinstance(value, str):

            if "," in value:
                value = value.split(",", 1)[1]

            return (
                base64.b64decode(value),
                "image/png"
            )

    # ===== URL 字段 =====
    for field in [
        "url",
        "image_url"
    ]:
        value = data.get(field)

        if value and isinstance(value, str):

            resp = requests.get(
                value,
                timeout=60
            )

            resp.raise_for_status()

            return (
                resp.content,
                resp.headers.get(
                    "content-type",
                    "image/png"
                )
            )

    # ===== 兼容数组格式 =====
    if isinstance(data.get("images"), list):

        first = data["images"][0]

        if isinstance(first, dict):
            return extract_image_bytes(first)

    if isinstance(data.get("output"), list):

        first = data["output"][0]

        if isinstance(first, dict):
            return extract_image_bytes(first)

    raise ValueError(
        f"无法解析图片响应: {str(data)[:500]}"
    )

def background_generate(task_id: str, image_base64: str):
    t_start = time.time()
    log(f"[{task_id}] START background_generate")

    try:
        tasks[task_id] = {"status": "processing"}

        t0 = time.time()
        processed = resize_image(image_base64)
        log(f"[{task_id}] resize_image took {time.time()-t0:.2f}s")

        t0 = time.time()
        vision_result = call_openrouter(processed, PROMPT_INSTRUCTION)
        log(f"[{task_id}] call_openrouter took {time.time()-t0:.2f}s")

        t0 = time.time()
        prompt_text = vision_result["choices"][0]["message"]["content"]
        log(f"[{task_id}] prompt: {prompt_text[:120]}")

        t0 = time.time()
        modal_response = requests.post(
            MODAL_API_URL,
            json={"prompt": prompt_text},
            timeout=600
        )
        modal_response.raise_for_status()
        image_data = modal_response.json()
        log(f"[{task_id}] Modal API took {time.time()-t0:.2f}s")

        log(f"[{task_id}] ====== MODAL RESPONSE ======")
        log(f"[{task_id}] {image_data}")

        t0 = time.time()
        img_bytes, content_type = extract_image_bytes(image_data)
        log(f"[{task_id}] extract_image_bytes took {time.time()-t0:.2f}s")

        log(f"[{task_id}] Image Size: {len(img_bytes)/1024:.1f} KB")

        image_store[task_id] = (img_bytes, content_type)

        tasks[task_id] = {
            "status": "completed",
            "result": {"image_url": f"/image/{task_id}"},
            "prompt": prompt_text
        }

        log(f"[{task_id}] COMPLETED in {time.time()-t_start:.2f}s")
    except Exception as e:
        tasks[task_id] = {
            "status": "failed",
            "error": str(e),
            "traceback": traceback.format_exc()
        }

        log(f"[{task_id}] FAILED after {time.time()-t_start:.2f}s: {e}")
        log(traceback.format_exc())

# =========================
# PROMPTS
# =========================

META_INSTRUCTION = """
你是专业图片分析助手。

输出 JSON 格式。

要求：

1. gender
2. age_estimate
3. hairstyle
4. clothing
5. pose
6. facial_expression
7. environment
8. lighting
9. camera_angle
10. text_content

直接输出 JSON。
不要 Markdown。
"""

PROMPT_INSTRUCTION = """
Analyze the image and generate a high quality AI image generation prompt.

Requirements:

- English only
- Suitable for FLUX and SDXL
- Include:
  subject,
  clothing,
  pose,
  environment,
  lighting,
  camera angle,
  artistic details

Output only prompt text.

No markdown.
No explanations.
"""

# =========================
# MODELS
# =========================

class ImageRequest(BaseModel):
    image_base64: str

# =========================
# HELPERS
# =========================

def verify_key(api_key: str):
    if API_KEY and api_key != API_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid API Key"
        )


def resize_image(base64_str, max_size=1024):

    if "," in base64_str:
        base64_str = base64_str.split(",", 1)[1]

    try:
        img_data = base64.b64decode(base64_str)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"base64 解码失败: {str(e)}")

    try:
        img = Image.open(BytesIO(img_data))
        img.verify()
        img = Image.open(BytesIO(img_data))
    except Exception as e:
        preview = img_data[:200]
        raise HTTPException(
            status_code=400,
            detail=f"无法识别图片格式: {str(e)} | 数据前200字节: {preview}"
        )

    if img.mode in ("RGBA", "P"):
        img = img.convert("RGB")

    img.thumbnail(
        (max_size, max_size),
        Image.Resampling.LANCZOS
    )

    buffer = BytesIO()

    img.save(
        buffer,
        format="JPEG",
        quality=85
    )

    return base64.b64encode(
        buffer.getvalue()
    ).decode()


OPENROUTER_FALLBACK_MODELS = [
    "nvidia/nemotron-nano-12b-v2-vl:free",
    "openrouter/free",
]


def call_openrouter(
    image_base64: str,
    instruction: str
):

    if not OPENROUTER_KEY:
        raise HTTPException(
            status_code=500,
            detail="OPENROUTER_API_KEY missing"
        )

    url = "https://openrouter.ai/api/v1/chat/completions"
    models = ["nex-agi/nex-n2-pro:free", *OPENROUTER_FALLBACK_MODELS]

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

    last_error = None

    for model in models:
        t0 = time.time()
        log(f"call_openrouter trying model: {model}")
        payload = {
            "model": model,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": instruction
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_base64}"
                            }
                        }
                    ]
                }
            ],
            "temperature": 0.1
        }

        try:
            response = requests.post(
                url,
                headers=headers,
                json=payload,
                timeout=90
            )

            elapsed = time.time() - t0
            log(f"call_openrouter model={model} status={response.status_code} took={elapsed:.2f}s")

            if response.status_code == 200:
                return response.json()

            last_error = (
                f"Model {model} failed: "
                f"{response.status_code} {response.text[:500]}"
            )

        except Exception as e:
            elapsed = time.time() - t0
            log(f"call_openrouter model={model} error after {elapsed:.2f}s: {e}")
            last_error = f"Model {model} error: {str(e)}"

    raise HTTPException(
        status_code=500,
        detail=f"所有模型均失败: {last_error}"
    )

# =========================
# ROUTES
# =========================

@app.get("/")
def home():
    return {
        "status": "running",
        "service": "XiMa API"
    }


@app.post("/analyze")
async def analyze(
    request: ImageRequest,
    x_api_key: str = Header(default="")
):

    verify_key(x_api_key)

    try:

        processed = resize_image(
            request.image_base64
        )

        result = call_openrouter(
            processed,
            META_INSTRUCTION
        )

        content = (
            result["choices"][0]
            ["message"]
            ["content"]
        )

        return {
            "success": True,
            "analysis": content
        }

    except Exception as e:

        raise HTTPException(
            status_code=500,
            detail=str(e)
        )


@app.post("/prompt")
async def prompt(
    request: ImageRequest,
    x_api_key: str = Header(default="")
):

    verify_key(x_api_key)

    try:

        processed = resize_image(
            request.image_base64
        )

        result = call_openrouter(
            processed,
            PROMPT_INSTRUCTION
        )

        prompt_text = (
            result["choices"][0]
            ["message"]
            ["content"]
        )

        return {
            "success": True,
            "prompt": prompt_text
        }

    except Exception as e:

        raise HTTPException(
            status_code=500,
            detail=str(e)
        )


@app.post("/generate")
async def generate(
    request: ImageRequest,
    x_api_key: str = Header(default="")
):

    verify_key(x_api_key)

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

    task_id = str(uuid.uuid4())
    tasks[task_id] = {"status": "pending"}

    thread = threading.Thread(
        target=background_generate,
        args=(task_id, request.image_base64)
    )
    thread.start()

    return {"success": True, "task_id": task_id}


@app.get("/task/{task_id}")
async def get_task(task_id: str):
    task = tasks.get(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    return task


@app.get("/image/{task_id}")
async def get_image(task_id: str):
    entry = image_store.get(task_id)
    if not entry:
        raise HTTPException(status_code=404, detail="Image not found")
    image_bytes, content_type = entry
    return Response(content=image_bytes, media_type=content_type)


if __name__ == "__main__":

    import uvicorn

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