"""OpenAI 图像后端（生成设计稿 + 三图合成）。

支持 OpenAI images/generations（文生图）与 images/edits（img2img 多参考图），
兼容两类网关返回：
  1) 同步：{"data": [{"b64_json" | "url"}]}
  2) 异步任务（ai-media.vip 等重负载自动转异步）：
     提交返回 202 + {"object":"image.task","task_id","poll_url","poll_after_ms"}，
     需轮询 GET {base}/images/tasks/{task_id} 直到 succeeded，再从成功响应取图。

所有请求跳过环境代理（NO_PROXY），适配用户挂 VPN 时直连国内/自建网关。
"""
import base64
import time
from pathlib import Path
from typing import Optional

import requests
import os

# 模型调用一律直连：用户常开 VPN（系统代理），网关多为国内/自建，走代理会被拦截或变慢。
# 环境变量级 NO_PROXY 双保险（requests/urllib3 均读取），Google 采集（pytrends）不受影响。
os.environ.setdefault('NO_PROXY', '*')
os.environ.setdefault('no_proxy', '*')

from .base import ImageBackend

# 直连策略：忽略环境代理（用户挂 VPN 时代理会拦截国内/自建网关的请求）
NO_PROXY = {"http": None, "https": None}

_FINAL_STATUS = ("succeeded", "completed", "done")
_FAIL_STATUS = ("failed", "error")


def _shrink_blob_to_2mb(img_path: str, blob: bytes, max_bytes: int = 2 * 1024 * 1024) -> bytes:
    """把大图压缩到 <max_bytes（默认 2MB）再上传：
    - 尺寸过大先缩放（合成输入 1600x2200 内足够）；
    - 有透明通道 → PNG 优化；否则 → JPEG 循环降质。
    失败时返回原 blob（不阻塞流程）。"""
    try:
        import io
        from PIL import Image
        img = Image.open(io.BytesIO(blob))
        max_w, max_h = 1600, 2200
        if img.width > max_w or img.height > max_h:
            img.thumbnail((max_w, max_h))
        has_alpha = img.mode in ("RGBA", "LA")
        if not has_alpha:
            img = img.convert("RGB")
        for q in (88, 70, 50, 35):
            buf = io.BytesIO()
            if has_alpha:
                img.save(buf, format="PNG", optimize=True)
            else:
                img.save(buf, format="JPEG", quality=q)
            if buf.tell() <= max_bytes:
                return buf.getvalue()
        # 保底：最小质量 PNG/JPEG
        buf = io.BytesIO()
        if has_alpha:
            img.save(buf, format="PNG", optimize=True)
        else:
            img.save(buf, format="JPEG", quality=30)
        return buf.getvalue()
    except Exception as e:  # noqa: BLE001
        print(f"[img] 图片压缩失败（用原图）: {e}")
        return blob


def _save_from_response(j: dict, out_path: str) -> str:
    """从同步/异步最终响应提取图片（data[].b64_json 或 url）并保存。"""
    data_item = (j.get("data") or [{}])[0]
    b64 = data_item.get("b64_json")
    if b64:
        Path(out_path).parent.mkdir(parents=True, exist_ok=True)
        Path(out_path).write_bytes(base64.b64decode(b64))
        return out_path
    url = data_item.get("url")
    if not url:
        raise RuntimeError(f"图像 API 返回无 b64_json/url: {str(j)[:200]}")
    img_resp = requests.get(url, proxies=NO_PROXY, timeout=120)
    img_resp.raise_for_status()
    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
    Path(out_path).write_bytes(img_resp.content)
    return out_path


