"""套图生成服务：图像 provider（豆包 Seedream / 通义万相）+ 任务执行器。

Provider 调用方式移植自 ecommerce-image-suite/scripts/generate.py：
  - doubao：火山方舟 images/generations，同步返回 URL；参考图走 image 字段（data URI）
  - tongyi：wan* 万象模型走异步任务轮询；qwen* 走同步 multimodal-generation
"""
from __future__ import annotations

import asyncio
import base64
import logging
import mimetypes
import re
from uuid import UUID

import httpx
from sqlalchemy import select

from config import get_settings
from db import get_session_factory
from models import Product, ProductAsset, Suite, SuiteImage, SUITE_RUNNING, SUITE_DONE, SUITE_PARTIAL, SUITE_FAILED, STATUS_OK, STATUS_FAILED
from services import storage
from services.prompt import build_prompt, build_context, type_name, wrap_prompt_for_gpt_edits

log = logging.getLogger("suite.generator")


class ApiError(RuntimeError):
    """带 HTTP 状态码的 API 错误（用于区分可重试的网关/限流错误）。"""

    def __init__(self, message: str, status: int = 0):
        super().__init__(message)
        self.status = status


_HTML_TITLE_RE = re.compile(r"<title[^>]*>(.*?)</title>", re.IGNORECASE | re.DOTALL)


def _raise_api_error(resp, provider: str):
    """HTTP 错误时抛出带 API 错误码/信息的异常（响应体里有真正的失败原因）。"""
    if resp.is_success:
        return
    text = resp.text or ""
    if "<html" in text[:300].lower() or text.lstrip()[:15].lower().startswith("<!doctype"):
        # Cloudflare/网关错误页：取 <title> 作摘要，避免整段 HTML 进错误信息
        m = _HTML_TITLE_RE.search(text)
        detail = (re.sub(r"\s+", " ", m.group(1)).strip() if m else "") or "网关返回 HTML 错误页（上游/CDN 故障）"
        raise ApiError(f"{provider} API HTTP {resp.status_code} — {detail}", resp.status_code)
    try:
        body = resp.json()
        detail = f"{body.get('code', '')}: {body.get('message', '')}".strip(': ')
    except Exception:  # noqa: BLE001
        detail = text[:200]
    raise ApiError(f"{provider} API HTTP {resp.status_code} — {detail or '无错误详情'}", resp.status_code)

# 参考图选择：material 用第 2 张（背面/细节），其余用第 1 张（正面）
TYPE_REF_INDEX = {
    "material": 1,
}
DEFAULT_REF_COUNT = 2  # 每次生图最多带的参考图数（正面 1 张 + 背面/细节 1 张）


def _image_size(provider: str, ratio: str, is_wan: bool = True, model: str = "") -> str:
    """平台比例 → provider 尺寸参数。3:4 竖版（Ozon/WB），1:1 方图（国内）。"""
    if provider == "doubao":
        return "1536x2048" if ratio == "3:4" else "2048x2048"
    if provider == "rightapi":
        # gpt-image 自定义尺寸约束：16 的倍数、长短边比 ≤ 3:1（1536x2048 合法）
        return "1536x2048" if ratio == "3:4" else "2048x2048"
    # tongyi：万象与千问的 size 语法相同（* 分隔），档位不同
    # wan2.6 系列总像素限制在 [1280², 1440²]，wan2.7 的 1536*2048/2048*2048 会超限
    if model.startswith("wan2.6"):
        return "1152*1536" if ratio == "3:4" else "1440*1440"
    if ratio == "3:4":
        return "1536*2048" if is_wan else "768*1024"
    return "2048*2048" if is_wan else "1024*1024"

_DOUBAO_ANTI_AI = (
    "authentic real-world photography, natural imperfections, genuine texture, "
    "no synthetic look, no CGI quality, no heavy post-processing"
)

DEFAULT_NEGATIVE_PROMPT = (
    "AI-generated look, artificial, CGI quality, 3D render, synthetic texture, "
    "plastic skin, mannequin-like, too perfect, oversaturated, HDR, heavy vignette, "
    "low resolution, blurry, deformed, bad anatomy, overexposed, underexposed, grainy, "
    "watermark, text distortion, bad typography, overlapping text, cheap look, cartoon"
)


# ── 参考图解析 ────────────────────────────────────────────────────────────

