"""ZenMux API 客户端。

直接 HTTP 调用 zenmux 的 OpenAI 兼容端点。为 Phase 4 的 Python 化脚本服务
（translate.py / polish.py）。所有调用都走 /api/v1/chat/completions。

设计原则：
- 独立于 opencode，可直接在 CLI / CI / cron 运行
- 内置重试（指数退避）、限流、token 统计、结构化日志
- 失败快速可见：打印请求 ID，便于在 zenmux 后台对账
- 默认读 secrets.env 里的 ZENMUX_API_KEY
"""

from __future__ import annotations

import json
import os
import sys
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import httpx

DEFAULT_BASE_URL = "https://zenmux.ai/api/v1"
DEFAULT_TIMEOUT = 300.0  # 翻译/润色单次调用可能 60s+，留够余量
MAX_RETRIES = 5
RETRYABLE_STATUSES = {408, 429, 500, 502, 503, 504, 520, 524}


_ANTHROPIC_MODEL_ALIASES = {
    "anthropic/claude-opus-4-7": "anthropic/claude-opus-4.7",
    "anthropic/claude-opus-4-6": "anthropic/claude-opus-4.6",
    "anthropic/claude-opus-4-5": "anthropic/claude-opus-4.5",
    "anthropic/claude-opus-4-1": "anthropic/claude-opus-4.1",
    "anthropic/claude-sonnet-4-6": "anthropic/claude-sonnet-4.6",
    "anthropic/claude-sonnet-4-5": "anthropic/claude-sonnet-4.5",
    "anthropic/claude-haiku-4-5": "anthropic/claude-haiku-4.5",
}
_MODELS_WITHOUT_TEMPERATURE = {
    "anthropic/claude-opus-4.7",
}


def normalize_zenmux_model(model: str) -> str:
    """Convert adapter-facing model IDs to ZenMux OpenAI API model IDs."""
    normalized = model.strip()
    if normalized.startswith("zenmux-anthropic/"):
        normalized = "anthropic/" + normalized.removeprefix("zenmux-anthropic/")
    elif normalized.startswith("zenmux/"):
        normalized = normalized.removeprefix("zenmux/")
    return _ANTHROPIC_MODEL_ALIASES.get(normalized, normalized)


def model_accepts_temperature(model: str) -> bool:
    """Return whether the ZenMux API accepts `temperature` for this model."""
    return normalize_zenmux_model(model) not in _MODELS_WITHOUT_TEMPERATURE


@dataclass
class UsageStats:
    """聚合一次脚本运行的 token 消耗。"""

    prompt_tokens: int = 0
    completion_tokens: int = 0
    cache_creation_tokens: int = 0
    cache_read_tokens: int = 0
    call_count: int = 0
    failed_calls: int = 0
    by_model: dict[str, dict[str, int]] = field(default_factory=dict)

    def add(self, model: str, usage: dict[str, Any]) -> None:
        p = usage.get("prompt_tokens", 0) or 0
        c = usage.get("completion_tokens", 0) or 0
        cc = (
            usage.get("cache_creation_input_tokens", 0)
            or usage.get("prompt_tokens_details", {}).get("cached_tokens", 0)
            or 0
        )
        cr = (
            usage.get("cache_read_input_tokens", 0)
            or 0
        )
        self.prompt_tokens += p
        self.completion_tokens += c
        self.cache_creation_tokens += cc
        self.cache_read_tokens += cr
        self.call_count += 1
        m = self.by_model.setdefault(
            model,
            {"prompt": 0, "completion": 0, "cache_creation": 0, "cache_read": 0, "calls": 0},
        )
        m["prompt"] += p
        m["completion"] += c
        m["cache_creation"] += cc
        m["cache_read"] += cr
        m["calls"] += 1

    def summary(self) -> str:
        lines = [
            f"Total calls: {self.call_count}  (failed: {self.failed_calls})",
            f"Prompt tokens:     {self.prompt_tokens:>12,}",
            f"Completion tokens: {self.completion_tokens:>12,}",
            f"Cache creation:    {self.cache_creation_tokens:>12,}",
            f"Cache read:        {self.cache_read_tokens:>12,}",
        ]
        for model, s in self.by_model.items():
            lines.append(
                f"  [{model}] calls={s['calls']} in={s['prompt']:,} "
                f"out={s['completion']:,} cache_w={s['cache_creation']:,} "
                f"cache_r={s['cache_read']:,}"
            )
        return "\n".join(lines)


class ZenMuxError(RuntimeError):
    pass


