"""Provider-agnostic LLM access (standard library only).

Groq   : very fast, good for the trade-by-trade second opinion.
Gemini : can search the web through Google Search grounding, so it is the research brain.

Models are DISCOVERED from each provider's model list (names change often), preferring the ones in config.
A provider that keeps failing is put in a 5 minute cool-down and the next one is used.
API keys only ever travel in request headers. They are never put in URLs, logs or error text.
"""
import json
import re
import threading
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field


class LLMError(Exception):
    pass


@dataclass
class Reply:
    text: str
    provider: str
    model: str
    latency_ms: int
    sources: list = field(default_factory=list)


def _http(method, url, headers, body, timeout):
    data = json.dumps(body).encode("utf-8") if body is not None else None
    h = {"Content-Type": "application/json", "User-Agent": "sailar-x/0.2 (+python-urllib)", **headers}
    req = urllib.request.Request(url, data, h, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return json.loads(r.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        raise LLMError(f"HTTP {e.code}: {e.read().decode('utf-8', 'replace')[:240]}") from None
    except Exception as e:
        raise LLMError(f"{type(e).__name__}: {e}") from None


def parse_json(text):
    """Pull a JSON object out of a model reply (handles ```json fences and chatter around it)."""
    t = (text or "").strip()
    t = re.sub(r"^```(?:json)?\s*|\s*```$", "", t, flags=re.I).strip()
    try:
        return json.loads(t)
    except Exception:
        pass
    a, b = t.find("{"), t.rfind("}")
    if 0 <= a < b:
        try:
            return json.loads(t[a:b + 1])
        except Exception:
            pass
    raise LLMError("the model did not return valid JSON")


EXCLUDE = re.compile(r"whisper|guard|tts|embed|image|audio|live|imagen|veo|robotics|moderation|safeguard|computer-use|aqa", re.I)


def _nums(name):
    return tuple(int(x) for x in re.findall(r"\d+", name)[:3]) or (0,)


def _pick(ids, prefer, bad=()):
    avail = [i for i in ids if i not in bad]
    for p in prefer:
        if p in avail:
            return p
    cands = [i for i in avail if not EXCLUDE.search(i)]
    if not cands:
        raise LLMError("no usable model on this account")
    return sorted(cands, key=lambda n: (("flash" in n or "versatile" in n or "instant" in n), "lite" not in n, _nums(n)), reverse=True)[0]


class _Provider:
    name = ""

    def __init__(self, key, prefer):
        self.key, self.prefer = key, list(prefer or [])
        self._model, self.bad = None, set()

    def _note_failure(self, err):
        if re.search(r"model|decommission|not found|does not exist|deprecat", str(err), re.I) and self._model:
            self.bad.add(self._model)
            self._model = None     # re-discover next time


class Groq(_Provider):
    name = "groq"
    BASE = "https://api.groq.com/openai/v1"

    def _hdr(self):
        return {"Authorization": f"Bearer {self.key}"}

    def model(self, timeout):
        if not self._model:
            ids = [m["id"] for m in _http("GET", self.BASE + "/models", self._hdr(), None, timeout).get("data", [])]
            self._model = _pick(ids, self.prefer, self.bad)
        return self._model

    def chat(self, system, user, json_mode, timeout, max_tokens, temperature, search=False):
        model = self.model(timeout)
        body = {"model": model, "temperature": temperature, "max_tokens": max_tokens,
                "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}]}
        if json_mode:
            body["response_format"] = {"type": "json_object"}
        try:
            r = _http("POST", self.BASE + "/chat/completions", self._hdr(), body, timeout)
        except LLMError as e:
            self._note_failure(e)
            raise
        try:
            return r["choices"][0]["message"]["content"] or "", model, []
        except Exception:
            raise LLMError("unexpected reply shape from Groq") from None


class Gemini(_Provider):
    name = "gemini"
    BASE = "https://generativelanguage.googleapis.com/v1beta"

    def _hdr(self):
        return {"x-goog-api-key": self.key}

    def model(self, timeout):
        if not self._model:
            data = _http("GET", self.BASE + "/models?pageSize=200", self._hdr(), None, timeout)
            ids = [m["name"].split("/", 1)[-1] for m in data.get("models", [])
                   if "generateContent" in m.get("supportedGenerationMethods", [])]
            self._model = _pick(ids, self.prefer, self.bad)
        return self._model

    def chat(self, system, user, json_mode, timeout, max_tokens, temperature, search=False):
        model = self.model(timeout)
        cfg = {"temperature": temperature, "maxOutputTokens": max_tokens}
        body = {"systemInstruction": {"parts": [{"text": system}]},
                "contents": [{"role": "user", "parts": [{"text": user}]}], "generationConfig": cfg}
        if search:
            body["tools"] = [{"google_search": {}}]      # Google Search grounding (cannot be combined with JSON mode)
        elif json_mode:
            cfg["responseMimeType"] = "application/json"
        try:
            r = _http("POST", f"{self.BASE}/models/{model}:generateContent", self._hdr(), body, timeout)
        except LLMError as e:
            self._note_failure(e)
            raise
        try:
            cand = r["candidates"][0]
            text = "".join(p.get("text", "") for p in cand.get("content", {}).get("parts", []) if not p.get("thought"))
        except Exception:
            raise LLMError("Gemini returned no answer (blocked or empty)") from None
        if not text.strip():
            raise LLMError(f"Gemini returned empty text (finish: {cand.get('finishReason')})")
        sources = []
        for ch in (cand.get("groundingMetadata") or {}).get("groundingChunks", []):
            w = ch.get("web") or {}
            if w.get("uri"):
                sources.append({"title": (w.get("title") or "")[:120], "uri": w["uri"]})
        return text, model, sources


class Router:
    def __init__(self, bc, env):
        self.bc, self.env, self.providers = bc, env, {}
        models = bc.get("models", {})
        if env.get("GROQ_API_KEY"):
            self.providers["groq"] = Groq(env["GROQ_API_KEY"], models.get("groq"))
        gk = env.get("GEMINI_API_KEY") or env.get("GOOGLE_API_KEY")
        if gk:
            self.providers["gemini"] = Gemini(gk, models.get("gemini"))
        self.order = [p for p in bc.get("providers", ["groq", "gemini"]) if p in self.providers]
        self.fails, self.until = {}, {}
        self.calls_today, self.day = 0, None
        self.last_error = ""
        self.lock = threading.Lock()

    def available(self):
        return bool(self.order)

    def has(self, name):
        return name in self.providers

    def _budget(self):
        today = time.strftime("%Y-%m-%d", time.gmtime())
        if self.day != today:
            self.day, self.calls_today = today, 0
        return self.calls_today < self.bc.get("max_calls_per_day", 300)

    def chat(self, system, user, json_mode=True, timeout=None, max_tokens=700, temperature=0.2, prefer=None, search=False):
        timeout = timeout or self.bc.get("review_timeout", 8)
        order = list(self.order)
        if prefer in order:
            order = [prefer] + [o for o in order if o != prefer]
        if search:
            order = [o for o in order if o == "gemini"]
        errors = []
        for name in order:
            if time.time() < self.until.get(name, 0):
                errors.append(f"{name}: cooling down")
                continue
            with self.lock:
                if not self._budget():
                    raise LLMError("the daily AI call budget is used up")
                self.calls_today += 1
            p, t0 = self.providers[name], time.time()
            try:
                text, model, sources = p.chat(system, user, json_mode, timeout, max_tokens, temperature, search)
                self.fails[name] = 0
                return Reply(text, name, model, int((time.time() - t0) * 1000), sources)
            except LLMError as e:
                f = self.fails.get(name, 0) + 1
                self.fails[name] = f
                if f >= 3:
                    self.until[name] = time.time() + 300
                errors.append(f"{name}: {e}")
        self.last_error = "; ".join(errors)[:300] or "no AI provider is configured"
        raise LLMError(self.last_error)

    def health(self):
        out = {}
        for n, p in self.providers.items():
            cool = max(0, int(self.until.get(n, 0) - time.time()))
            out[n] = {"model": p._model, "fails": self.fails.get(n, 0), "cooling_s": cool,
                      "ok": cool == 0 and self.fails.get(n, 0) < 3}
        if self.env and self.env.get("SERPER_API_KEY"):
            out["serper"] = {"model": "google-search", "fails": 0, "cooling_s": 0, "ok": True}
        elif self.env and self.env.get("TAVILY_API_KEY"):
            out["tavily"] = {"model": "web-search", "fails": 0, "cooling_s": 0, "ok": True}
        return out