def _wait_task(base_url: str, headers: dict, task_id: str, poll_after_ms: int,
               timeout: int = 60) -> dict:
    """轮询异步图像任务直到完成；uncertain（结果暂时不确定）继续等，不重复提交。
    timeout 默认 60s：异步路径不稳定（ai-media 网关常 uncertain/task not found），
    超时即抛异常由调用方重试同步提交。"""
    deadline = time.time() + timeout
    last = ""
    while time.time() < deadline:
        time.sleep(max(int(poll_after_ms or 2000), 2000) / 1000.0)
        try:
            resp = requests.get(f"{base_url}/images/tasks/{task_id}", headers=headers,
                                timeout=60, proxies=NO_PROXY)
            if resp.status_code >= 400:
                continue
            j = resp.json()
        except Exception as e:  # noqa: BLE001
            print(f"[img] 任务轮询异常（重试）: {e}")
            continue
        st = j.get("status")
        if st != last:
            print(f"[img] 异步任务 {task_id[:20]}… status={st}")
            last = st
        if st in _FINAL_STATUS:
            return j
        if st in _FAIL_STATUS:
            raise RuntimeError(f"图像异步任务失败: {j.get('error')} {j.get('error_code')}")
        # queued / running / uncertain → 继续等
    raise RuntimeError(f"图像异步任务超时（{timeout}s）")


def _resolve_task_or_sync(j: dict, base_url: str, headers: dict, out_path: str) -> str:
    """提交后的统一处理：异步任务则轮询，然后提取图片保存。"""
    if j.get("object") == "image.task" or j.get("task_id") or j.get("id", "").startswith("imgtask"):
        tid = j.get("task_id") or j.get("id") or ""
        if not tid:
            raise RuntimeError(f"异步任务无 task_id: {str(j)[:200]}")
        j = _wait_task(base_url, headers, tid, int(j.get("poll_after_ms") or 2000))
    return _save_from_response(j, out_path)