def _bytes_to_data_uri(data: bytes, mime: str) -> str:
    return f"data:{mime};base64,{base64.b64encode(data).decode()}"


async def _resolve_ref_bytes(url: str) -> tuple[bytes, str]:
    """参考图 URL → (bytes, mime)。本地 media 文件直读磁盘；远程 URL 带 Referer 下载。

    生图 API 的服务器无法访问 127.0.0.1，代理 URL 也不能直接透传，
    所以统一在本地解析成原始字节再进请求体（data URI 或 multipart）。
    """
    if url.startswith("data:"):
        head, _, b64 = url.partition(",")
        mime = head[5:].split(";", 1)[0] or "image/jpeg"
        return base64.b64decode(b64), mime
    path = storage.local_path(url)
    if path is not None:
        mime = mimetypes.guess_type(path.name)[0] or "image/jpeg"
        return path.read_bytes(), mime
    if url.startswith(("http://", "https://")):
        from api.proxy import guess_referer
        data, ctype = await storage.download_bytes(url, referer=guess_referer(url))
        if not ctype.startswith("image/"):
            ctype = "image/jpeg"
        return data, ctype
    raise FileNotFoundError(f"无法解析参考图: {url}")


async def _resolve_ref(url: str) -> str:
    data, mime = await _resolve_ref_bytes(url)
    return _bytes_to_data_uri(data, mime)


# ── Provider：豆包 Seedream（火山方舟）────────────────────────────────────

async def generate_doubao(prompt: str, ref_images: list[str], size: str = "2048x2048", model: str | None = None) -> bytes:
    s = get_settings()
    if not s.ark_api_key:
        raise RuntimeError("未配置 ARK_API_KEY（.env）")
    body = {
        "model": model or s.ark_image_model,
        "prompt": prompt.rstrip(". ") + ". " + _DOUBAO_ANTI_AI,
        "size": size,
        "response_format": "url",
        "watermark": False,
        "n": 1,
    }
    if ref_images:
        body["image"] = [await _resolve_ref(u) for u in ref_images]
    async with httpx.AsyncClient(timeout=s.request_timeout, verify=False) as client:
        resp = await client.post(
            s.ark_base_url,
            headers={"Authorization": f"Bearer {s.ark_api_key}", "Content-Type": "application/json"},
            json=body,
        )
        _raise_api_error(resp, "豆包")
        img_url = resp.json()["data"][0]["url"]
        dl = await client.get(img_url, timeout=s.request_timeout)
        dl.raise_for_status()
        return dl.content


# ── Provider：通义万相 / 千问（DashScope）────────────────────────────────

def _is_wan_model(model: str) -> bool:
    return model.lower().startswith("wan")


def _is_t2i_model(model: str) -> bool:
    """纯文生图模型（如 wan2.6-t2i）：不接受参考图，商品一致性只能靠文案描述。"""
    return "t2i" in model.lower()


async def _tongyi_poll_task(client: httpx.AsyncClient, key: str, task_id: str, max_wait: int) -> str:
    poll_url = "https://dashscope.aliyuncs.com/api/v1/tasks/" + task_id
    elapsed, interval = 0, 3
    while elapsed < max_wait:
        resp = await client.get(poll_url, headers={"Authorization": f"Bearer {key}"}, timeout=30)
        resp.raise_for_status()
        result = resp.json()
        status = result.get("output", {}).get("task_status", "")
        if status == "SUCCEEDED":
            choices = result["output"].get("choices", [])
            if choices:
                content = choices[0].get("message", {}).get("content", [])
                if content:
                    return content[0].get("image", "")
            results = result["output"].get("results", [])
            if results:
                return results[0].get("url") or results[0].get("b64_image", "")
            raise RuntimeError(f"通义任务成功但无结果: {result}")
        if status in ("FAILED", "UNKNOWN"):
            raise RuntimeError(f"通义任务失败: {result}")
        await asyncio.sleep(interval)
        elapsed += interval
        interval = min(interval + 2, 10)
    raise TimeoutError(f"通义异步任务超时 ({max_wait}s): task_id={task_id}")


