"""
Grok Image API 客户端
支持两个接口：
  - POST /v1/images/generations  文生图
  - POST /v1/images/edits        图生图（带参考图）

上游 API 格式与 OpenAI Images API 兼容。
"""

import asyncio
import base64
import json
import time
from concurrent.futures import ThreadPoolExecutor
from io import BytesIO
from typing import List, Optional

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

from ..utils.config import get_api_key_or_raise, get_base_url_by_route
from ..utils.image_utils import tensor_to_pil

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"

_MODEL_NAME_MAP = {
    "Grok Image": "grok-imagine-image",
    "Grok Image Pro": "grok-imagine-image-quality",
}

_REQUEST_TIMEOUT = 900
_MAX_BODY_BYTES = 20 * 1024 * 1024

_MAX_RETRIES = 3
_RETRY_DELAY = 5


class GrokImageClient:

    def __init__(self, route: str = None):
        self.api_key = get_api_key_or_raise("O1KEY_API_KEY")
        self.base_url = get_base_url_by_route(route)

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

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

    # ── 图像工具 ──────────────────────────────────────────────────────────────

    @staticmethod
    def _shrink_png_to_limit(png_bytes: bytes, max_bytes: int, label: str = "") -> bytes:
        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
            next_w = max(1, int(w * scale))
            next_h = max(1, int(h * scale))
            if (next_w, next_h) == (w, h):
                raise RuntimeError(
                    f"Grok Image 无法将图像缩小到 {max_bytes} 字节以内"
                )
            w, h = next_w, next_h
            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 Grok 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 _pil_list_to_tensor(images: List[Image.Image]) -> torch.Tensor:
        if not images:
            placeholder = Image.new("RGB", (512, 512), (128, 128, 128))
            images = [placeholder]
        tensors = []
        for img in images:
            arr = np.array(img.convert("RGB")).astype(np.float32) / 255.0
            tensors.append(torch.from_numpy(arr))
        return torch.stack(tensors, dim=0)

    @classmethod
    def _build_edit_body(
        cls,
        prompt: str,
        model: str,
        aspect_ratio: str,
        resolution: str,
        image_list: List[torch.Tensor],
    ) -> dict:
        """构建最多三张参考图的编辑请求，并确保完整 JSON 不超过 20MB。"""
        if len(image_list) > 3:
            raise ValueError("Grok Image 最多支持 3 张参考图")

        reference_images = []
        for index, tensor in enumerate(image_list, start=1):
            pil_images = tensor_to_pil(tensor)
            if not pil_images:
                raise ValueError(f"无法读取参考图{index}")
            reference_images.append(pil_images[0])
        if not reference_images:
            raise ValueError("图像编辑至少需要 1 张参考图")

        body: dict = {
            "model": _MODEL_NAME_MAP.get(model, model),
            "prompt": prompt,
            "response_format": "b64_json",
        }
        if aspect_ratio and aspect_ratio != "auto":
            body["aspect_ratio"] = aspect_ratio
        if resolution:
            body["resolution"] = resolution

        # base64 大约是原始字节的 4/3。预留 1MB 给提示词和 JSON 字段，
        # 同时保留旧逻辑的单图 PNG 最大 10MB 上限。
        base_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
        available = _MAX_BODY_BYTES - base_size - 1024 * 1024
        if available <= 0:
            raise ValueError("提示词和请求参数已超过 Grok Image 20MB 请求体限制")
        per_image_limit = min(
            _MAX_BODY_BYTES // 2,
            max(1, int(available * 0.75) // len(reference_images)),
        )

        png_images = []
        for index, image in enumerate(reference_images, start=1):
            buffer = BytesIO()
            image.save(buffer, format="PNG")
            png_images.append(
                cls._shrink_png_to_limit(
                    buffer.getvalue(),
                    per_image_limit,
                    label=f"参考图{index}",
                )
            )

        def _set_images() -> None:
            encoded = [base64.b64encode(data).decode("ascii") for data in png_images]
            if len(encoded) == 1:
                # 单图保持当前 o1key 兼容格式，不改变既有请求行为。
                body.pop("images", None)
                body["image"] = encoded[0]
            else:
                body.pop("image", None)
                body["images"] = [
                    {
                        "type": "image_url",
                        "url": f"data:image/png;base64,{value}",
                    }
                    for value in encoded
                ]

        _set_images()
        body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
        for _ in range(8):
            if body_size <= _MAX_BODY_BYTES:
                return body
            shrink_ratio = max(0.1, (_MAX_BODY_BYTES / body_size) * 0.95)
            png_images = [
                cls._shrink_png_to_limit(
                    data,
                    max(1, int(len(data) * shrink_ratio)),
                    label=f"参考图{index}",
                )
                for index, data in enumerate(png_images, start=1)
            ]
            _set_images()
            body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))

        raise RuntimeError(
            f"Grok Image 参考图缩放后请求体仍超过 20MB：{body_size / 1024 / 1024:.2f}MB"
        )

    @staticmethod
    def _redact_edit_body(body: dict) -> dict:
        result = dict(body)
        image = result.get("image")
        if isinstance(image, str) and len(image) > 50:
            result["image"] = image[:50] + "..."
        images = result.get("images")
        if isinstance(images, list):
            result["images"] = [
                {
                    **item,
                    "url": item.get("url", "")[:50] + "...",
                }
                if isinstance(item, dict) else item
                for item in images
            ]
        return result

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

    @staticmethod
    async def _poll_interrupt():
        while True:
            await asyncio.sleep(0.5)
            if _INTERRUPT_AVAILABLE and processing_interrupted():
                return

    @staticmethod
    async def _run_with_interrupt(coro):
        if not _INTERRUPT_AVAILABLE:
            return await coro
        request_task = asyncio.ensure_future(coro)
        interrupt_task = asyncio.ensure_future(GrokImageClient._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()

    # ── 响应解析 ──────────────────────────────────────────────────────────────

    async def _parse_response(self, resp_json: dict, session: aiohttp.ClientSession) -> List[Image.Image]:
        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 not data_list:
            raise RuntimeError(f"API 响应中未找到 data 字段")
        images: List[Image.Image] = []
        for idx, item in enumerate(data_list):
            b64 = item.get("b64_json", "")
            url = item.get("url", "")
            if b64:
                img_bytes = base64.b64decode(b64)
                img = Image.open(BytesIO(img_bytes))
                images.append(img)
            elif url and url.startswith("http"):
                async with session.get(url, allow_redirects=True) as r:
                    if r.status != 200:
                        raise RuntimeError(f"图像下载失败 HTTP {r.status}")
                    img_bytes = await r.read()
                images.append(Image.open(BytesIO(img_bytes)))
            else:
                print(f"[o1key Grok Image] 警告：第 {idx + 1} 条数据无有效图像，已跳过")
        return images

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

    async def _generate_async(
        self,
        prompt: str,
        model: str,
        aspect_ratio: str,
        resolution: str,
        n: int,
    ) -> List[Image.Image]:
        api_model = _MODEL_NAME_MAP.get(model, model)
        body: dict = {
            "model": api_model,
            "prompt": prompt,
            "aspect_ratio": aspect_ratio if aspect_ratio else "auto",
            "resolution": resolution if resolution else "1k",
            "response_format": "b64_json",
        }

        url = f"{self.base_url}{_ENDPOINT_GENERATIONS}"
        log_body = {k: v for k, v in body.items()}
        print(f"[o1key Grok Image] 请求 URL: {url}")
        print(f"[o1key Grok Image] 请求体: {json.dumps(log_body, ensure_ascii=False)}")

        results = []
        for i in range(n):
            images = await self._do_request_with_retry(url, body)
            results.extend(images)
            if n > 1:
                print(f"[o1key Grok Image] 第 {i+1}/{n} 张完成")
        return results

    # ── 图生图（edits 接口）───────────────────────────────────────────────────

    async def _edit_async(
        self,
        prompt: str,
        model: str,
        aspect_ratio: str,
        resolution: str,
        n: int,
        image_list: List[torch.Tensor],
    ) -> List[Image.Image]:
        body = self._build_edit_body(
            prompt=prompt,
            model=model,
            aspect_ratio=aspect_ratio,
            resolution=resolution,
            image_list=image_list,
        )

        url = f"{self.base_url}{_ENDPOINT_EDITS}"
        log_body = self._redact_edit_body(body)
        print(f"[o1key Grok Image] 请求 URL: {url}")
        print(f"[o1key Grok Image] 请求体: {json.dumps(log_body, ensure_ascii=False)}")

        results = []
        for i in range(n):
            images = await self._do_request_with_retry(url, body)
            results.extend(images)
            if n > 1:
                print(f"[o1key Grok Image] 第 {i+1}/{n} 张完成")
        return results

    # ── 带重试的请求 ────────────────────────────────────────────────────────

    async def _do_request_with_retry(self, url: str, body: dict) -> List[Image.Image]:
        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:
                last_error = None
                for attempt in range(1, _MAX_RETRIES + 1):
                    t0 = time.time()
                    async with session.post(url, json=body, headers=self._json_headers()) as resp:
                        elapsed = time.time() - t0
                        text = await resp.text()

                        if resp.status == 429 or resp.status in (502, 503, 504):
                            last_error = f"HTTP {resp.status}"
                            print(f"[o1key Grok Image] 重试 {attempt}/{_MAX_RETRIES}（{last_error}）")
                            await asyncio.sleep(_RETRY_DELAY * attempt)
                            continue

                        if resp.status == 400 and "high load" in text.lower():
                            last_error = "high load"
                            print(f"[o1key Grok Image] 重试 {attempt}/{_MAX_RETRIES}（服务繁忙）")
                            await asyncio.sleep(_RETRY_DELAY * attempt)
                            continue

                        if resp.status != 200:
                            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(f"请求失败 HTTP {resp.status}: {msg}")

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

                    print(f"[o1key Grok Image] API 响应耗时 {elapsed:.1f}s")
                    return await self._parse_response(resp_json, session)

                raise RuntimeError(f"重试 {_MAX_RETRIES} 次后仍失败: {last_error}")

        return await self._run_with_interrupt(_do_request())

    # ── 同步入口 ──────────────────────────────────────────────────────────────

    def run_sync(
        self,
        prompt: str,
        model: str,
        aspect_ratio: str,
        resolution: str,
        n: int,
        image_list: Optional[List[torch.Tensor]] = None,
    ) -> List[Image.Image]:
        if image_list:
            coro = self._edit_async(
                prompt=prompt, model=model, aspect_ratio=aspect_ratio,
                resolution=resolution, n=n, image_list=image_list,
            )
        else:
            coro = self._generate_async(
                prompt=prompt, model=model, aspect_ratio=aspect_ratio,
                resolution=resolution, n=n,
            )

        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("Grok Image 请求超时，请检查网络或稍后重试")

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

    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}"