"""阶段 3：对 Trellis 抽帧结果做 IP-Adapter + ControlNet 重绘，还原原画风。

输入：
- original_image：用户上传的原图（提供风格 / IP-Adapter 参考）
- structure_image：Stage 2 抽出的 normal 帧（提供几何 / ControlNet 结构约束）

输出：
- 单张保留原画风且视角正确的 PNG，落盘到 GENERATED_DIR/ 下

模型：chigozienri/ip_adapter-sdxl-controlnet-depth（配置见 config.IMAGE_MODELS）
- $0.07/张，~72s，硬件 L40S
- `image` 字段 = IP-Adapter 参考图（风格源）
- `controlnet_input` 字段 = 结构参考图（当前一期传 normal 帧作为 depth 近似）

调参建议（可在 IMAGE_MODELS 默认值基础上覆盖）：
- scale (IP-Adapter 权重): 0.7-0.8，越高越贴近原图风格
- controlnet_conditioning_scale (结构权重): 0.7-0.9，越高越保几何
"""

from __future__ import annotations

import asyncio
import logging
import os
import uuid
from pathlib import Path
from typing import Any

import httpx

from app.config import get_image_model_config, get_style_restore_model_id
from app.services.image_gen import GENERATED_DIR, _make_replicate_client, to_data_uri

logger = logging.getLogger(__name__)

_replicate_client = _make_replicate_client()


async def _download_image(url: str, prefix: str) -> Path:
    """把 Replicate 返回的图片下载到 generated/。"""
    filename = f"{prefix}_{uuid.uuid4().hex[:12]}.png"
    filepath = GENERATED_DIR / filename
    proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY")
    async with httpx.AsyncClient(proxy=proxy, timeout=httpx.Timeout(180.0)) as client:
        resp = await client.get(url, follow_redirects=True)
        resp.raise_for_status()
        filepath.write_bytes(resp.content)
    return filepath


def _first_url(output: Any) -> str | None:
    """Replicate 输出有时是 list[FileOutput]、有时是单个；统一取第一张。"""
    if output is None:
        return None
    if isinstance(output, list):
        if not output:
            return None
        return str(output[0])
    return str(output)


async def restore(
    original_image: str,
    structure_image: str,
    azimuth: int = 0,
    elevation: int = 0,
    params: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """对单张 normal 帧做 IP-Adapter + ControlNet 重绘。

    Args:
        original_image: 用户原图，本地路径 / data URI / URL 均可
        structure_image: Stage 2 抽出的 normal 帧，同上
        azimuth / elevation: 仅用于日志与返回元数据，不参与模型调用
        params: 覆盖默认参数的字典（scale / controlnet_conditioning_scale / prompt 等）

    Returns:
        {
          "success": bool,
          "url": "/generated/restyled_xxx.png" | None,
          "azimuth": int,
          "elevation": int,
          "error": str | None,
        }
    """
    config = get_image_model_config(get_style_restore_model_id())
    replicate_model_id = config["model_id"]
    merged_params = dict(config.get("default_params", {}))
    if params:
        merged_params.update(params)

    GENERATED_DIR.mkdir(parents=True, exist_ok=True)

    try:
        image_uri = to_data_uri(original_image)
        structure_uri = to_data_uri(structure_image)

        input_params: dict[str, Any] = {
            "image": image_uri,                      # IP-Adapter 参考图 = 用户原图
            "controlnet_input": structure_uri,       # 结构参考 = normal 帧
            **merged_params,
        }

        logger.info(
            "Stage 3 重绘 az=%d el=%d: scale=%s, cn_scale=%s",
            azimuth, elevation,
            merged_params.get("scale"),
            merged_params.get("controlnet_conditioning_scale"),
        )

        output = await _replicate_client.async_run(
            replicate_model_id, input=input_params, wait=False
        )

        url = _first_url(output)
        if not url or not (url.startswith("http://") or url.startswith("https://")):
            return {
                "success": False,
                "url": None,
                "azimuth": azimuth,
                "elevation": elevation,
                "error": f"模型未返回有效图片 URL: {str(output)[:200]}",
            }

        local = await _download_image(url, prefix=f"restyled_az{azimuth}_el{elevation}")
        return {
            "success": True,
            "url": f"/generated/{local.name}",
            "azimuth": azimuth,
            "elevation": elevation,
            "error": None,
        }

    except Exception as e:
        detail = str(e) or f"{type(e).__name__}: {repr(e)}"
        logger.error(
            "Stage 3 重绘失败 az=%d el=%d: %s", azimuth, elevation, detail, exc_info=True
        )
        return {
            "success": False,
            "url": None,
            "azimuth": azimuth,
            "elevation": elevation,
            "error": detail,
        }


async def restore_batch(
    original_image: str,
    frames: list[dict[str, Any]],
    params: dict[str, Any] | None = None,
    concurrency: int = 3,
) -> list[dict[str, Any]]:
    """对一批 frames 做 Stage 3 重绘，带并发控制。

    Args:
        original_image: 用户原图（所有帧共用同一张风格参考）
        frames: [{"az", "el", "normal_png", ...}, ...]，normal_png 作为结构参考
        params: 透传给 restore() 的覆盖参数
        concurrency: 最大并发数；Replicate 并发打太猛容易撞 rate limit，3 比较稳

    Returns:
        列表，顺序与 frames 对齐，元素与 restore() 的返回一致。
    """
    if not frames:
        return []

    sem = asyncio.Semaphore(max(1, concurrency))

    async def _one(f: dict[str, Any]) -> dict[str, Any]:
        async with sem:
            return await restore(
                original_image=original_image,
                structure_image=f["normal_png"],
                azimuth=int(f.get("az", 0)),
                elevation=int(f.get("el", 0)),
                params=params,
            )

    return await asyncio.gather(*[_one(f) for f in frames])