async def generate_tongyi(prompt: str, ref_images: list[str], size: str = "2048*2048", model: str | None = None) -> bytes:
    s = get_settings()
    if not s.dashscope_api_key:
        raise RuntimeError("未配置 DASHSCOPE_API_KEY（.env）")
    model = model or s.dashscope_model
    is_wan = _is_wan_model(model)
    url = s.dashscope_base_url or (
        "https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation"
        if is_wan
        else "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
    )

    # t2i 模型不接受参考图：content 只有文本，商品一致性依赖 prompt 里的标题/卖点描述
    content: list[dict] = []
    if not _is_t2i_model(model):
        content = [{"image": await _resolve_ref(u)} for u in ref_images]
    content.append({"text": prompt})

    params = {"size": size, "n": 1, "watermark": False}
    if not is_wan:
        params["prompt_extend"] = False
    params["negative_prompt"] = DEFAULT_NEGATIVE_PROMPT[:500]

    headers = {"Authorization": f"Bearer {s.dashscope_api_key}", "Content-Type": "application/json"}
    if is_wan:
        headers["X-DashScope-Async"] = "enable"

    body = {"model": model, "input": {"messages": [{"role": "user", "content": content}]}, "parameters": params}

    async with httpx.AsyncClient(timeout=s.request_timeout, verify=False) as client:
        resp = await client.post(url, headers=headers, json=body)
        _raise_api_error(resp, "通义")
        data = resp.json()
        if is_wan:
            task_id = data.get("output", {}).get("task_id", "")
            if not task_id:
                raise RuntimeError(f"通义万象未返回 task_id: {data}")
            img_url = await _tongyi_poll_task(client, s.dashscope_api_key, task_id, s.poll_max_wait)
            if img_url.startswith("data:") or len(img_url) > 500:
                return base64.b64decode(img_url.split(",", 1)[-1] if "," in img_url else img_url)
            dl = await client.get(img_url, timeout=s.request_timeout)
            dl.raise_for_status()
            return dl.content
        img_url = data["output"]["choices"][0]["message"]["content"][0]["image"]
        dl = await client.get(img_url, timeout=s.request_timeout)
        dl.raise_for_status()
        return dl.content


# ── Provider：RightAPI（gpt-image，OpenAI 兼容中转）──────────────────────

# 可重试的状态码：中转限流/网关抖动（该中转限流时返回 Cloudflare 502 而非 429）
RETRYABLE_STATUS = {429, 500, 502, 503, 504}

# 中转对 input_fidelity 参数的支持探测：None=未探测，True=支持，False=不支持（已降级）
_rightapi_fidelity_supported: bool | None = None


async def _rightapi_request(s, prompt: str, ref_images: list[str], size: str, model: str) -> bytes:
    """gpt-image 系列：有参考图走 /v1/images/edits（multipart），无参考图走 /v1/images/generations。

    OpenAI 兼容协议：响应固定 b64_json（不支持 response_format 参数，传了报 400）；
    同步调用无任务轮询，高质量档单张 1-5 分钟，超时按文档建议兜底 600s。
    input_fidelity=high 强制高保真保留输入图细节（商品一致性关键参数，仅 edits 端点）；
    中转若不认该参数（400），自动去掉重试并记住，后续请求不再带。
    """
    global _rightapi_fidelity_supported
    base = s.rightapi_base_url.rstrip("/")
    headers = {"Authorization": f"Bearer {s.rightapi_api_key}"}
    use_fidelity = bool(ref_images) and s.rightapi_input_fidelity and _rightapi_fidelity_supported is not False

    async with httpx.AsyncClient(timeout=max(s.request_timeout, 600), verify=False) as client:
        common = {
            "model": model,
            "prompt": prompt,
            "size": size,
            "quality": s.rightapi_image_quality,
            "output_format": "jpeg",  # 与落盘 .jpg 后缀一致
            "n": 1,
        }
        if use_fidelity:
            common["input_fidelity"] = s.rightapi_input_fidelity
        if ref_images:
            files = []
            for i, u in enumerate(ref_images):
                data, mime = await _resolve_ref_bytes(u)
                files.append(("image[]", (f"ref-{i + 1}.{mime.split('/')[-1]}", data, mime)))
            resp = await client.post(f"{base}/v1/images/edits", headers=headers, files=files, data=common)
            # 中转不认 input_fidelity：去掉参数重试一次（仅一次探测）
            if resp.status_code == 400 and use_fidelity and "input_fidelity" in resp.text:
                _rightapi_fidelity_supported = False
                log.warning("RightAPI 不支持 input_fidelity 参数，已自动去掉并降级（后续请求不再带）")
                common.pop("input_fidelity", None)
                resp = await client.post(f"{base}/v1/images/edits", headers=headers, files=files, data=common)
            elif resp.is_success and use_fidelity:
                _rightapi_fidelity_supported = True
        else:
            resp = await client.post(
                f"{base}/v1/images/generations",
                headers={**headers, "Content-Type": "application/json"},
                json=common,
            )
        _raise_api_error(resp, "RightAPI")
        item = resp.json()["data"][0]
        b64 = item.get("b64_json") or ""
        if b64:
            return base64.b64decode(b64.split(",", 1)[-1] if "," in b64 else b64)
        img_url = item.get("url") or ""
        if not img_url:
            raise RuntimeError(f"RightAPI 响应里没有图片数据: {item}")
        dl = await client.get(img_url, timeout=s.request_timeout)
        dl.raise_for_status()
        return dl.content


