from __future__ import annotations
from typing import Callable, Any, Optional
import json

from .client import Client
from .vendor import default_ep, capabilities
from .exceptions import ValidationError, LayBotError
from .base import Base, looks_like_laybot
from .decoder import iter_frames
from .gemini import build_gemini_payload, gemini_to_openai_like, gemini_text


def _usage_norm_from_openai_like(resp_json: dict) -> dict:
    """
    OpenAI-like usage 归一化（不破坏原 usage 字段）。
    兼容 prompt caching:
      usage.prompt_tokens_details.cached_tokens
    """
    usage = resp_json.get("usage") or {}
    prompt = usage.get("prompt_tokens")
    completion = usage.get("completion_tokens")
    total = usage.get("total_tokens")

    cached = 0
    ptd = usage.get("prompt_tokens_details")
    if isinstance(ptd, dict) and isinstance(ptd.get("cached_tokens"), int):
        cached = int(ptd["cached_tokens"])

    return {
        "prompt_tokens": int(prompt) if isinstance(prompt, int) else None,
        "completion_tokens": int(completion) if isinstance(completion, int) else None,
        "cached_prompt_tokens": cached,
        "total_tokens": int(total) if isinstance(total, int) else None,
        "vendor": "openai_like",
    }


class Chat:
    def __init__(self, api_key: str | Client,
                 *, vendor: str = "laybot", base: str | None = None,
                 timeout: dict | None = None, proxies: dict | None = None, verify=None,
                 query: dict | None = None, headers: dict | None = None):
        self.vendor = vendor.lower()
        self.cli: Client = api_key if isinstance(api_key, Client) else Client(
            api_key, vendor=self.vendor, base=base, timeout=timeout,
            proxies=proxies, verify=verify, query=query, headers=headers
        )
        self.is_laybot = looks_like_laybot(self.cli.base) or self.vendor == "laybot"
        self.base = Base(self.cli, is_laybot=self.is_laybot)

    def capabilities(self) -> dict:
        """返回 vendor 能力矩阵（目前重点是 cache）。"""
        return capabilities(self.vendor)

    def completions(self, body: dict,
                    on_stream: Optional[Callable[[dict, bool], Any]] = None) -> Optional[dict]:
        if "model" not in body or "messages" not in body:
            raise ValidationError("model & messages required")

        if self.vendor == "gemini":
            return self._gemini(body, on_stream=on_stream)

        # openai/laybot compatible
        def_path = default_ep(self.vendor, "chat") or "/chat/completions"
        path, payload = self.base.ready(body, "chat", def_path)

        # OpenAI prompt caching 相关参数：仅透传，不强校验
        # - prompt_cache_key / prompt_cache_retention
        # 调用方可直接放在 body 内，ready() 只负责 capability 等
        if not payload.get("stream"):
            rsp = self.cli.post(path, payload).json()
            if isinstance(rsp, dict) and "usage_norm" not in rsp:
                rsp["usage_norm"] = _usage_norm_from_openai_like(rsp)
            return rsp

        # stream: 仍按 OpenAI SSE 逻辑
        resp = self.cli.post(path, payload, stream=True)
        for typ, raw in iter_frames(resp):
            if typ == "done":
                on_stream and on_stream({}, True)
                break
            try:
                payload_json = json.loads(raw)
            except Exception:
                continue
            on_stream and on_stream(payload_json, False)
        return None

    def _gemini(self, body: dict, on_stream=None) -> Optional[dict]:
        model = body["model"]
        stream = bool(body.get("stream"))
        payload = build_gemini_payload(body)

        ep_key = "chat_stream" if stream else "chat"
        ep_tpl = default_ep("gemini", ep_key)
        if not ep_tpl:
            raise LayBotError("gemini endpoint missing in vendor table")
        path = ep_tpl.replace("{MODEL}", model)

        params = body.get("query", None)  # 允许单次请求覆盖 query
        if params is None:
            params = {}

        if not stream:
            resp_json = self.cli.post(path, payload, params=params).json()
            return gemini_to_openai_like(resp_json)

        resp = self.cli.post(path, payload, stream=True, params=params)

        def emit_delta(text: str):
            if not on_stream or not text:
                return
            on_stream({"choices": [{"index": 0, "delta": {"content": text}}]}, False)

        for typ, raw in iter_frames(resp):
            if typ == "done":
                on_stream and on_stream({}, True)
                break
            try:
                frame = json.loads(raw)
            except Exception:
                continue
            text = gemini_text(frame)
            emit_delta(text)

        return None
