"""Provider layer — one OpenAI-compatible client covers Groq (and 15+ other vendors),
with a zeroclaw-style fallback chain and Qwen think-tag stripping."""
from __future__ import annotations

import json
import re
from typing import Any

from openai import AsyncOpenAI

from .config import Config

THINK_RE = re.compile(r"<think>.*?</think>\s*", re.DOTALL)


def strip_think(text: str) -> str:
    """Remove closed <think> blocks; salvage unclosed ones (model hit token limit mid-think)."""
    text = THINK_RE.sub("", text)
    if "<think>" in text:
        head, _, tail = text.partition("<think>")
        # keep whatever real answer surrounds the orphan block; prefer text before it
        text = head.strip() or tail.strip()
    # models occasionally parrot the legacy no-think control token
    return text.replace("/no_think", "").strip()


class LLM:
    def __init__(self, cfg: Config):
        self.cfg = cfg
        self.client = AsyncOpenAI(base_url=cfg.base_url, api_key=cfg.api_key)
        self.chain = [cfg.primary_model, *cfg.fallback_models]
        self._no_reasoning_param: set[str] = set()  # models that reject reasoning_effort

    async def chat(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        json_mode: bool = False,
        temperature: float = 0.3,
        max_tokens: int = 4096,
    ) -> Any:
        """Returns the raw completion message. Walks the fallback chain on errors."""
        last_err: Exception | None = None
        for model in self.chain:
            kwargs: dict = dict(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
            )
            if tools:
                kwargs["tools"] = tools
                kwargs["tool_choice"] = "auto"
            if json_mode:
                kwargs["response_format"] = {"type": "json_object"}
            # native reasoning off-switch beats prompt hints; fall back if the model rejects it
            attempts = [dict(kwargs, extra_body={"reasoning_effort": "none"}), kwargs] \
                if self.cfg.no_think and model not in self._no_reasoning_param else [kwargs]
            for i, attempt in enumerate(attempts):
                try:
                    resp = await self.client.chat.completions.create(**attempt)
                    msg = resp.choices[0].message
                    if msg.content:
                        msg.content = strip_think(msg.content)
                    return msg
                except Exception as e:
                    last_err = e
                    if len(attempts) == 2 and i == 0 and "reasoning" in str(e).lower():
                        self._no_reasoning_param.add(model)
                        continue  # retry same model without the param
                    break  # real failure — next model in chain
        raise RuntimeError(f"all models in chain failed: {last_err}") from last_err

    async def text(self, system: str, user: str, **kw) -> str:
        msg = await self.chat(
            [{"role": "system", "content": system}, {"role": "user", "content": user}], **kw
        )
        return msg.content or ""

    async def json(self, system: str, user: str, **kw) -> dict:
        raw = await self.text(system, user, json_mode=True, **kw)
        try:
            return json.loads(raw)
        except json.JSONDecodeError:
            m = re.search(r"\{.*\}", raw, re.DOTALL)
            if m:
                return json.loads(m.group(0))
            raise