async def generate_rightapi(prompt: str, ref_images: list[str], size: str = "2048x2048", model: str | None = None) -> bytes:
    """带重试的 RightAPI 入口：429/5xx/超时按递增间隔重试。

    实测该中转对同 key 连续请求有分钟级冷却（成功一张后紧接着的请求会被网关秒拒 502），
    60s → 120s → 240s 的退避基本能等到窗口放开。
    """
    s = get_settings()
    if not s.rightapi_api_key:
        raise RuntimeError("未配置 RIGHTAPI_API_KEY（.env）")
    model = model or s.rightapi_image_model
    attempts = max(1, s.rightapi_max_retries)

    last_exc: Exception | None = None
    for i in range(attempts):
        try:
            return await _rightapi_request(s, prompt, ref_images, size, model)
        except ApiError as exc:
            last_exc = exc
            if exc.status not in RETRYABLE_STATUS:
                raise  # 参数错误等不可重试，立即失败
        except (httpx.TimeoutException, httpx.TransportError) as exc:
            last_exc = exc  # 网络抖动/超时可重试
        if i == attempts - 1:
            break
        wait = s.rightapi_retry_wait * (2 ** i)
        log.warning("RightAPI 第 %d/%d 次请求失败（%s），%ds 后重试", i + 1, attempts, last_exc, wait)
        await asyncio.sleep(wait)
    raise last_exc  # type: ignore[misc]


GENERATORS = {"doubao": generate_doubao, "tongyi": generate_tongyi, "rightapi": generate_rightapi}


# ── 任务执行器 ────────────────────────────────────────────────────────────

def _order_refs(refs: list[str], type_id: str) -> list[str]:
    """参考图槽位选择 + 截断：material 偏好第 2 张，其余用第 1 张。"""
    preferred = TYPE_REF_INDEX.get(type_id)
    if preferred is not None and len(refs) > preferred:
        refs = [refs[preferred]] + [r for i, r in enumerate(refs) if i != preferred]
    return refs[:DEFAULT_REF_COUNT]


def _refs_for_job(images: list[dict], job: dict) -> list[str]:
    """无状态路径：按方案项选参考图。

    优先 variant_name 精确匹配（「主图·粉色」用粉色那张 SKU 图）；
    匹配不到则回退 main 组第一张（再退到任意第一张）。
    """
    variant = job.get("variant_name")
    if variant:
        matched = [i["url"] for i in images if i.get("variant_name") == variant]
        if matched:
            return matched[:DEFAULT_REF_COUNT]
    mains = [i["url"] for i in images if i.get("group_key") == "main"]
    others = [i["url"] for i in images if i.get("group_key") != "main"]
    pool = mains or others or [i["url"] for i in images]
    if not pool:
        raise RuntimeError("任务没有参考图")
    return _order_refs(pool, job.get("kind", ""))


async def _select_ref_images(db, product_id: UUID, type_id: str) -> list[str]:
    """商品路径：主图组前几张。转存完成的用本地文件，未完成的直接用源站 URL。"""
    assets = (await db.scalars(
        select(ProductAsset).where(
            ProductAsset.product_id == product_id,
            ProductAsset.group_key == "main",
            ProductAsset.type == "img",
        ).order_by(ProductAsset.sort_order)
    )).all()
    refs = [a.stored_url or a.source_url for a in assets if (a.stored_url or a.source_url)]
    if not refs:
        raise RuntimeError("商品没有可用参考图（未采集主图）")
    return _order_refs(refs, type_id)


