"""
gemini_factory.py — 每任务构建独立的 AI client 实例
解决全局单例问题（扒 banana gemini_service.py __init__，改造为每任务局部实例）

铁律（基石B）：
- 调用方只传 task_id，不传 key
- 本模块在 worker 内部查库 → Fernet 解密 → 构建 client
- 解密结果只活在局部变量，函数返回后即销毁
- 绝不打印 / 记录 / 传递明文 key
"""
from __future__ import annotations
import asyncio
import logging
import os
from dataclasses import dataclass, field
from typing import Any

import httpx

logger = logging.getLogger(__name__)


@dataclass
class AIClients:
    """
    一个任务专用的 AI client 集合。
    worker 在任务开始时构建，任务结束后释放（局部变量，不存 Redis/DB）。
    """
    # httpx AsyncClient 懒加载 + 按事件循环缓存：Celery 每任务多次 asyncio.run，
    # 持久 client 会绑死到首个已关闭的 loop → "Event loop is closed"。
    # 故只存 token/base，按当前运行 loop 缓存 client，loop 变了就重建。
    _gpt_token: str | None = field(default=None, repr=False)
    _gpt_base: str | None = field(default=None, repr=False)
    _gpt_client: httpx.AsyncClient | None = field(default=None, repr=False)
    _gpt_client_loop_id: int | None = field(default=None, repr=False)
    # 备用 OpenAI 兼容中转站（codeproxy）：apiports 503 时真正切过去（独立 base+key）
    _alt_token: str | None = field(default=None, repr=False)
    _alt_base: str | None = field(default=None, repr=False)
    # 多 base/token 的 client 池：key=(base,token的id)，按 loop 失效重建
    _client_pool: dict = field(default_factory=dict, repr=False)
    _pool_loop_id: int | None = field(default=None, repr=False)
    _gemini_key: str | None = field(default=None, repr=False)   # 局部变量不打印
    _model_image: str = "gpt-image-2"
    _model_text: str = "claude-sonnet-4-5"  # apiports无gpt-4o-mini,文案用claude中文质量好

    def _client(self) -> httpx.AsyncClient:
        """主通道(apiports) client，按当前事件循环缓存"""
        return self._client_for(self._gpt_base, self._gpt_token)

    def _client_for(self, base: str | None, token: str | None) -> httpx.AsyncClient:
        """按 (base, token) 返回可用 client；loop 变化则整池重建（避免跨 loop 复用）"""
        if not token:
            raise RuntimeError("GPT client 未初始化（缺 token）")
        loop_id = id(asyncio.get_running_loop())
        if self._pool_loop_id != loop_id:
            self._client_pool = {}
            self._pool_loop_id = loop_id
        ck = (base or "", token)
        if ck not in self._client_pool:
            self._client_pool[ck] = httpx.AsyncClient(
                headers={"Authorization": f"Bearer {token}"},
                base_url=base or None,
                timeout=120.0,
            )
        return self._client_pool[ck]

    # ── ImageClient 协议实现（供 image_gen.py 使用）────────

    def _gpt_target(self, provider: str | None) -> tuple[str, str | None, httpx.AsyncClient]:
        """按 provider 选 (base, token, client)；codeproxy 走备用站独立 base+key"""
        if provider == "codeproxy" and self._alt_token:
            base = (self._alt_base or os.environ.get("CODEPROXY_BASE_URL") or "").rstrip("/")
            return base, self._alt_token, self._client_for(base, self._alt_token)
        base = (os.environ.get("IMAGE_API_BASE") or os.environ.get("APIPORTS_BASE_URL") or "").rstrip("/")
        return base, self._gpt_token, self._client_for(self._gpt_base, self._gpt_token)

    async def gpt_edits(self, prompt: str, reference_images: list[bytes], size: str, provider: str | None = None) -> bytes:
        """GPT edits endpoint（带产品参考图，禁纯文生图）"""
        import io
        files: list[tuple] = [("prompt", (None, prompt))]
        for i, img in enumerate(reference_images):
            files.append(("image[]", (f"ref_{i}.png", io.BytesIO(img), "image/png")))
        files.append(("size", (None, size)))
        files.append(("model", (None, self._model_image)))
        base, _, client = self._gpt_target(provider)
        resp = await client.post(f"{base}/images/edits", files=files, timeout=120.0)
        resp.raise_for_status()
        return _extract_image_bytes(resp.json())

    async def gpt_generate(self, prompt: str, size: str, provider: str | None = None) -> bytes:
        """GPT 纯文生图（仅 ALLOW_TEXT_ONLY_IMAGE=true 时用）"""
        base, _, client = self._gpt_target(provider)
        payload = {"model": self._model_image, "prompt": prompt, "n": 1, "size": size}
        resp = await client.post(f"{base}/images/generations", json=payload, timeout=120.0)
        resp.raise_for_status()
        return _extract_image_bytes(resp.json())

    async def gemini_generate(self, prompt: str, reference_images: list[bytes], model: str) -> bytes:
        """Gemini 生图（备用通道）"""
        if not self._gemini_key:
            raise RuntimeError("Gemini key 未初始化")
        gemini_base = os.environ.get("GEMINI_API_URL", "https://generativelanguage.googleapis.com/v1beta")
        url = f"{gemini_base}/models/{model}:generateContent?key={self._gemini_key}"
        parts: list[dict] = [{"text": prompt}]
        for img in reference_images:
            import base64
            parts.append({"inline_data": {"mime_type": "image/png", "data": base64.b64encode(img).decode()}})
        payload = {"contents": [{"role": "user", "parts": parts}], "generationConfig": {"responseModalities": ["IMAGE", "TEXT"]}}
        async with httpx.AsyncClient() as client:
            resp = await client.post(url, json=payload, timeout=120.0)
        resp.raise_for_status()
        return _extract_gemini_image(resp.json())

    async def chat_complete(self, messages: list[dict], model: str | None = None, max_tokens: int = 4096, temperature: float = 0.75) -> str:
        """文字生成（文案生成用）"""
        base = (os.environ.get("IMAGE_API_BASE") or os.environ.get("APIPORTS_BASE_URL") or "").rstrip("/")
        payload = {"model": model or self._model_text, "messages": messages, "max_tokens": max_tokens, "temperature": temperature}
        # 单批≤4条文案正常 40-55s 返回；apiports 网关 ~60s 上限。客户端超时设 75s：
        # 略高于网关上限即可，过长（如180s）会在 apiports 卡顿时干等，拖慢整体。
        timeout = float(os.environ.get("TEXT_LLM_TIMEOUT", "75"))
        resp = await self._client().post(f"{base}/chat/completions", json=payload, timeout=timeout)
        resp.raise_for_status()
        return resp.json()["choices"][0]["message"]["content"] or ""

    async def gpt_vision_analyze(self, prompt: str, images: list[bytes], model: str | None = None) -> str:
        """
        GPT/Claude vision 读产品图，返回 JSON 字符串。
        messages content 混合 text + image_url(base64)，OpenAI vision 格式。
        model 默认最强档（claude-opus-4-8），绝不偷降级。
        最多传 4 张图，避免超 token。
        """
        import base64
        content: list[dict] = [{"type": "text", "text": prompt}]
        for img in images[:4]:
            b64 = base64.b64encode(img).decode()
            content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
            })
        used_model = model or os.environ.get("MODEL_TEXT", "claude-opus-4-8")
        messages = [{"role": "user", "content": content}]
        base = (os.environ.get("IMAGE_API_BASE") or os.environ.get("APIPORTS_BASE_URL") or "").rstrip("/")
        payload = {
            "model": used_model,
            "messages": messages,
            "max_tokens": 2048,
            "temperature": 0.2,
        }
        resp = await self._client().post(f"{base}/chat/completions", json=payload, timeout=90.0)
        resp.raise_for_status()
        return resp.json()["choices"][0]["message"]["content"] or ""

    # duck-type: text_variants._call_llm 用的属性
    @property
    def _model(self) -> str:
        return self._model_text

    async def aclose(self) -> None:
        # client 可能绑在已关闭的 loop（Celery 多次 asyncio.run），aclose 也可能报
        # "Event loop is closed"，吞掉即可——进程级连接随 loop 关闭自然释放。
        for c in list(self._client_pool.values()):
            try:
                await c.aclose()
            except Exception:
                pass
        self._client_pool = {}
        self._pool_loop_id = None
        self._gpt_client = None
        self._gpt_client_loop_id = None


