"""Client for calling an LLM directly via Groq, bypassing both the EXL AI Gateway (blocked by
its own WAF from this AWS account's egress IP, confirmed by direct testing — not an assumption)
and AWS Bedrock (Anthropic models aren't enabled in this account without a separate use-case
form that only an account admin can submit).

Groq exposes a native OpenAI-compatible endpoint directly (api.groq.com/openai/v1), so this
uses the OpenAI SDK pointed at it — the exact same pattern the original gateway client used,
just a different base_url. litellm was tried first as a multi-provider translation layer, but
it produced an intermittent, unreproducible "model_not_found" error for an unrelated,
already-deprecated model name mid-request — calling Groq's own endpoint directly removes that
extra routing layer and its failure mode entirely.
"""
import os
import re
import time

import openai
from openai import OpenAI

MODEL_ID = os.environ.get("MODEL_ID", "openai/gpt-oss-120b")


def load_env(path=None):
    """Minimal .env loader so this runs without an extra dependency for one file."""
    path = path or os.path.join(os.path.dirname(__file__), ".env")
    if not os.path.exists(path):
        return
    for line in open(path):
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        os.environ.setdefault(key.strip(), value.strip())


load_env()

GROQ_API_KEY = os.environ.get("GROQ_API_KEY")


def _parse_retry_after(error_text):
    """Groq's 429 body says e.g. 'Please try again in 892.5ms' or '...in 1.2s' — use that
    exact wait instead of a guessed backoff."""
    match = re.search(r"try again in ([\d.]+)(ms|s)\b", error_text)
    if not match:
        return None
    value, unit = float(match.group(1)), match.group(2)
    seconds = value / 1000 if unit == "ms" else value
    return seconds + 0.2  # small safety margin


class GatewayClient:
    """Named GatewayClient for drop-in compatibility with agent_core.py's existing calls —
    it no longer talks to the EXL gateway, but the request/response shape agent_core.py
    depends on (message.content, message.tool_calls[].function.{name,arguments}) is
    unchanged, so nothing else needed to change when this swap happened.
    """

    def __init__(self, model=MODEL_ID, api_key=GROQ_API_KEY):
        if not api_key:
            raise RuntimeError("GROQ_API_KEY must be set (see .env)")
        self.model = model
        self._client = OpenAI(base_url="https://api.groq.com/openai/v1", api_key=api_key)

    def complete(self, messages, tools=None, tool_choice=None, **kwargs):
        kwargs = dict(kwargs)
        if tools:
            kwargs["tools"] = tools
            if tool_choice:
                kwargs["tool_choice"] = tool_choice
        # Fact extraction and verification judging are classification tasks, not creative
        # writing — temperature=0 found necessary after the real multi-turn flow intermittently
        # failed to extract a fact (vulnerability_consent) that an isolated identical-input call
        # extracted correctly, on the default (non-zero) sampling temperature.
        kwargs.setdefault("temperature", 0)
        # Groq's free tier caps tokens-per-minute (8000 TPM observed) — this agent makes
        # several calls per turn with sizeable schemas, so a real multi-turn conversation
        # run can hit that cap mid-conversation. Observed directly: a 429 here, left
        # unhandled, otherwise looks indistinguishable from a model reasoning failure
        # (a fact silently not getting extracted) rather than the capacity issue it is —
        # worth retrying with the wait time Groq itself reports, not guessing one.
        last_exc = None
        for attempt in range(5):
            try:
                response = self._client.chat.completions.create(
                    model=self.model, messages=messages, **kwargs)
                return response.choices[0].message
            except openai.RateLimitError as exc:
                last_exc = exc
                wait_s = _parse_retry_after(str(exc)) or (2 * (attempt + 1))
                time.sleep(wait_s)
        raise last_exc