"""
GPT Image API 客户端
新版节点请求走异步任务接口：
  - POST /async/v1/generateImage
  - GET  /async/v1/tasks/{task_id}

旧同步接口保留兼容代码，但 GPT Image 节点不再使用：
  - POST /v1/images/generations/
  - POST /v1/images/edits

设计原则：
  - 对 ComfyUI 节点暴露同步入口，内部提交异步任务并轮询
  - 图片和蒙版以 data:image/png;base64,... 放入 JSON 请求体
  - 响应优先读取 data.images[].url
"""

import asyncio
import base64
import binascii
import json
import math
import os
import time
from concurrent.futures import ThreadPoolExecutor
from io import BytesIO
from typing import Any, Callable, Dict, List, Optional

import aiohttp
import numpy as np
import torch
from PIL import Image, UnidentifiedImageError

from ..utils.config import get_api_key_or_raise, get_api_base_url
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
from ..utils.http_error import HTTP_ERROR_MESSAGES, get_friendly_message
from ..utils.http2_client import (
    format_response_body_diagnostics,
    read_response_body_with_diagnostics,
    response_task_id,
    validate_response_task_id,
)
from ..utils.o1key_image_catalog import (
    GPT_IMAGE_BACKGROUND_OPTIONS,
    GPT_IMAGE_MODEL_OPTIONS,
    GPT_IMAGE_OUTPUT_FORMAT_OPTIONS,
    UNIFIED_IMAGE_ROUTE_OPTIONS,
)

try:
    from comfy.model_management import processing_interrupted, InterruptProcessingException
    _INTERRUPT_AVAILABLE = True
except ImportError:
    _INTERRUPT_AVAILABLE = False
    InterruptProcessingException = RuntimeError
    processing_interrupted = lambda: False

# ── 接口端点 ──────────────────────────────────────────────────────────────────
_ENDPOINT_GENERATIONS = "/v1/images/generations/"
_ENDPOINT_EDITS       = "/v1/images/edits"
_ENDPOINT_ASYNC_GENERATE = "/async/v1/generateImage"
_ENDPOINT_ASYNC_TASK = "/async/v1/tasks/{task_id}"

# ── 主模型与模型线路映射（界面值 → API 模型名）──────────────────────────────
GPT_IMAGE_ROUTE_OPTIONS = UNIFIED_IMAGE_ROUTE_OPTIONS
GPT_IMAGE_MODEL_MATRIX = {
    ("gpt-image-2", "畅速"): "gpt-image-2-c-sp",
    ("gpt-image-2", "直连"): "gpt-image-2-c-sd",
    ("gpt-image-2", "专线"): "gpt-image-2",
    ("gpt-image-2.5-sunburst", "畅速"): "gpt-image-2.5-sunburst-sp",
    ("gpt-image-2.5-sunburst", "直连"): "gpt-image-2.5-sunburst-sd",
    ("gpt-image-2.5-sunburst", "专线"): "gpt-image-2.5-sunburst",
    ("gpt-image-2.5-flare", "畅速"): "gpt-image-2.5-flare-sp",
    ("gpt-image-2.5-flare", "直连"): "gpt-image-2.5-flare-sd",
    ("gpt-image-2.5-flare", "专线"): "gpt-image-2.5-flare",
}
_MODEL_NAME_MAP = {
    "畅速": "gpt-image-2-c-sp",
    "直连": "gpt-image-2-c-sd",
    "专线": "gpt-image-2",
    # 兼容旧工作流和直接调用节点的旧参数值。
    "gpt-image-2-特价": "gpt-image-2-c-sp",
    "gpt-image-2-官方": "gpt-image-2",
    "gpt-image-2-次卡": "gpt-image-2-c-sp",
    "gpt-image-2-按量": "gpt-image-2",
}
def resolve_gpt_image_model(model_name: str, route: str) -> str:
    """由主模型和线路解析实际 API 模型名，并兼容旧的单参数调用。"""
    # 旧节点曾把线路或旧显示模型名直接作为“模型”参数传入。
    if model_name in GPT_IMAGE_ROUTE_OPTIONS or model_name in _MODEL_NAME_MAP:
        return _MODEL_NAME_MAP.get(model_name, model_name)
    model = GPT_IMAGE_MODEL_MATRIX.get((model_name, route))
    # 兼容外部调用把旧线路显示值放在 route 参数中的情况。
    if model is None and route in _MODEL_NAME_MAP:
        model = _MODEL_NAME_MAP[route]
    if model is None:
        raise ValueError(f"模型 '{model_name}' 不支持模型线路 '{route}'")
    return model

# ── 超时 ──────────────────────────────────────────────────────────────────────
_REQUEST_TIMEOUT = 900   # 秒
_ASYNC_POLL_SCHEDULE = [5.0, 20.0]
_ASYNC_POLL_INTERVAL = 3.0


def _attach_original_image_info(image: Image.Image, data: bytes) -> Image.Image:
    """Keep provider bytes alongside decoded pixels for format-preserving saves."""
    image_format = str(image.format or "").upper()
    if image_format == "JPG":
        image_format = "JPEG"
    if image_format:
        image.format = image_format
        setattr(image, "_o1key_original_format", image_format)
    setattr(image, "_o1key_original_bytes", data)
    return image
_ASYNC_MAX_WAIT = 600.0
# Only use these after a task has been submitted: status/result GET requests
# are idempotent, while retrying task submission can create duplicate charges.
_RESPONSE_READ_MAX_RETRIES = 3
_RESPONSE_READ_BASE_DELAY = 1.0
_RESPONSE_READ_MAX_DELAY = 8.0
_RETRYABLE_RESULT_HTTP_STATUSES = {408, 425, 429, 500, 502, 503, 504}

REQUEST_LOG_ENABLED = False
POLL_LOG_ENABLED = False
RESPONSE_LOG_ENABLED = os.environ.get("O1KEY_RESPONSE_LOG", "0").strip().lower() not in {
    "0", "false", "no", "off"
}


class _IncompleteInlineImageError(RuntimeError):
    """The task succeeded, but an inline result image is incomplete."""


class _RetryableResultHTTPError(OSError):
    """An idempotent task query or result download returned a transient status."""