async def run_suite(suite_id: str) -> None:
    """后台执行套图任务：逐张生成 → 落盘 → 记录；单张失败不中断。

    两条路径：
      - 无状态（product_id 为空）：上下文与参考图来自请求自带的 context / ref_images
      - 商品路径（兼容旧流程）：从 product + product_assets 取
    """
    settings = get_settings()
    async with get_session_factory()() as db:
        suite = await db.get(Suite, UUID(suite_id))
        if suite is None:
            return

        product = None
        if suite.product_id:
            product = await db.get(Product, suite.product_id)
            if product is None:
                suite.status = SUITE_FAILED
                suite.error = "商品不存在"
                await db.commit()
                return

        suite.status = SUITE_RUNNING
        await db.commit()

        provider_name = suite.provider or settings.image_provider
        generator = GENERATORS.get(provider_name)
        if generator is None:
            suite.status = SUITE_FAILED
            suite.error = f"未知 provider: {provider_name}"
            await db.commit()
            return

        raw = suite.context if not product else (product.raw or {})
        ctx = build_context(raw or {}, fallback_name=product.name if product else "")
        model = suite.model or {
            "tongyi": settings.dashscope_model,
            "rightapi": settings.rightapi_image_model,
        }.get(provider_name, settings.ark_image_model)
        is_wan = provider_name == "tongyi" and _is_wan_model(model)
        size = _image_size(provider_name, suite.ratio, is_wan=is_wan, model=model)

        # 任务列表：方案（逐张）优先，旧路径按 types
        if suite.plan:
            jobs = [dict(j) for j in suite.plan]
        else:
            jobs = [
                {"kind": t, "title": type_name(t), "detail": "", "prompt_hint": "", "variant_name": None}
                for t in (suite.types or [])
            ]

        ok, failed = 0, 0
        failures: list[str] = []
        for job in jobs:
            type_id = job["kind"]
            image_row = SuiteImage(
                suite_id=suite.id,
                type_id=type_id,
                name=job.get("title") or type_name(type_id),
                status=STATUS_FAILED,
            )
            db.add(image_row)
            await db.flush()
            try:
                prompt = build_prompt(
                    type_id, ctx, suite.style_set, suite.lang,
                    extra=job, style_prompt=suite.style_prompt, requirements=suite.requirements,
                )
                # gpt-image edits 语义：商品冻结契约前置，防止风格词改商品
                if provider_name == "rightapi":
                    prompt = wrap_prompt_for_gpt_edits(prompt)
                if product:
                    refs = await _select_ref_images(db, product.id, type_id)
                else:
                    refs = _refs_for_job(list(suite.ref_images or []), job)
                data = await generator(prompt, refs, size=size, model=model)
                # 部分中转不遵守 output_format（要 jpeg 回 PNG），按魔数定扩展名
                ext = ".png" if data[:8] == b"\x89PNG\r\n\x1a\n" else ".jpg"
                key = storage.write_bytes(data, key_prefix=f"suites/{suite.id}", ext=ext)
                image_row.stored_url = storage.public_url(key)
                image_row.status = STATUS_OK
                ok += 1
            except Exception as exc:  # noqa: BLE001
                log.exception("套图 %s 类型 %s 生成失败", suite_id, type_id)
                err = str(exc)[:500]
                image_row.error = err
                failures.append(f"{job.get('title') or type_name(type_id)}：{err[:200]}")
                failed += 1
            await db.commit()

        suite.status = SUITE_DONE if failed == 0 else (SUITE_PARTIAL if ok > 0 else SUITE_FAILED)
        if failed:
            uniq = list(dict.fromkeys(failures))  # 去重保序
            detail = "；".join(uniq[:6])
            if len(uniq) > 6:
                detail += f"；…等共 {failed} 张失败"
            if ok == 0:
                suite.error = f"全部生成失败。{detail}"
            else:
                suite.error = f"部分生成失败（{failed} 张）。{detail}"
        from datetime import datetime, timezone
        suite.finished_at = datetime.now(timezone.utc)
        if product:
            product.stage = "generated"  # 商品路径才有的阶段升级
        await db.commit()