def build_ai_clients(plain_key: str, gemini_key: str | None = None) -> AIClients:
    """
    用解密后的明文 key 构建 AIClients。
    只在 Celery worker 函数体内调用，plain_key 是局部变量。
    httpx client 不在此预创建（避免绑死到调用方 loop），首次 await 时按 loop 懒建。
    调用完成后 caller 负责 await clients.aclose()。
    """
    gpt_base = (
        os.environ.get("IMAGE_API_BASE")          # 旧变量名
        or os.environ.get("APIPORTS_BASE_URL")    # .env 实际变量名
        or ""
    ).rstrip("/")
    # 备用站 codeproxy：系统级 key（非用户录入），apiports 503 时切过去保生图成功
    alt_base = (os.environ.get("CODEPROXY_BASE_URL") or "").rstrip("/")
    alt_token = os.environ.get("CODEPROXY_KEY") or None
    return AIClients(
        _gpt_token=plain_key,
        _gpt_base=gpt_base or None,
        _alt_base=alt_base or None,
        _alt_token=alt_token,
        _gemini_key=gemini_key,
        _model_image=os.environ.get("IMAGE_MODEL") or os.environ.get("MODEL_IMAGE", "gpt-image-2"),
        _model_text=os.environ.get("MODEL_TEXT", "claude-opus-4-8"),
    )


# ── 图片响应解析工具 ─────────────────────────────────────

def _extract_image_bytes(resp_json: dict) -> bytes:
    """从 OpenAI images API 响应提取图片 bytes（b64 或 url）"""
    import base64
    data = resp_json.get("data", [{}])
    if not data:
        raise ValueError("图片 API 返回空 data")
    item = data[0]
    if "b64_json" in item:
        return base64.b64decode(item["b64_json"])
    if "url" in item:
        resp = httpx.get(item["url"], timeout=30.0)
        resp.raise_for_status()
        return resp.content
    raise ValueError(f"无法解析图片响应：{list(item.keys())}")


def _extract_gemini_image(resp_json: dict) -> bytes:
    """从 Gemini generateContent 响应提取图片 bytes"""
    import base64
    candidates = resp_json.get("candidates", [])
    for cand in candidates:
        parts = cand.get("content", {}).get("parts", [])
        for part in parts:
            if "inlineData" in part:
                return base64.b64decode(part["inlineData"]["data"])
    raise ValueError("Gemini 响应中未找到图片数据")