class GptImageClient:
    """
    GPT Image API 客户端

    接口说明：
      async generateImage：JSON 提交，返回 task_id
      tasks/{task_id}：轮询任务状态，成功后读取 data.images[].url

    旧 generations / edits 同步接口保留为兼容代码。
    """

    def __init__(self):
        self.api_key  = get_api_key_or_raise("O1KEY_API_KEY")
        self.base_url = get_api_base_url()
        # Batch nodes can turn these off while keeping the same client API for
        # single-image diagnostics.
        self.response_log_enabled = RESPONSE_LOG_ENABLED
        self.poll_log_enabled = POLL_LOG_ENABLED

    @staticmethod
    def _response_retry_delay(attempt: int) -> float:
        return min(
            _RESPONSE_READ_BASE_DELAY * (2 ** attempt),
            _RESPONSE_READ_MAX_DELAY,
        )

    async def _get_json_with_response_retry(
        self,
        session: aiohttp.ClientSession,
        url: str,
        headers: dict,
        label: str,
        expected_task_id: Optional[str] = None,
    ) -> dict:
        """Read a JSON GET response, retrying interrupted response bodies."""
        for attempt in range(_RESPONSE_READ_MAX_RETRIES + 1):
            try:
                async with session.get(url, headers=headers) as resp:
                    status = resp.status
                    body, diagnostics = await read_response_body_with_diagnostics(resp)
                    if self.response_log_enabled:
                        self._log_original_response_body(
                            f"{label} status={status}",
                            body.decode("utf-8", errors="replace"),
                        )
                    if status != 200:
                        print(
                            f"[o1key GPT Image] 任务查询传输追踪"
                            f" | requested_task_id={expected_task_id or '<unknown>'}"
                            f" | response_task_id=<not-parsed>"
                            f" | {format_response_body_diagnostics(diagnostics)}"
                            f" | json=not-parsed"
                        )
                        text = body.decode("utf-8", errors="replace")
                        message = self._extract_error_message(text, status)
                        if status in _RETRYABLE_RESULT_HTTP_STATUSES:
                            raise _RetryableResultHTTPError(message)
                        raise RuntimeError(message)
                    try:
                        payload = json.loads(body)
                    except json.JSONDecodeError:
                        print(
                            f"[o1key GPT Image] 任务查询传输追踪"
                            f" | requested_task_id={expected_task_id or '<unknown>'}"
                            f" | response_task_id=<unavailable>"
                            f" | {format_response_body_diagnostics(diagnostics)}"
                            f" | json=invalid"
                        )
                        raise

                    actual_task_id = response_task_id(payload)
                    if expected_task_id is not None:
                        if actual_task_id is not None and actual_task_id != expected_task_id:
                            print(
                                f"[o1key GPT Image] 任务查询传输追踪"
                                f" | requested_task_id={expected_task_id}"
                                f" | response_task_id={actual_task_id}"
                                f" | task_id_check=mismatch"
                                f" | {format_response_body_diagnostics(diagnostics)}"
                                f" | json=valid"
                            )
                        validate_response_task_id(payload, expected_task_id)
                    return payload
            except (
                aiohttp.ClientError,
                asyncio.TimeoutError,
                OSError,
                json.JSONDecodeError,
            ) as exc:
                if attempt >= _RESPONSE_READ_MAX_RETRIES:
                    raise RuntimeError(
                        f"{label}响应读取失败，已重试 "
                        f"{_RESPONSE_READ_MAX_RETRIES} 次：{exc}"
                    ) from None

                delay = self._response_retry_delay(attempt)
                print(
                    f"[o1key GPT Image] {label}响应不完整 "
                    f"({type(exc).__name__})，{delay:.1f}s 后重试 "
                    f"({attempt + 1}/{_RESPONSE_READ_MAX_RETRIES})"
                    f" | requested_task_id={expected_task_id or '<unknown>'} | {exc}"
                )
                await asyncio.sleep(delay)

        raise RuntimeError(f"{label}响应读取失败")  # pragma: no cover

    async def _download_image_with_response_retry(
        self,
        session: aiohttp.ClientSession,
        url: str,
        label: str,
    ) -> tuple[Image.Image, int, float]:
        """Download and fully decode an image, retrying interrupted GET reads."""
        download_started = time.perf_counter()
        for attempt in range(_RESPONSE_READ_MAX_RETRIES + 1):
            try:
                async with session.get(url, allow_redirects=True) as resp:
                    image_bytes, diagnostics = await read_response_body_with_diagnostics(resp)
                    if resp.status != 200:
                        print(
                            f"[o1key GPT Image] 结果下载传输追踪"
                            f" | {format_response_body_diagnostics(diagnostics)}"
                        )
                        message = f"{label}下载失败 HTTP {resp.status}"
                        if resp.status in _RETRYABLE_RESULT_HTTP_STATUSES:
                            raise _RetryableResultHTTPError(message)
                        raise RuntimeError(message)

                image = Image.open(BytesIO(image_bytes))
                image.load()
                _attach_original_image_info(image, image_bytes)
                return image, len(image_bytes), time.perf_counter() - download_started
            except (
                aiohttp.ClientError,
                asyncio.TimeoutError,
                OSError,
                UnidentifiedImageError,
            ) as exc:
                if attempt >= _RESPONSE_READ_MAX_RETRIES:
                    raise RuntimeError(
                        f"{label}下载失败，已重试 "
                        f"{_RESPONSE_READ_MAX_RETRIES} 次：{exc}"
                    ) from None

                delay = self._response_retry_delay(attempt)
                print(
                    f"[o1key GPT Image] {label}下载中断 "
                    f"({type(exc).__name__})，{delay:.1f}s 后重试 "
                    f"({attempt + 1}/{_RESPONSE_READ_MAX_RETRIES})..."
                )
                await asyncio.sleep(delay)

        raise RuntimeError(f"{label}下载失败")  # pragma: no cover

    @staticmethod
    def _format_transfer_size(size: int) -> str:
        if size < 1024 * 1024:
            return f"{size / 1024:.1f} KiB"
        return f"{size / 1024 / 1024:.2f} MiB"

    @staticmethod
    def _format_transfer_rate(size: int, elapsed: float) -> str:
        if elapsed <= 0:
            return "∞ MiB/s"
        return f"{size / 1024 / 1024 / elapsed:.2f} MiB/s"

    # ── 认证头 ────────────────────────────────────────────────────────────────

    def _auth_headers(self) -> dict:
        return {"Authorization": f"Bearer {self.api_key}"}

    def _json_headers(self) -> dict:
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

    @staticmethod
    def _new_multipart_form() -> aiohttp.FormData:
        try:
            return aiohttp.FormData(default_to_multipart=True)
        except TypeError:
            form = aiohttp.FormData()
            form._is_multipart = True
            return form

    @staticmethod
    def _add_form_fields(form: aiohttp.FormData, fields: dict) -> None:
        for key, value in fields.items():
            if value is None:
                continue
            if isinstance(value, bool):
                value = "true" if value else "false"
            elif isinstance(value, (dict, list)):
                value = json.dumps(value, ensure_ascii=False)
            form.add_field(key, str(value))

    # ── 图像转换工具 ──────────────────────────────────────────────────────────

    # ── 请求体大小限制 ────────────────────────────────────────────────────────
    _MAX_BODY_BYTES = 20 * 1024 * 1024   # 20 MB
    _UNIFIED_BODY_LIMIT_BYTES = 18 * 1024 * 1024
    _SMART_RESIZE_MIN_LONG_EDGE = 256
    _SMART_RESIZE_MAX_ATTEMPTS = 32

    @staticmethod
    def _shrink_png_to_limit(png_bytes: bytes, max_bytes: int, label: str = "") -> bytes:
        """
        若 PNG bytes 超过 max_bytes，按等比缩放反复压缩直到满足限制。
        每次将面积缩小至约 80%（线性尺寸缩小至约 89.4%）。
        """
        if len(png_bytes) <= max_bytes:
            return png_bytes

        img = Image.open(BytesIO(png_bytes))
        w, h = img.size
        original_size = len(png_bytes)
        step = 0

        while len(png_bytes) > max_bytes:
            scale = 0.894   # sqrt(0.8)，面积缩小 20%
            w = max(1, int(w * scale))
            h = max(1, int(h * scale))
            img = img.resize((w, h), Image.LANCZOS)
            buf = BytesIO()
            img.save(buf, format="PNG")
            png_bytes = buf.getvalue()
            step += 1

        tag = f" ({label})" if label else ""
        print(
            f"[o1key GPT Image] 图像{tag}超出 {max_bytes // (1024*1024)}MB 限制，"
            f"已等比缩放 {step} 次：{original_size // 1024}KB → {len(png_bytes) // 1024}KB "
            f"（{w}×{h}）"
        )
        return png_bytes

    @staticmethod
    def _png_bytes_to_data_url(png_bytes: bytes) -> str:
        b64 = base64.b64encode(png_bytes).decode("ascii")
        return f"data:image/png;base64,{b64}"

    @staticmethod
    def _json_body_size(body: dict) -> int:
        return len(json.dumps(body, ensure_ascii=False, separators=(",", ":")).encode("utf-8"))

    @staticmethod
    def _format_body_size(size: int) -> str:
        if size < 1024 * 1024:
            return f"{size / 1024:.2f}KB"
        return f"{size / 1024 / 1024:.2f}MB"

    @staticmethod
    def _shorten_data_urls_for_log(obj, max_len: int = 200):
        if isinstance(obj, dict):
            return {
                key: GptImageClient._shorten_data_urls_for_log(value, max_len)
                for key, value in obj.items()
            }
        if isinstance(obj, list):
            return [GptImageClient._shorten_data_urls_for_log(item, max_len) for item in obj]
        if isinstance(obj, str) and len(obj) > max_len:
            header, separator, data = obj.partition(",")
            if separator and header.lower().startswith("data:") and ";base64" in header.lower():
                return f"{header},<base64 data, {len(data)} chars>"
            candidate = obj.strip()
            if len(candidate) % 4 == 0:
                try:
                    base64.b64decode(candidate, validate=True)
                except Exception:
                    pass
                else:
                    return f"<base64 data, {len(candidate)} chars>"
        return obj

    def _log_original_request_body(self, label: str, body: dict) -> None:
        if not REQUEST_LOG_ENABLED:
            return
        body_size = self._json_body_size(body)
        print(
            f"\n{'=' * 60}\n"
            f"[o1key GPT Image] 原始请求体日志 | {label} | "
            f"请求体积: {self._format_body_size(body_size)} "
            f"(data URL 已折叠显示 base64 长度)\n"
            f"{json.dumps(self._shorten_data_urls_for_log(body), ensure_ascii=False, indent=2)}\n"
            f"{'=' * 60}\n"
        )

    def _log_original_response_body(self, label: str, text: str) -> None:
        if not self.response_log_enabled:
            return
        size = len(text.encode("utf-8"))
        try:
            safe_body = json.dumps(
                self._shorten_data_urls_for_log(json.loads(text)),
                ensure_ascii=False,
                indent=2,
            )
        except Exception:
            safe_body = f"<unparseable response body, {size} bytes; content omitted>"
        print(
            f"\n{'=' * 60}\n"
            f"[o1key GPT Image] 原始返回响应体日志 | {label} | "
            f"响应体积: {self._format_body_size(size)} (base64 已折叠)\n"
            f"{safe_body}\n"
            f"{'=' * 60}\n"
        )

    @staticmethod
    def _extract_error_message(payload_or_text, status_code: int = 0) -> str:
        payload = payload_or_text
        if isinstance(payload_or_text, str):
            try:
                payload = json.loads(payload_or_text)
            except Exception:
                return get_friendly_message(status_code, payload_or_text)

        if isinstance(payload, dict):
            error = payload.get("error")
            if isinstance(error, str) and error.strip():
                return get_friendly_message(0, error.strip())
            if isinstance(error, dict):
                msg = error.get("message") or error.get("msg") or error.get("error")
                if msg:
                    return get_friendly_message(0, str(msg))
                return json.dumps(error, ensure_ascii=False)

            msg = payload.get("message") or payload.get("msg")
            if msg:
                return get_friendly_message(0, str(msg))

        return get_friendly_message(status_code, str(payload_or_text))

    @staticmethod
    def _extract_error_detail(payload: dict) -> dict:
        if not isinstance(payload, dict):
            return {}
        detail = payload.get("error_detail")
        return detail if isinstance(detail, dict) else {}

    @staticmethod
    def _coerce_progress_percent(value: Any) -> Optional[int]:
        if value is None:
            return None
        if isinstance(value, str):
            text = value.strip().rstrip("%")
            if not text:
                return None
            try:
                value = float(text)
            except ValueError:
                return None
        elif isinstance(value, (int, float)):
            value = float(value)
        else:
            return None

        if 0 <= value <= 1:
            value *= 100
        return max(0, min(100, int(round(value))))

    @staticmethod
    def _emit_progress(progress_callback: Optional[Callable[[int], None]], pct: int) -> None:
        if progress_callback is None:
            return
        try:
            progress_callback(pct)
        except Exception as error:
            print(f"[o1key GPT Image] progress callback failed: {error}")

    @staticmethod
    def _resize_png_bytes(source_image: Image.Image, scale: float) -> bytes:
        if scale < 0.999:
            width, height = source_image.size
            new_width = max(1, round(width * scale))
            new_height = max(1, round(height * scale))
            image = source_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        else:
            image = source_image

        buf = BytesIO()
        image.save(buf, format="PNG")
        return buf.getvalue()

    @staticmethod
    def _pil_to_png_bytes(image: Image.Image) -> bytes:
        buf = BytesIO()
        image.save(buf, format="PNG")
        return buf.getvalue()

    def _fit_png_assets_to_body_limit(
        self,
        assets: List[Dict[str, Any]],
        build_body,
        *,
        resize_mode: Optional[str] = None,
    ) -> Dict[str, bytes]:
        """Fit exact JSON body size while preserving unified-node resize semantics.

        ``resize_mode=None`` preserves the legacy GPT client behavior: automatically
        fit to the provider's 20 MiB limit.  The unified image node passes an
        explicit mode and therefore uses the same 18 MiB safety ceiling as Nano
        Banana.  Smart mode resizes the largest encoded asset group first and
        always resamples from the original pixels.  A reference image and its mask
        share a group so their dimensions remain aligned.
        """
        if resize_mode not in {None, "不缩放", "智能缩放"}:
            raise ValueError("缩放图片参数无效")

        body_limit = (
            self._MAX_BODY_BYTES
            if resize_mode is None
            else self._UNIFIED_BODY_LIMIT_BYTES
        )
        effective_mode = "智能缩放" if resize_mode is None else resize_mode
        asset_bytes = {asset["key"]: asset["bytes"] for asset in assets}
        initial_size = self._json_body_size(build_body(asset_bytes))
        if initial_size <= body_limit:
            return asset_bytes

        limit_mib = body_limit / (1024 * 1024)
        if effective_mode == "不缩放":
            raise ValueError(
                f"GPT Image: 请求体 {initial_size / (1024 * 1024):.2f} MiB 超过 "
                f"{limit_mib:.0f} MiB 上限；当前设置为“不缩放”，"
                "请在上游缩小图片或改用“智能缩放”"
            )
        if not assets:
            raise ValueError(
                f"GPT Image: 请求体 {initial_size / (1024 * 1024):.2f} MiB 超过 "
                f"{limit_mib:.0f} MiB 上限，且没有可缩放图片"
            )

        originals: Dict[str, Image.Image] = {}
        groups: Dict[str, List[str]] = {}
        labels: Dict[str, str] = {}
        try:
            for asset in assets:
                with Image.open(BytesIO(asset["bytes"])) as image:
                    image.load()
                    originals[asset["key"]] = image.copy()
                group = str(asset.get("resize_group") or asset["key"])
                groups.setdefault(group, []).append(asset["key"])
                labels[asset["key"]] = str(asset.get("label") or asset["key"])

            scales = {group: 1.0 for group in groups}
            final_sizes = {key: image.size for key, image in originals.items()}
            body_size = initial_size

            for _attempt in range(self._SMART_RESIZE_MAX_ATTEMPTS):
                if body_size <= body_limit:
                    changes = [
                        f"{labels[key]} {originals[key].width}×{originals[key].height}"
                        f"→{final_sizes[key][0]}×{final_sizes[key][1]}"
                        for key in originals
                        if final_sizes[key] != originals[key].size
                    ]
                    print(
                        f"[o1key GPT Image] 智能缩放完成 | "
                        f"{initial_size / (1024 * 1024):.2f}→"
                        f"{body_size / (1024 * 1024):.2f} MiB | "
                        + "，".join(changes)
                    )
                    return asset_bytes

                candidates = []
                for group, keys in groups.items():
                    min_scale = max(
                        min(1.0, self._SMART_RESIZE_MIN_LONG_EDGE / max(originals[key].size))
                        for key in keys
                    )
                    current_scale = scales[group]
                    if current_scale <= min_scale + 1e-6:
                        continue
                    encoded_length = sum(
                        4 * ((len(asset_bytes[key]) + 2) // 3)
                        for key in keys
                    )
                    candidates.append((encoded_length, group, min_scale))
                if not candidates:
                    break

                current_data_length, group, min_scale = max(candidates)
                overflow = body_size - body_limit
                desired_data_length = max(1024, current_data_length - overflow - 32 * 1024)
                estimated_area_ratio = min(0.90, desired_data_length / current_data_length)
                proposed_scale = scales[group] * math.sqrt(max(0.01, estimated_area_ratio)) * 0.98
                new_scale = max(min_scale, min(scales[group] * 0.92, proposed_scale))

                candidate_bytes = dict(asset_bytes)
                candidate_sizes = dict(final_sizes)
                for key in groups[group]:
                    candidate_bytes[key] = self._resize_png_bytes(originals[key], new_scale)
                    with Image.open(BytesIO(candidate_bytes[key])) as candidate_image:
                        candidate_sizes[key] = candidate_image.size

                if all(candidate_sizes[key] == final_sizes[key] for key in groups[group]):
                    new_scale = max(min_scale, scales[group] * 0.85)
                    for key in groups[group]:
                        candidate_bytes[key] = self._resize_png_bytes(originals[key], new_scale)
                        with Image.open(BytesIO(candidate_bytes[key])) as candidate_image:
                            candidate_sizes[key] = candidate_image.size

                asset_bytes = candidate_bytes
                final_sizes = candidate_sizes
                scales[group] = new_scale
                body_size = self._json_body_size(build_body(asset_bytes))

            final_size = self._json_body_size(build_body(asset_bytes))
            raise ValueError(
                f"GPT Image: 智能缩放后请求体仍为 {final_size / (1024 * 1024):.2f} MiB，"
                f"无法安全降至 {limit_mib:.0f} MiB；请减少参考图或在上游进一步缩小"
            )
        finally:
            for image in originals.values():
                image.close()

    @staticmethod
    def _tensor_to_png_bytes(tensor: torch.Tensor) -> bytes:
        """
        单张 ComfyUI IMAGE tensor [1, H, W, C] 或 [H, W, C] → PNG bytes
        """
        if tensor.dim() == 4:
            tensor = tensor.squeeze(0)              # [H, W, C]
        arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
        img = Image.fromarray(arr)
        buf = BytesIO()
        img.save(buf, format="PNG")
        return buf.getvalue()

    @staticmethod
    def _mask_tensor_to_rgba_png_bytes(mask: torch.Tensor, image_size: tuple) -> bytes:
        """
        ComfyUI MASK tensor [1, H, W] 或 [H, W] → RGBA PNG bytes
        白色区域（mask=1）→ 透明（alpha=0），即 API 将在此处生成新内容。
        """
        if mask.dim() == 3:
            mask = mask.squeeze(0)                  # [H, W]

        h, w = mask.shape
        ih, iw = image_size

        # 尺寸不一致时给出提示（API 侧也会报错）
        if (h, w) != (ih, iw):
            raise ValueError(
                f"蒙版尺寸 ({h}×{w}) 与图像尺寸 ({ih}×{iw}) 不一致，请保持相同尺寸"
            )

        alpha = ((1.0 - mask.cpu().numpy()) * 255).clip(0, 255).astype(np.uint8)
        rgba  = np.zeros((h, w, 4), dtype=np.uint8)
        rgba[:, :, 3] = alpha                       # 只设 alpha，RGB 全 0

        buf = BytesIO()
        Image.fromarray(rgba, mode="RGBA").save(buf, format="PNG")
        return buf.getvalue()

    @staticmethod
    def _pil_list_to_tensor(images: List[Image.Image]) -> torch.Tensor:
        """
        PIL Image 列表 → ComfyUI IMAGE tensor [B, H, W, C]，值域 [0, 1]
        RGBA 自动转换为 RGBA（保留透明通道）
        """
        if not images:
            placeholder = Image.new("RGBA", (512, 512), (128, 128, 128, 255))
            images = [placeholder]

        tensors = []
        source_metadata = []
        for img in images:
            source_metadata.append({
                "format": getattr(img, "_o1key_original_format", None) or img.format,
                "bytes": getattr(img, "_o1key_original_bytes", None),
                "modified": False,
            })
            arr = np.array(img.convert("RGBA")).astype(np.float32) / 255.0
            tensors.append(torch.from_numpy(arr))

        # 批量模式下 API 可能返回不同尺寸，统一 resize 到最大尺寸
        max_h = max(t.shape[0] for t in tensors)
        max_w = max(t.shape[1] for t in tensors)
        aligned = []
        for index, t in enumerate(tensors):
            if t.shape[0] != max_h or t.shape[1] != max_w:
                source_metadata[index]["modified"] = True
                t = t.permute(2, 0, 1).unsqueeze(0)  # [1, C, H, W]
                t = torch.nn.functional.interpolate(
                    t, size=(max_h, max_w), mode="bilinear", align_corners=False
                )
                t = t.squeeze(0).permute(1, 2, 0)  # [H, W, C]
            aligned.append(t)

        result = torch.stack(aligned, dim=0)        # [B, H, W, 4]
        result._o1key_source_metadata = source_metadata
        return result

    # ── 响应解析（通用） ─────────────────────────────────────────────────────

    async def _parse_response(
        self,
        resp_json: dict,
        session: aiohttp.ClientSession,
    ) -> List[Image.Image]:
        """
        解析 data 列表，优先取 b64_json，回退到 url 下载
        """
        if "error" in resp_json:
            err = resp_json["error"]
            msg = (
                err.get("message") or err.get("msg") or json.dumps(err, ensure_ascii=False)
                if isinstance(err, dict)
                else str(err)
            )
            raise RuntimeError(f"API 返回错误: {msg}")

        data_list = resp_json.get("data")
        if data_list is None:
            data_list = resp_json.get("images")
        if data_list is None and (resp_json.get("b64_json") or resp_json.get("url")):
            data_list = [resp_json]
        if not data_list:
            raise RuntimeError(
                f"API 响应中未找到 data 字段，完整响应：\n"
                f"{json.dumps(resp_json, ensure_ascii=False, indent=2)}"
            )

        images: List[Image.Image] = []
        for idx, item in enumerate(data_list):
            b64 = item.get("b64_json", "")
            url = item.get("url", "")

            if b64:
                # 优先 base64（无需二次下载）
                try:
                    img_bytes = base64.b64decode(b64)
                    img = Image.open(BytesIO(img_bytes))
                    img.load()
                    _attach_original_image_info(img, img_bytes)
                    images.append(img)
                    print(f"[o1key GPT Image] 第 {idx + 1} 张 base64 解码完成 "
                          f"({img.size[0]}×{img.size[1]})")
                except Exception as e:
                    raise RuntimeError(f"第 {idx + 1} 张 base64 解码失败: {e}")

            elif url and url.startswith("http"):
                # 回退：下载 URL
                img, byte_count, elapsed = await self._download_image_with_response_retry(
                    session, url, f"第 {idx + 1} 张图片"
                )
                images.append(img)
                print(
                    f"[o1key GPT Image] 下载图像 {idx + 1}/{len(data_list)} 完成 | "
                    f"{img.size[0]}×{img.size[1]} | {self._format_transfer_size(byte_count)} | "
                    f"{elapsed:.2f}s | {self._format_transfer_rate(byte_count, elapsed)}"
                )
            else:
                print(f"[o1key GPT Image] 警告：第 {idx + 1} 条数据既无 b64_json 也无 url，已跳过")

        return images

    @staticmethod
    def _decode_b64_image(b64: str, label: str) -> Image.Image:
        try:
            img_bytes = base64.b64decode(b64)
            img = Image.open(BytesIO(img_bytes))
            img.load()
            _attach_original_image_info(img, img_bytes)
            print(f"[o1key GPT Image] {label} base64 解码完成 ({img.size[0]}×{img.size[1]})")
            return img
        except Exception as e:
            raise RuntimeError(f"{label} base64 解码失败: {e}") from None

    @staticmethod
    def _decode_inline_result_image(b64: str, label: str) -> tuple[Image.Image, int]:
        """Strictly decode and fully load a completed task's inline image."""
        try:
            if not isinstance(b64, str) or not b64:
                raise ValueError("empty base64 payload")
            image_bytes = base64.b64decode(b64, validate=True)
            if not image_bytes:
                raise ValueError("decoded image is empty")
            image = Image.open(BytesIO(image_bytes))
            image.load()
            _attach_original_image_info(image, image_bytes)
            return image, len(image_bytes)
        except (binascii.Error, ValueError, OSError, UnidentifiedImageError) as exc:
            raise _IncompleteInlineImageError(f"{label}不完整或无法解码: {exc}") from None

    async def _append_images_from_payload(
        self,
        payload: dict,
        session: aiohttp.ClientSession,
        images: List[Image.Image],
        event_name: str = "",
    ) -> bool:
        if isinstance(payload, dict) and "error" in payload:
            err = payload["error"]
            msg = (
                err.get("message") or err.get("msg") or json.dumps(err, ensure_ascii=False)
                if isinstance(err, dict)
                else str(err)
            )
            raise RuntimeError(get_friendly_message(500, msg)) from None

        if not isinstance(payload, dict):
            return False

        event_type = payload.get("type") or event_name
        if "partial_image" in event_type:
            return False

        for key in ("data", "images"):
            data_list = payload.get(key)
            if isinstance(data_list, list):
                parsed = await self._parse_response({"data": data_list}, session)
                images.extend(parsed)
                return True

        if payload.get("b64_json") or payload.get("url"):
            parsed = await self._parse_response({"data": [payload]}, session)
            images.extend(parsed)
            return True

        image_obj = payload.get("image")
        if isinstance(image_obj, dict) and (image_obj.get("b64_json") or image_obj.get("url")):
            parsed = await self._parse_response({"data": [image_obj]}, session)
            images.extend(parsed)
            return True

        return False

    async def _parse_edit_stream_response(
        self,
        resp: aiohttp.ClientResponse,
        session: aiohttp.ClientSession,
    ) -> List[Image.Image]:
        """
        Parse /v1/images/edits SSE events and return final completed images.
        Partial images are intentionally ignored so the node output stays unchanged.
        """
        images: List[Image.Image] = []
        buffer = ""
        event_name = ""
        data_lines = []
        partial_count = 0
        debug_body_parts = []

        async def _handle_event():
            nonlocal event_name, data_lines, partial_count, images
            if not data_lines:
                event_name = ""
                return

            data_str = "\n".join(data_lines).strip()
            event_name = event_name.strip()
            data_lines = []

            if not data_str or data_str == "[DONE]":
                return

            try:
                payload = json.loads(data_str)
            except Exception:
                raise RuntimeError(get_friendly_message(500, data_str)) from None

            if isinstance(payload, dict):
                event_type = payload.get("type") or event_name
            else:
                event_type = event_name
            if "partial_image" in event_type:
                partial_count += 1
                return

            await self._append_images_from_payload(payload, session, images, event_name)

        async for raw_chunk in resp.content.iter_any():
            chunk_text = raw_chunk.decode("utf-8", errors="ignore")
            debug_body_parts.append(chunk_text)
            buffer += chunk_text
            while "\n" in buffer:
                line, buffer = buffer.split("\n", 1)
                line = line.rstrip("\r")
                if line == "":
                    await _handle_event()
                    event_name = ""
                    continue
                if line.startswith(":"):
                    continue
                if line.startswith("event:"):
                    event_name = line[len("event:"):].strip()
                elif line.startswith("data:"):
                    data_lines.append(line[len("data:"):].lstrip())

        if buffer.strip():
            data_lines.append(buffer.strip())
        await _handle_event()

        if partial_count:
            print(f"[o1key GPT Image] 流式中间图 {partial_count} 张（已忽略，仅输出最终图）")
        if not images:
            raise RuntimeError("流式响应结束，但未收到最终图片")

        return images

    async def _parse_stream_text_response(
        self,
        text: str,
        session: aiohttp.ClientSession,
    ) -> List[Image.Image]:
        images: List[Image.Image] = []
        partial_count = 0
        event_name = ""
        data_lines = []

        async def _handle_event():
            nonlocal event_name, data_lines, partial_count, images
            if not data_lines:
                event_name = ""
                return
            data_str = "\n".join(data_lines).strip()
            event_name = event_name.strip()
            data_lines = []
            if not data_str or data_str == "[DONE]":
                return
            payload = json.loads(data_str)
            if isinstance(payload, dict):
                event_type = payload.get("type") or event_name
            else:
                event_type = event_name
            if "partial_image" in event_type:
                partial_count += 1
                return
            await self._append_images_from_payload(payload, session, images, event_name)

        for raw_line in text.splitlines():
            line = raw_line.rstrip("\r")
            if line == "":
                await _handle_event()
                event_name = ""
                continue
            if line.startswith(":"):
                continue
            if line.startswith("event:"):
                event_name = line[len("event:"):].strip()
            elif line.startswith("data:"):
                data_lines.append(line[len("data:"):].lstrip())
        await _handle_event()

        if partial_count:
            print(f"[o1key GPT Image] 流式中间图 {partial_count} 张（已忽略，仅输出最终图）")
        if not images:
            raise RuntimeError("流式响应结束，但未收到最终图片")
        return images

    async def _parse_success_response(
        self,
        resp: aiohttp.ClientResponse,
        session: aiohttp.ClientSession,
        label: str = "",
    ) -> List[Image.Image]:
        content_type = resp.headers.get("Content-Type", "").lower()
        if "event-stream" in content_type:
            return await self._parse_edit_stream_response(resp, session)

        text = await resp.text()
        self._log_original_response_body(f"{label or 'sync'} status={resp.status}", text)
        stripped = text.lstrip()
        if stripped.startswith("data:") or stripped.startswith("event:"):
            return await self._parse_stream_text_response(text, session)

        try:
            resp_json = json.loads(text)
        except Exception:
            raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
        return await self._parse_response(resp_json, session)

    # ── 中断轮询 ──────────────────────────────────────────────────────────────

    @staticmethod
    async def _poll_interrupt():
        """每 0.5s 轮询一次 ComfyUI 中断标志"""
        while True:
            await asyncio.sleep(0.5)
            if _INTERRUPT_AVAILABLE and processing_interrupted():
                return

    @staticmethod
    async def _run_with_interrupt(coro):
        """
        将异步任务与中断轮询并发执行。
        如果用户点击取消，cancel 掉 coro 并抛出 InterruptProcessingException。
        """
        if not _INTERRUPT_AVAILABLE:
            return await coro

        request_task = asyncio.ensure_future(coro)
        interrupt_task = asyncio.ensure_future(GptImageClient._poll_interrupt())

        done, pending = await asyncio.wait(
            [request_task, interrupt_task],
            return_when=asyncio.FIRST_COMPLETED,
        )

        for t in pending:
            t.cancel()
            try:
                await t
            except (asyncio.CancelledError, Exception):
                pass

        if interrupt_task in done and request_task not in done:
            raise InterruptProcessingException()

        return request_task.result()

    # ── 新版异步 GPT Image 接口 ──────────────────────────────────────────────

    def _build_async_generate_body(
        self,
        prompt: str,
        model: str,
        quality: str,
        size: Optional[str],
        n: int,
        image_list: Optional[List[torch.Tensor]] = None,
        mask_tensor: Optional[torch.Tensor] = None,
        output_format: str = "png",
        background: str = "auto",
        moderation: Optional[str] = None,
        resize_mode: Optional[str] = None,
    ) -> dict:
        if int(n) != 1:
            raise ValueError("GPT Image 单次请求的 n 必须为 1；多图请使用外层并发")
        output_format = str(output_format or "png").strip().lower()
        background = str(background or "auto").strip().lower()
        moderation = str(moderation).strip().lower() if moderation is not None else None
        if output_format not in GPT_IMAGE_OUTPUT_FORMAT_OPTIONS:
            raise ValueError("GPT Image 输出格式无效")
        if background not in GPT_IMAGE_BACKGROUND_OPTIONS:
            raise ValueError("GPT Image 背景参数无效")
        if moderation not in {None, "low"}:
            raise ValueError("GPT Image 内容审查强度无效")
        if background == "transparent" and output_format == "jpeg":
            raise ValueError("GPT Image 透明背景仅支持 PNG 或 WebP 输出格式")
        api_model = _MODEL_NAME_MAP.get(model, model)
        assets: List[Dict[str, Any]] = []
        image_keys: List[str] = []
        mask_key = None

        if image_list:
            for idx_img, img_tensor in enumerate(image_list):
                pil_images = tensor_to_pil(img_tensor)
                if not pil_images:
                    continue
                key = f"image_{idx_img}"
                image_keys.append(key)
                assets.append({
                    "key": key,
                    "label": f"参考图{idx_img + 1}",
                    "bytes": self._pil_to_png_bytes(pil_images[0]),
                    "resize_group": key,
                })

        if mask_tensor is not None:
            if not image_list:
                raise ValueError("提供了蒙版但未提供图片，请同时提供图片和蒙版")

            first_tensor = image_list[0]
            if first_tensor.dim() == 3:
                first_tensor = first_tensor.unsqueeze(0)
            image_size = (first_tensor.shape[1], first_tensor.shape[2])
            mask_key = "mask"
            assets.append({
                "key": mask_key,
                "label": "蒙版",
                "bytes": self._mask_tensor_to_rgba_png_bytes(mask_tensor, image_size),
                "resize_group": image_keys[0],
            })

        def _make_body(asset_bytes: Dict[str, bytes]) -> dict:
            body = {
                "model": api_model,
                "prompt": prompt,
                "images": [
                    self._png_bytes_to_data_url(asset_bytes[key])
                    for key in image_keys
                    if key in asset_bytes
                ],
                "quality": quality,
                "n": int(n),
                "output_format": output_format,
                "background": background,
            }
            if size:
                body["size"] = size
            if mask_key and mask_key in asset_bytes:
                body["mask"] = {
                    "image_url": self._png_bytes_to_data_url(asset_bytes[mask_key])
                }
            if moderation == "low":
                body["moderation"] = "low"
            return body

        original_asset_bytes = {asset["key"]: asset["bytes"] for asset in assets}
        original_body = _make_body(original_asset_bytes)
        self._log_original_request_body("async generateImage", original_body)

        asset_bytes = self._fit_png_assets_to_body_limit(
            assets,
            _make_body,
            resize_mode=resize_mode,
        )
        body = _make_body(asset_bytes)
        body_size = self._json_body_size(body)
        body_limit = (
            self._MAX_BODY_BYTES
            if resize_mode is None
            else self._UNIFIED_BODY_LIMIT_BYTES
        )
        if body_size > body_limit:
            raise RuntimeError(
                f"请求体超过 {body_limit / (1024 * 1024):.0f} MiB："
                f"{body_size / (1024 * 1024):.2f} MiB"
            )

        return body

    async def _submit_generate_image_task(
        self,
        session: aiohttp.ClientSession,
        payload: dict,
        task_submitted_callback: Optional[Callable[[str, str, float], None]] = None,
    ) -> str:
        url = f"{self.base_url}{_ENDPOINT_ASYNC_GENERATE}"

        t0 = time.time()
        async with session.post(
            url,
            json=payload,
            headers=self._json_headers(),
        ) as resp:
            elapsed = time.time() - t0
            text = await resp.text()
            self._log_original_response_body(
                f"submit generateImage status={resp.status}",
                text,
            )

            if resp.status not in (200, 201, 202):
                raise RuntimeError(self._extract_error_message(text, resp.status))

            try:
                data = json.loads(text)
            except Exception:
                raise RuntimeError(f"提交响应 JSON 解析失败，原始内容：{text[:500]}") from None

            task_id = data.get("task_id")
            if not task_id:
                raise RuntimeError(f"提交响应中未找到 task_id: {data}")

            status = data.get("status", "")
            if task_submitted_callback:
                task_submitted_callback(task_id, status, elapsed)
            else:
                print(f"[o1key GPT Image] 异步任务已提交 | task_id={task_id} | status={status} | 耗时 {elapsed:.1f}s")
            return task_id

    async def _poll_generate_image_task(
        self,
        session: aiohttp.ClientSession,
        task_id: str,
        progress_callback: Optional[Callable[[int], None]] = None,
        initial_delay: bool = True,
    ) -> dict:
        url = f"{self.base_url}{_ENDPOINT_ASYNC_TASK.format(task_id=task_id)}"
        start_time = time.time()

        poll_count = 0
        last_poll_at = start_time

        while True:
            if poll_count == 0 and not initial_delay:
                next_poll_at = start_time
            elif poll_count < len(_ASYNC_POLL_SCHEDULE):
                next_poll_at = start_time + _ASYNC_POLL_SCHEDULE[poll_count]
            else:
                next_poll_at = last_poll_at + _ASYNC_POLL_INTERVAL

            sleep_time = next_poll_at - time.time()
            if sleep_time > 0:
                await asyncio.sleep(sleep_time)

            last_poll_at = time.time()
            elapsed = last_poll_at - start_time
            if elapsed > _ASYNC_MAX_WAIT:
                raise RuntimeError(f"任务 {task_id} 超时（>{int(_ASYNC_MAX_WAIT)}秒），请稍后用 task_id 查询结果")

            poll_count += 1
            task = await self._get_json_with_response_retry(
                session,
                url,
                self._auth_headers(),
                label=f"任务 {task_id} 查询",
                expected_task_id=task_id,
            )

            status = task.get("status", "UNKNOWN")
            progress = task.get("progress")
            progress_pct = self._coerce_progress_percent(progress)
            if self.poll_log_enabled:
                progress_text = f" | progress={progress}" if progress is not None else ""
                print(f"[o1key GPT Image] 查询任务 #{poll_count} | task_id={task_id} | status={status}{progress_text}")

            if status == "SUCCESS":
                self._emit_progress(progress_callback, 100)
                return task
            if progress_pct is not None and progress_pct < 100:
                self._emit_progress(progress_callback, progress_pct)
            if status == "FAILURE":
                error_message = self._extract_error_message(task, 500) or "生成失败"
                detail = self._extract_error_detail(task)
                code = detail.get("code")
                category = detail.get("category")
                parts = [p for p in (
                    f"code={code}" if code else "",
                    f"category={category}" if category else "",
                    f"task_id={task_id}",
                ) if p]
                raise RuntimeError(f"{error_message} ({', '.join(parts)})")
            if status not in ("SUBMITTED", "IN_PROGRESS"):
                raise RuntimeError(f"未知任务状态 {status}: {task}")

    async def _parse_async_task_images(
        self,
        task: dict,
        session: aiohttp.ClientSession,
        result_url_callback: Optional[Callable[[str], None]] = None,
        log_downloads: bool = True,
        log_prefix: str = "[o1key GPT Image]",
    ) -> List[Image.Image]:
        data = task.get("data", {})
        image_items = data.get("images") if isinstance(data, dict) else None

        if not isinstance(image_items, list) or not image_items:
            raise RuntimeError(f"任务结果中未找到 data.images: {task}")

        images: List[Image.Image] = []
        result_urls: List[str] = []
        try:
            for idx, item in enumerate(image_items, 1):
                if not isinstance(item, dict):
                    continue

                url = item.get("url") or item.get("image_url")
                b64 = item.get("b64_json", "")

                if url and isinstance(url, str) and url.startswith("data:image"):
                    header, separator, b64_data = url.partition(",")
                    if not separator or ";base64" not in header.lower():
                        raise _IncompleteInlineImageError(
                            f"第 {idx} 张内联 data URL 缺少有效 Base64 头"
                        )
                    started = time.perf_counter()
                    img, byte_count = self._decode_inline_result_image(
                        b64_data,
                        f"第 {idx} 张内联 data URL 图片",
                    )
                    images.append(img)
                    if log_downloads:
                        print(
                            f"{log_prefix} 结果图像 {idx}/{len(image_items)} 内联 data URL，无网络下载 | "
                            f"{img.size[0]}×{img.size[1]} | {self._format_transfer_size(byte_count)} | "
                            f"解码={time.perf_counter() - started:.2f}s"
                        )
                elif url and isinstance(url, str) and url.startswith("http"):
                    if log_downloads:
                        print(f"{log_prefix} 下载图像 {idx}/{len(image_items)} 开始")
                    img, byte_count, elapsed = await self._download_image_with_response_retry(
                        session, url, f"第 {idx} 张图片"
                    )
                    images.append(img)
                    result_urls.append(url)
                    if log_downloads:
                        print(
                            f"{log_prefix} 下载图像 {idx}/{len(image_items)} 完成 | "
                            f"{img.size[0]}×{img.size[1]} | {self._format_transfer_size(byte_count)} | "
                            f"{elapsed:.2f}s | {self._format_transfer_rate(byte_count, elapsed)}"
                        )
                elif b64:
                    started = time.perf_counter()
                    img, byte_count = self._decode_inline_result_image(
                        b64,
                        f"第 {idx} 张内联 Base64 图片",
                    )
                    images.append(img)
                    if log_downloads:
                        print(
                            f"{log_prefix} 结果图像 {idx}/{len(image_items)} 内联 base64，无网络下载 | "
                            f"{img.size[0]}×{img.size[1]} | {self._format_transfer_size(byte_count)} | "
                            f"解码={time.perf_counter() - started:.2f}s"
                        )
                else:
                    print(f"[o1key GPT Image] 警告：第 {idx} 条结果既无 url 也无 b64_json，已跳过")
        except Exception:
            for image in images:
                image.close()
            raise

        if not images:
            raise RuntimeError("任务成功但没有可用图片结果")

        if result_url_callback:
            for url in dict.fromkeys(result_urls):
                result_url_callback(url)

        return images

    @staticmethod
    def _inline_result_retry_delay(task_id: str, attempt: int) -> float:
        base_delay = GptImageClient._response_retry_delay(attempt)
        stable_jitter = (sum(task_id.encode("utf-8")) % 500) / 1000.0
        return base_delay + stable_jitter

    async def _parse_completed_task_images_with_retry(
        self,
        task: dict,
        session: aiohttp.ClientSession,
        task_id: str,
        progress_callback: Optional[Callable[[int], None]] = None,
        result_url_callback: Optional[Callable[[str], None]] = None,
        log_downloads: bool = True,
        log_prefix: str = "[o1key GPT Image]",
    ) -> tuple[dict, List[Image.Image]]:
        """Re-fetch a completed task when its inline image is incomplete."""
        current_task = task
        for attempt in range(_RESPONSE_READ_MAX_RETRIES + 1):
            try:
                images = await self._parse_async_task_images(
                    current_task,
                    session,
                    result_url_callback=result_url_callback,
                    log_downloads=log_downloads,
                    log_prefix=log_prefix,
                )
                return current_task, images
            except _IncompleteInlineImageError as exc:
                if attempt >= _RESPONSE_READ_MAX_RETRIES:
                    raise RuntimeError(
                        f"GPT Image: task_id={task_id} 的内联结果在 "
                        f"{attempt + 1} 次获取后仍不完整: {exc}"
                    ) from None

                delay = self._inline_result_retry_delay(task_id, attempt)
                print(
                    f"[o1key GPT Image] 内联结果不完整，{delay:.1f}s 后重新查询同一任务 "
                    f"| task_id={task_id} | retry={attempt + 1}/{_RESPONSE_READ_MAX_RETRIES}"
                )
                await asyncio.sleep(delay)
                current_task = await self._poll_generate_image_task(
                    session,
                    task_id,
                    progress_callback=progress_callback,
                    initial_delay=False,
                )

        raise RuntimeError(f"GPT Image: task_id={task_id} 的内联结果读取失败")  # pragma: no cover

    async def _generate_image_task_async(
        self,
        prompt: str,
        model: str,
        quality: str,
        size: Optional[str],
        n: int,
        seed: int,
        image_tensor: Optional[List[torch.Tensor]] = None,
        mask_tensor: Optional[torch.Tensor] = None,
        output_format: str = "png",
        background: str = "auto",
        moderation: Optional[str] = None,
        progress_callback: Optional[Callable[[int], None]] = None,
        result_url_callback: Optional[Callable[[str], None]] = None,
        log_downloads: bool = True,
        log_prefix: str = "[o1key GPT Image]",
        task_submitted_callback: Optional[Callable[[str, str, float], None]] = None,
        log_request_start: bool = True,
        task_completed_callback: Optional[Callable[[str, int, float, List[str]], None]] = None,
        resize_mode: Optional[str] = None,
    ) -> List[Image.Image]:
        body = self._build_async_generate_body(
            prompt=prompt,
            model=model,
            quality=quality,
            size=size,
            n=n,
            image_list=image_tensor,
            mask_tensor=mask_tensor,
            output_format=output_format,
            background=background,
            moderation=moderation,
            resize_mode=resize_mode,
        )

        mode = "图像编辑" if mask_tensor is not None else ("图生图" if image_tensor else "文生图")
        body_size = self._json_body_size(body)
        if log_request_start:
            print(
                f"[o1key GPT Image] {mode} | 新异步接口 | 模型={model} | "
                f"quality={quality} | size={size} | n={n} | body={body_size // 1024}KB"
            )

        connector = aiohttp.TCPConnector(ssl=False, force_close=True)
        timeout = aiohttp.ClientTimeout(total=_ASYNC_MAX_WAIT + 120)
        task_started = time.time()
        async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
            task_id = await self._submit_generate_image_task(
                session, body, task_submitted_callback=task_submitted_callback,
            )
            task = await self._poll_generate_image_task(
                session,
                task_id,
                progress_callback=progress_callback,
            )
            upstream_urls: List[str] = []
            def _on_result_url(url: str) -> None:
                upstream_urls.append(url)
                if result_url_callback:
                    result_url_callback(url)
            task, images = await self._parse_completed_task_images_with_retry(
                task,
                session,
                task_id,
                progress_callback=progress_callback,
                result_url_callback=_on_result_url,
                log_downloads=log_downloads,
                log_prefix=log_prefix,
            )
            if task_completed_callback:
                task_completed_callback(task_id, len(images), time.time() - task_started, upstream_urls)
            return images

    async def _generate_special_price_images_parallel(
        self,
        prompt: str,
        model: str,
        quality: str,
        size: Optional[str],
        n: int,
        seed: int,
        image_tensor: Optional[List[torch.Tensor]] = None,
        mask_tensor: Optional[torch.Tensor] = None,
        output_format: str = "png",
        background: str = "auto",
        moderation: Optional[str] = None,
        progress_callback: Optional[Callable[[int], None]] = None,
        result_url_callback: Optional[Callable[[str], None]] = None,
        log_downloads: bool = True,
        log_prefix: str = "[o1key GPT Image]",
        task_submitted_callback: Optional[Callable[[str, str, float], None]] = None,
        log_request_start: bool = True,
        task_completed_callback: Optional[Callable[[str, int, float, List[str]], None]] = None,
        resize_mode: Optional[str] = None,
    ) -> List[Image.Image]:
        """多图生成：并发提交 n 个单图任务，每个请求固定传 n=1。"""
        request_count = max(1, int(n))
        progress_values = [0] * request_count

        def _task_progress_callback(task_index: int):
            if progress_callback is None:
                return None

            def _update(value: int):
                try:
                    progress_values[task_index] = max(0, min(100, int(value)))
                except (TypeError, ValueError):
                    return
                progress_callback(int(sum(progress_values) / request_count))

            return _update

        if log_request_start:
            print(
                f"[o1key GPT Image] 多图并发模式 | 并发请求={request_count} | 每个请求 n=1"
            )
        tasks = [
            asyncio.create_task(
                self._generate_image_task_async(
                    prompt=prompt,
                    model=model,
                    quality=quality,
                    size=size,
                    n=1,
                    seed=seed,
                    image_tensor=image_tensor,
                    mask_tensor=mask_tensor,
                    output_format=output_format,
                    background=background,
                    moderation=moderation,
                    progress_callback=_task_progress_callback(task_index),
                    result_url_callback=result_url_callback,
                    log_downloads=log_downloads,
                    log_prefix=log_prefix,
                    task_submitted_callback=task_submitted_callback,
                    log_request_start=log_request_start,
                    task_completed_callback=task_completed_callback,
                    resize_mode=resize_mode,
                )
            )
            for task_index in range(request_count)
        ]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        images: List[Image.Image] = []
        errors = []
        for task_index, result in enumerate(results, 1):
            if isinstance(result, BaseException):
                errors.append(f"第{task_index}个请求: {str(result).splitlines()[0]}")
            else:
                images.extend(result)

        if progress_callback is not None:
            progress_callback(100)

        if errors:
            print(
                f"[o1key GPT Image] 多图并发完成 | 成功={request_count - len(errors)} "
                f"| 失败={len(errors)} | {'; '.join(errors)}"
            )
        if not images:
            detail = "; ".join(errors) or "未返回图片"
            raise RuntimeError(f"多图并发请求全部失败：{detail}")

        return images

    # ── 文生图 / 图生图（generations 接口）───────────────────────────────────

    async def _generate_async(
        self,
        prompt: str,
        model: str,
        quality: str,
        size: str,
        n: int,
        seed: int,
        image_list: Optional[List[torch.Tensor]] = None,
    ) -> List[Image.Image]:
        """
        调用 /v1/images/generations/ 接口。
        当传入 image_list 时，以 data URI 格式内联图像（图生图）。
        """
        # 模型名映射：UI 显示名 → API 参数名
        api_model = _MODEL_NAME_MAP.get(model, model)

        body: dict = {
            "model":      api_model,
            "prompt":     prompt,
            "quality":    quality,
            "n":          n,
            "moderation": "low",
            "partial_images": 0,
        }

        body["size"] = size if size else "auto"

        image_files = []

        # 图生图：multipart 方式上传参考图
        if image_list is not None:
            for idx_img, img_tensor in enumerate(image_list):
                pil_images = tensor_to_pil(img_tensor)
                img = pil_images[0]
                buf = BytesIO()
                img.save(buf, format="PNG")
                png_bytes = buf.getvalue()
                # 单张图像预算：20MB 按图数平摊，至少保留 1MB 给其他字段
                per_image_budget = max(
                    1024 * 1024,
                    (self._MAX_BODY_BYTES - 1024 * 1024) // len(image_list),
                )
                # base64 膨胀约 4/3，所以 PNG 目标上限 = budget * 3/4
                png_budget = int(per_image_budget * 3 / 4)
                label = f"第{idx_img + 1}张" if len(image_list) > 1 else ""
                png_bytes = self._shrink_png_to_limit(png_bytes, png_budget, label)
                image_files.append(png_bytes)
            mode = f"图生图（参考图 {len(image_files)} 张）"
        else:
            mode = "文生图"

        url = f"{self.base_url}{_ENDPOINT_GENERATIONS}"
        print(f"[o1key GPT Image] {mode} | 模型={model} | quality={quality} | "
              f"size={size} | n={n}")

        connector = aiohttp.TCPConnector(ssl=False, force_close=True)
        timeout   = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)

        def _build_multipart_form() -> aiohttp.FormData:
            form = self._new_multipart_form()
            self._add_form_fields(form, body)
            image_field = "image[]" if len(image_files) > 1 else "image"
            for idx_img, png_bytes in enumerate(image_files):
                form.add_field(
                    image_field,
                    png_bytes,
                    filename=f"image_{idx_img + 1}.png",
                    content_type="image/png",
                )
            return form

        async def _do_request():
            async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
                t0 = time.time()
                async with session.post(
                    url,
                    data=_build_multipart_form(),
                    headers=self._auth_headers(),
                ) as resp:
                    elapsed = time.time() - t0

                    if resp.status != 200:
                        text = await resp.text()
                        if resp.status in HTTP_ERROR_MESSAGES:
                            raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
                        try:
                            err_json = json.loads(text)
                            err_obj  = err_json.get("error", {})
                            msg = (
                                err_obj.get("message") or err_obj.get("msg") or text
                                if isinstance(err_obj, dict)
                                else str(err_obj) or text
                            )
                        except Exception:
                            msg = text
                        raise RuntimeError(get_friendly_message(resp.status, msg))

                    print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
                    return await self._parse_success_response(resp, session, "GENERATIONS")

        return await self._run_with_interrupt(_do_request())

    # ── 图像编辑（edits 接口，multipart/form-data）──────────────────────

    async def _edit_async(
        self,
        prompt: str,
        model: str,
        quality: str,
        size: str,
        n: int,
        seed: int,
        image_list: List[torch.Tensor],
        mask_tensor: Optional[torch.Tensor] = None,
    ) -> List[Image.Image]:
        """
        调用 /v1/images/edits 接口（multipart/form-data）。
        """
        # 模型名映射：UI 显示名 → API 参数名
        api_model = _MODEL_NAME_MAP.get(model, model)

        # 统一 tensors 为 [1,H,W,C] 格式，支持不同尺寸
        normalized_tensors = []
        for t in image_list:
            if t.dim() == 3:
                t = t.unsqueeze(0)   # [H,W,C] → [1,H,W,C]
            normalized_tensors.append(t)
        num_images = len(normalized_tensors)

        image_files = []

        # 多图：用 multipart image/image[] 字段逐张上传
        # 预算：20MB 按图数平摊，蒙版预留 1MB
        mask_reserve = 1024 * 1024 if mask_tensor is not None else 0
        per_image_budget = max(
            1024 * 1024,
            (self._MAX_BODY_BYTES - mask_reserve) // num_images,
        )
        for i, frame in enumerate(normalized_tensors):
            img_bytes = self._tensor_to_png_bytes(frame)
            label = f"第{i + 1}张" if num_images > 1 else ""
            img_bytes = self._shrink_png_to_limit(img_bytes, per_image_budget, label)
            image_files.append(img_bytes)

        # 蒙版尺寸校验以第一张图为基准
        first_tensor = normalized_tensors[0]
        ih, iw = first_tensor.shape[1], first_tensor.shape[2]

        mask_png = None
        if mask_tensor is not None:
            mask_png = self._mask_tensor_to_rgba_png_bytes(mask_tensor, (ih, iw))
            mode = "图像编辑（带蒙版）"
        else:
            mode = "图像编辑（无蒙版）"

        form_fields = {
            "model": api_model,
            "prompt": prompt,
            "partial_images": 0,
            "n": n,
            "quality": quality,
            "size": size if size else "auto",
            "output_format": "png",
            "background": "opaque",
            "moderation": "low",
        }

        def _build_multipart_form() -> aiohttp.FormData:
            form = self._new_multipart_form()
            self._add_form_fields(form, form_fields)

            for idx_img, img_bytes in enumerate(image_files):
                form.add_field(
                    "image[]",
                    img_bytes,
                    filename=f"image_{idx_img + 1}.png",
                    content_type="image/png",
                )
            if mask_png is not None:
                form.add_field(
                    "mask",
                    mask_png,
                    filename="mask.png",
                    content_type="image/png",
                )
            return form

        url = f"{self.base_url}{_ENDPOINT_EDITS}"
        print(f"[o1key GPT Image] {mode} | 模型={model} | 参考图={num_images}张 | "
              f"quality={quality} | size={size} | n={n}")

        connector = aiohttp.TCPConnector(ssl=False, force_close=True)
        timeout   = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)

        async def _do_request():
            async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
                t0 = time.time()
                async with session.post(
                    url,
                    data=_build_multipart_form(),
                    headers=self._auth_headers(),
                ) as resp:
                    elapsed = time.time() - t0

                    if resp.status != 200:
                        text = await resp.text()
                        if resp.status in HTTP_ERROR_MESSAGES:
                            raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
                        try:
                            err_json = json.loads(text)
                            err_obj  = err_json.get("error", {})
                            msg = (
                                err_obj.get("message") or err_obj.get("msg") or text
                                if isinstance(err_obj, dict)
                                else str(err_obj) or text
                            )
                        except Exception:
                            msg = text
                        raise RuntimeError(get_friendly_message(resp.status, msg))

                    print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
                    return await self._parse_success_response(resp, session, "EDITS")

        return await self._run_with_interrupt(_do_request())

    # ── 同步统一入口（供节点调用）────────────────────────────────────────────

    def run_sync(
        self,
        prompt: str,
        model: str,
        quality: str,
        size: str,
        n: int,
        seed: int,
        image_tensor: Optional[List[torch.Tensor]] = None,
        mask_tensor: Optional[torch.Tensor] = None,
    ) -> List[Image.Image]:
        """
        同步入口，在独立线程中运行事件循环，避免与 ComfyUI 主循环冲突。

        路由逻辑：
          - 无 image_tensor  → generations 接口（文生图，multipart/form-data）
          - 有 image_tensor  → edits 接口（图生图/编辑，multipart/form-data）
        """
        use_edits = (image_tensor is not None)

        if use_edits:
            coro = self._edit_async(
                prompt=prompt, model=model, quality=quality,
                size=size, n=n, seed=seed,
                image_list=image_tensor, mask_tensor=mask_tensor,
            )
        else:
            coro = self._generate_async(
                prompt=prompt, model=model, quality=quality,
                size=size, n=n, seed=seed,
                image_list=image_tensor,
            )

        def _run():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                return loop.run_until_complete(coro)
            finally:
                loop.close()

        with ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(_run)
            try:
                return future.result(timeout=_REQUEST_TIMEOUT + 30)
            except TimeoutError:
                raise RuntimeError(
                    f"o1key GPT Image 请求超时（>{_REQUEST_TIMEOUT}s），请检查网络或稍后重试"
                )

    async def generate_image_async(
        self,
        prompt: str,
        model: str,
        quality: str,
        size: Optional[str],
        n: int,
        seed: int,
        image_tensor: Optional[List[torch.Tensor]] = None,
        mask_tensor: Optional[torch.Tensor] = None,
        output_format: str = "png",
        background: str = "auto",
        moderation: Optional[str] = None,
        progress_callback: Optional[Callable[[int], None]] = None,
        special_price_parallel: bool = False,
        result_url_callback: Optional[Callable[[str], None]] = None,
        log_downloads: bool = True,
        log_prefix: str = "[o1key GPT Image]",
        task_submitted_callback: Optional[Callable[[str, str, float], None]] = None,
        log_request_start: bool = True,
        task_completed_callback: Optional[Callable[[str, int, float, List[str]], None]] = None,
        resize_mode: Optional[str] = None,
    ) -> List[Image.Image]:
        """异步公共入口；多图始终拆成并发的 n=1 单图请求。"""
        requested_n = max(1, int(n))
        if requested_n > 1:
            return await self._generate_special_price_images_parallel(
                prompt=prompt,
                model=model,
                quality=quality,
                size=size,
                n=requested_n,
                seed=seed,
                image_tensor=image_tensor,
                mask_tensor=mask_tensor,
                output_format=output_format,
                background=background,
                moderation=moderation,
                progress_callback=progress_callback,
                result_url_callback=result_url_callback,
                log_downloads=log_downloads,
                log_prefix=log_prefix,
                task_submitted_callback=task_submitted_callback,
                log_request_start=log_request_start,
                task_completed_callback=task_completed_callback,
                resize_mode=resize_mode,
            )
        return await self._generate_image_task_async(
            prompt=prompt,
            model=model,
            quality=quality,
            size=size,
            n=requested_n,
            seed=seed,
            image_tensor=image_tensor,
            mask_tensor=mask_tensor,
            output_format=output_format,
            background=background,
            moderation=moderation,
            progress_callback=progress_callback,
            result_url_callback=result_url_callback,
            log_downloads=log_downloads,
            log_prefix=log_prefix,
            task_submitted_callback=task_submitted_callback,
            log_request_start=log_request_start,
            task_completed_callback=task_completed_callback,
            resize_mode=resize_mode,
        )

    def generate_image_async_sync(
        self,
        prompt: str,
        model: str,
        quality: str,
        size: Optional[str],
        n: int,
        seed: int,
        image_tensor: Optional[List[torch.Tensor]] = None,
        mask_tensor: Optional[torch.Tensor] = None,
        output_format: str = "png",
        background: str = "auto",
        moderation: Optional[str] = None,
        progress_callback: Optional[Callable[[int], None]] = None,
        special_price_parallel: bool = False,
        result_url_callback: Optional[Callable[[str], None]] = None,
        log_downloads: bool = True,
        log_prefix: str = "[o1key GPT Image]",
        task_submitted_callback: Optional[Callable[[str, str, float], None]] = None,
        log_request_start: bool = True,
        task_completed_callback: Optional[Callable[[str, int, float, List[str]], None]] = None,
        resize_mode: Optional[str] = None,
    ) -> List[Image.Image]:
        """
        新版异步任务入口，供节点调用。
        旧 run_sync 保留兼容，但 GPT Image 节点不再使用旧同步接口。
        """
        coro = self.generate_image_async(
            prompt=prompt,
            model=model,
            quality=quality,
            size=size,
            n=n,
            seed=seed,
            image_tensor=image_tensor,
            mask_tensor=mask_tensor,
            output_format=output_format,
            background=background,
            moderation=moderation,
            progress_callback=progress_callback,
            special_price_parallel=special_price_parallel,
            result_url_callback=result_url_callback,
            log_downloads=log_downloads,
            log_prefix=log_prefix,
            task_submitted_callback=task_submitted_callback,
            log_request_start=log_request_start,
            task_completed_callback=task_completed_callback,
            resize_mode=resize_mode,
        )

        def _run():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                return loop.run_until_complete(self._run_with_interrupt(coro))
            finally:
                loop.close()

        with ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(_run)
            try:
                return future.result(timeout=_ASYNC_MAX_WAIT + 150)
            except TimeoutError:
                raise RuntimeError(
                    f"o1key GPT Image 异步任务超时（>{int(_ASYNC_MAX_WAIT)}秒），请稍后用 task_id 查询结果"
                )

    # ── 余额查询 ──────────────────────────────────────────────────────────────

    async def _query_balance_async(self) -> dict:
        url = f"{self.base_url}/api/usage/token"
        connector = aiohttp.TCPConnector(ssl=False, force_close=True)
        timeout   = aiohttp.ClientTimeout(total=10)
        async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
            async with session.get(url, headers=self._auth_headers()) as resp:
                if resp.status != 200:
                    raise RuntimeError(f"余额查询失败 HTTP {resp.status}")
                return await resp.json()

    def query_balance_sync(self) -> dict:
        def _run():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                return loop.run_until_complete(self._query_balance_async())
            finally:
                loop.close()

        with ThreadPoolExecutor(max_workers=1) as executor:
            return executor.submit(_run).result(timeout=15)

    @staticmethod
    def format_balance_info(balance_data: dict) -> str:
        data             = balance_data.get("data", {})
        api_name         = data.get("name", "未知")
        total_available  = data.get("total_available", 0)
        balance_in_dollars = total_available / 500000
        return f"当前余额：{balance_in_dollars:.2f} | API：{api_name}"