class ZenMuxClient:
    """ZenMux 轻量客户端。

    只暴露一个方法 `chat_complete()`，屏蔽 httpx 细节。
    """

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: float = DEFAULT_TIMEOUT,
        log_file: Path | None = None,
    ) -> None:
        self.api_key = api_key or os.environ.get("ZENMUX_API_KEY")
        if not self.api_key:
            raise ZenMuxError(
                "ZENMUX_API_KEY not set. Source secrets.env or pass api_key explicitly."
            )
        self.base_url = (base_url or os.environ.get("ZENMUX_BASE_URL") or DEFAULT_BASE_URL).rstrip("/")
        self.timeout = timeout
        self.log_file = log_file
        self.usage = UsageStats()
        self._client = httpx.Client(timeout=timeout)
        self._log_lock = threading.Lock()
        self._usage_lock = threading.Lock()

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "ZenMuxClient":
        return self

    def __exit__(self, *_args: Any) -> None:
        self.close()

    def _log(self, payload: dict[str, Any]) -> None:
        if not self.log_file:
            return
        self.log_file.parent.mkdir(parents=True, exist_ok=True)
        with self._log_lock:
            with self.log_file.open("a", encoding="utf-8") as f:
                f.write(json.dumps(payload, ensure_ascii=False) + "\n")

    def chat_complete(
        self,
        model: str,
        system: str,
        user: str,
        *,
        temperature: float = 0.3,
        max_tokens: int = 16000,
        extra_messages: list[dict[str, str]] | None = None,
        web_search: bool = False,
        web_search_options: dict[str, Any] | None = None,
        tag: str = "",
    ) -> str:
        """一次非流式对话补全。

        Args:
            model: 完整 model id，例如 `anthropic/claude-sonnet-4.6`（zenmux slug 不带 `zenmux/` 前缀，因为 baseURL 已经定位到 zenmux）。
            system: system prompt
            user: user message
            temperature, max_tokens: 常规参数
            extra_messages: 插在 system 之后、user 之前的额外消息（few-shot 等）
            tag: 给这次调用打标签，便于日志里识别（如 "translate:ch03"）

        Returns:
            assistant 的纯文本内容。如失败抛 ZenMuxError。
        """
        messages: list[dict[str, str]] = [{"role": "system", "content": system}]
        if extra_messages:
            messages.extend(extra_messages)
        messages.append({"role": "user", "content": user})

        api_model = normalize_zenmux_model(model)
        body: dict[str, Any] = {
            "model": api_model,
            "messages": messages,
            "max_tokens": max_tokens,
        }
        if model_accepts_temperature(api_model):
            body["temperature"] = temperature
        if web_search:
            body["web_search_options"] = web_search_options or {}
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        url = f"{self.base_url}/chat/completions"

        last_error: str = ""
        for attempt in range(MAX_RETRIES):
            t0 = time.time()
            try:
                resp = self._client.post(url, json=body, headers=headers)
                elapsed = time.time() - t0
            except httpx.RequestError as e:
                last_error = f"network: {e}"
                elapsed = time.time() - t0
                self._log({
                    "tag": tag, "attempt": attempt, "elapsed": elapsed,
                    "error": last_error,
                })
                time.sleep(2 ** attempt)
                continue

            if resp.status_code == 200:
                try:
                    data = resp.json()
                except Exception as e:
                    raise ZenMuxError(f"invalid JSON from zenmux: {e}; body={resp.text[:500]}")
                usage = data.get("usage", {}) or {}
                with self._usage_lock:
                    self.usage.add(api_model, usage)
                content = ""
                choices = data.get("choices") or []
                if choices:
                    msg = choices[0].get("message") or {}
                    content = msg.get("content") or ""
                self._log({
                    "tag": tag, "model": api_model, "requested_model": model, "attempt": attempt,
                    "elapsed": round(elapsed, 2),
                    "usage": usage,
                    "out_chars": len(content),
                    "status": 200,
                })
                if not content.strip():
                    # zenmux 偶尔返 200 但 content 空；视作可重试
                    last_error = "empty content"
                    time.sleep(2 ** attempt)
                    continue
                return content

            # 非 200
            retryable = resp.status_code in RETRYABLE_STATUSES
            last_error = f"HTTP {resp.status_code}: {resp.text[:500]}"
            self._log({
                "tag": tag, "attempt": attempt, "elapsed": round(elapsed, 2),
                "status": resp.status_code, "error": last_error,
                "retryable": retryable,
            })
            if not retryable:
                self.usage.failed_calls += 1
                raise ZenMuxError(last_error)
            sleep_for = min(60, (2 ** attempt) + (attempt * 0.5))
            time.sleep(sleep_for)

        self.usage.failed_calls += 1
        raise ZenMuxError(f"max retries exhausted. last error: {last_error}")

    def chat_complete_with_meta(
        self,
        model: str,
        system: str,
        user: str,
        *,
        temperature: float = 0.3,
        max_tokens: int = 16000,
        extra_messages: list[dict[str, str]] | None = None,
        web_search: bool = False,
        web_search_options: dict[str, Any] | None = None,
        tag: str = "",
    ) -> dict[str, Any]:
        """Return content plus metadata from one completion call."""
        messages: list[dict[str, str]] = [{"role": "system", "content": system}]
        if extra_messages:
            messages.extend(extra_messages)
        messages.append({"role": "user", "content": user})

        api_model = normalize_zenmux_model(model)
        body: dict[str, Any] = {
            "model": api_model,
            "messages": messages,
            "max_tokens": max_tokens,
        }
        if model_accepts_temperature(api_model):
            body["temperature"] = temperature
        if web_search:
            body["web_search_options"] = web_search_options or {}

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        url = f"{self.base_url}/chat/completions"

        last_error = ""
        for attempt in range(MAX_RETRIES):
            t0 = time.time()
            try:
                resp = self._client.post(url, json=body, headers=headers)
                elapsed = time.time() - t0
            except httpx.RequestError as e:
                last_error = f"network: {e}"
                elapsed = time.time() - t0
                self._log({"tag": tag, "attempt": attempt, "elapsed": elapsed, "error": last_error})
                time.sleep(2 ** attempt)
                continue

            if resp.status_code != 200:
                retryable = resp.status_code in RETRYABLE_STATUSES
                last_error = f"HTTP {resp.status_code}: {resp.text[:500]}"
                self._log({
                    "tag": tag,
                    "attempt": attempt,
                    "elapsed": round(elapsed, 2),
                    "status": resp.status_code,
                    "error": last_error,
                    "retryable": retryable,
                })
                if not retryable:
                    self.usage.failed_calls += 1
                    raise ZenMuxError(last_error)
                sleep_for = min(60, (2 ** attempt) + (attempt * 0.5))
                time.sleep(sleep_for)
                continue

            try:
                data = resp.json()
            except Exception as e:
                raise ZenMuxError(f"invalid JSON from zenmux: {e}; body={resp.text[:500]}")

            usage = data.get("usage", {}) or {}
            with self._usage_lock:
                self.usage.add(api_model, usage)

            message = ((data.get("choices") or [{}])[0].get("message") or {})
            content = message.get("content") or ""
            annotations = message.get("annotations") or []
            urls: list[str] = []
            for ann in annotations:
                if not isinstance(ann, dict):
                    continue
                citation = ann.get("url_citation") or {}
                url_item = citation.get("url")
                if url_item:
                    urls.append(url_item)
            self._log({
                "tag": tag,
                "model": api_model,
                "requested_model": model,
                "attempt": attempt,
                "elapsed": round(elapsed, 2),
                "usage": usage,
                "out_chars": len(content),
                "status": 200,
                "web_search": web_search,
                "citations": len(urls),
            })
            if not content.strip():
                last_error = "empty content"
                time.sleep(2 ** attempt)
                continue
            return {
                "content": content,
                "usage": usage,
                "citations": urls,
                "raw": data,
            }

        self.usage.failed_calls += 1
        raise ZenMuxError(f"max retries exhausted. last error: {last_error}")


def load_secrets(env_path: Path | None = None) -> None:
    """从 secrets.env 把 key 塞到 os.environ，便于脚本直接运行。

    格式宽松：`KEY=VALUE` 每行一条，`#` 开头是注释，忽略空行。
    """
    if env_path is None:
        # 默认在仓库根找 secrets.env
        here = Path(__file__).resolve()
        for parent in [here.parent, *here.parents]:
            cand = parent / "secrets.env"
            if cand.exists():
                env_path = cand
                break
    if not env_path or not env_path.exists():
        return
    for line in env_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, _, v = line.partition("=")
        k = k.strip()
        v = v.strip().strip('"').strip("'")
        if k and v and k not in os.environ:
            os.environ[k] = v


if __name__ == "__main__":
    # 冒烟测试：python -m scripts.lib.zenmux_client
    load_secrets()
    with ZenMuxClient() as c:
        out = c.chat_complete(
            model="anthropic/claude-haiku-4.5",
            system="You reply in exactly one English word.",
            user="Say hello.",
            max_tokens=20,
            tag="smoke",
        )
        print("reply:", out)
        print(c.usage.summary(), file=sys.stderr)