class OpenAIImageBackend(ImageBackend):
    name = "openai"

    def __init__(self):
        self._cfg: dict = {}

    def bind_config(self, cfg: dict):
        self._cfg = cfg or {}

    def _base_url(self) -> str:
        return str(self._cfg.get("base_url") or "https://api.openai.com/v1").rstrip("/")

    def print(self, prompt: str, base_image: str, out_path: str, negative: str = "",
              extra_images=None, size: str = "") -> str:
        """img2img 编辑：参考图 base_image（+ 可选 extra_images 多参考图）按 prompt 生成新图。

        - 平铺服装图：base_image=平铺衣服底图(图3)，extra_images=[印花设计稿(图2)]
        - 三图模特合成：base_image=模特图(图1)，extra_images=[印花设计稿(图2), 平铺底图(图3)]
          提交顺序即图1→图2→图3，与提示词中的图片角色一一对应。
        size: 显式尺寸覆盖（如 "1504x2000"）；留空用配置 size（默认 1024x1024）。
        """
        cfg = self._cfg
        api_key = cfg.get("api_key", "")
        model = cfg.get("model", "gpt-image-1")
        if not api_key:
            raise RuntimeError("compose.api_key 未配置，无法调用 OpenAI 图像后端。")

        full_prompt = prompt + (f"\nNegative: {negative}" if negative else "")
        base_url = self._base_url()
        headers = {"Authorization": f"Bearer {api_key}"}
        # 必须把文件读成 bytes 再提交（句柄随 with 关闭会导致 "read of closed file"）；
        # 输入图（模特/设计/底图）先压缩校验：>2MB 压缩到 <2MB 再上传，避免大图拖垮网关导致超时
        file_payloads = []
        for img_path in [base_image] + list(extra_images or []):
            blob = Path(img_path).read_bytes()
            if len(blob) > 2 * 1024 * 1024:
                shrunk = _shrink_blob_to_2mb(img_path, blob)
                print(f"[img] 输入图压缩: {Path(img_path).name} {len(blob)//1024}KB → {len(shrunk)//1024}KB（<2MB）")
                blob = shrunk
            file_payloads.append((Path(img_path).name, blob))
        files = [("image", (name, blob, "image/png")) for name, blob in file_payloads]
        data = {
            "prompt": full_prompt,
            "n": 1,
            "size": size or cfg.get("size", "1024x1024"),
            "model": model,
            "execution_mode": "sync",   # 强制同步（ai-media 等网关默认重负载转异步任务，不稳定）
        }
        # 提交重试：异步路径不稳定 → 失败重试同步提交（最多 3 次）；
        # 内容政策拦截（content_policy_violation）多为网关误判 → 等待后重试
        last_err: Optional[str] = None
        for attempt in range(3):
            resp = requests.post(f"{base_url}/images/edits", headers=headers, files=files,
                                 data=data, timeout=300, proxies=NO_PROXY)
            if resp.status_code >= 400:
                body = resp.text or ""
                if "content_policy" in body and attempt < 2:
                    print(f"[img] 内容政策拦截（可能误判），等待后重试 {attempt + 2}/3")
                    time.sleep(2)
                    continue
                if attempt == 0:
                    data.pop("execution_mode", None)   # 网关不认识该参数（官方 OpenAI）→ 去掉重试
                    continue
                raise RuntimeError(f"图像 API {resp.status_code}: {body[:300]}")
            try:
                return _resolve_task_or_sync(resp.json(), base_url, headers, out_path)
            except Exception as e:  # noqa: BLE001
                last_err = str(e)
                print(f"[img] 第 {attempt + 1} 次提交异步失败，重试同步提交: {e}")
        raise RuntimeError(f"图像合成多次提交均失败: {last_err}")

    def generate(self, prompt: str, out_path: str, negative: str = "", size: str = "") -> str:
        """纯文生图：生成白底纯印花设计稿（standalone pure print design）。
        size: 显式尺寸覆盖（印花设计统一 1024x1024）；留空用配置 size。
        background: 配置 compose.background="transparent" 时传 background 参数 → 透明背景 PNG
        （gpt-image-1/2 等模型支持；网关不支持该参数时会被忽略或由网关兜底）。"""
        cfg = self._cfg
        api_key = cfg.get("api_key", "")
        model = cfg.get("model", "gpt-image-1")
        if not api_key:
            raise RuntimeError("compose.api_key 未配置，无法调用 OpenAI 图像后端。")

        full_prompt = prompt + (f"\nNegative: {negative}" if negative else "")
        base_url = self._base_url()
        headers = {"Authorization": f"Bearer {api_key}"}
        data = {
            "prompt": full_prompt,
            "n": 1,
            "size": size or cfg.get("size", "1024x1024"),
            "model": model,
            "response_format": "b64_json",
            "execution_mode": "sync",   # 强制同步（ai-media 等网关默认重负载转异步任务，不稳定）
        }
        bg = str(cfg.get("background") or "").strip()
        if bg:
            data["background"] = bg   # 如 "transparent"（透明背景 PNG）
        last_err: Optional[str] = None
        for attempt in range(3):
            resp = requests.post(f"{base_url}/images/generations", headers=headers, json=data,
                                 timeout=300, proxies=NO_PROXY)
            if resp.status_code >= 400:
                body = resp.text or ""
                if "content_policy" in body and attempt < 2:
                    print(f"[img] 内容政策拦截（可能误判），等待后重试 {attempt + 2}/3")
                    time.sleep(2)
                    continue
                if attempt == 0:
                    data.pop("execution_mode", None)   # 网关不认识该参数（官方 OpenAI）→ 去掉重试
                    continue
                raise RuntimeError(f"图像 API {resp.status_code}: {body[:300]}")
            try:
                return _resolve_task_or_sync(resp.json(), base_url, headers, out_path)
            except Exception as e:  # noqa: BLE001
                last_err = str(e)
                print(f"[img] 第 {attempt + 1} 次生成异步失败，重试同步提交: {e}")
        raise RuntimeError(f"图像生成多次提交均失败: {last_err}")
