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

from .client import Client
from .vendor import default_ep, query_key_name
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

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 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)

        if not payload.get("stream"):
            return self.cli.post(path, payload).json()

        # stream: 仍按 OpenAI SSE 逻辑（原来的 stream_post 简化合并到这里）
        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)

        # endpoint 渲染
        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)

        # 鉴权：你的代理无鉴权 => 不强制带 key
        # 若未来需要：可以在 Client(query={"key":xxx}) 或 headers={"x-goog-api-key":xxx} 设置
        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)

        # stream：兼容 SSE/JSONL/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
