"""
Groq API Key Rotation Utility
──────────────────────────────────────────────────────────────────────────────
Loads ALL env vars named GROQ_API_KEY or GROQ_API_KEY_* at startup.
When any route hits a 429 rate-limit error, calling rotate() silently
advances to the next key. Users never see a rate-limit error as long as
at least one key still has remaining quota.

Rotation is round-robin and thread-safe (uses a threading.Lock so concurrent
FastAPI workers don't race on the index).

Usage:
    from utils.groq_client import call_groq_json, HEAVY_MODEL, LIGHT_MODEL

    data = call_groq_json(system_prompt, user_prompt, json_mode=True)
"""

import os
import json
import time
import threading
from pathlib import Path
from types import SimpleNamespace
from groq import Groq

# ─── E2E fixture mode ─────────────────────────────────────────────────────────
# When E2E_FIXTURE_MODE=1, every call below returns a canned fixture instead of
# hitting the real Groq API — no cost, deterministic, works offline. Used by
# the E2E/surface-discovery test suite (see docs/qa/test-strategy.md and the
# project-e2e-testing-design memory), never by production or by the unit tests
# in backend/tests/ (those stay pure-function tests with no I/O).
# Each call site passes a `fixture` key identifying which canned response it
# wants; fixture mode with no key given is a bug at the call site, not a
# silent fallback, so it raises loudly instead of guessing.
FIXTURE_MODE = os.environ.get("E2E_FIXTURE_MODE") == "1"
_FIXTURES_DIR = Path(__file__).resolve().parent.parent / "tests" / "fixtures" / "groq"


# DEMO_MODE=1 (with fixture mode, used by `npm run demo`): every AI answer reads
# "Test Environment". Only fields the app needs to keep working are left alone.
DEMO_MODE = FIXTURE_MODE and os.environ.get("DEMO_MODE") == "1"
DEMO_TEXT = "Test Environment"
_DEMO_KEEP_KEYS = {"id", "type", "column", "targetDate", "date", "status"}


def _demo_value(value, key=None):
    if isinstance(value, dict):
        return {k: _demo_value(v, k) for k, v in value.items()}
    if isinstance(value, list):
        strings = [v for v in value if isinstance(v, str)]
        if strings and len(strings) == len(value):
            # Numbered so choice lists and React keys stay distinct
            return [f"{DEMO_TEXT} {i + 1}" for i in range(len(value))] if len(value) > 1 else [DEMO_TEXT]
        return [_demo_value(v, key) for v in value]
    if isinstance(value, str) and key not in _DEMO_KEEP_KEYS:
        return DEMO_TEXT
    return value


def _load_json_fixture(name: str) -> dict:
    path = _FIXTURES_DIR / f"{name}.json"
    if not path.exists():
        raise RuntimeError(
            f"[E2E_FIXTURE_MODE] no fixture file for '{name}' at {path} — "
            "add one under backend/tests/fixtures/groq/"
        )
    data = json.loads(path.read_text(encoding="utf-8"))
    return _demo_value(data) if DEMO_MODE else data


def _load_text_fixture(name: str) -> str:
    path = _FIXTURES_DIR / f"{name}.txt"
    if not path.exists():
        raise RuntimeError(
            f"[E2E_FIXTURE_MODE] no fixture file for '{name}' at {path} — "
            "add one under backend/tests/fixtures/groq/"
        )
    return DEMO_TEXT if DEMO_MODE else path.read_text(encoding="utf-8")

# ─── Load all keys ────────────────────────────────────────────────────────────
# Accepts: GROQ_API_KEY, GROQ_API_KEY_1 ... GROQ_API_KEY_N (any suffix)
_KEYS: list[str] = []
for k in sorted(os.environ.keys()):
    if k == 'GROQ_API_KEY' or k.startswith('GROQ_API_KEY_'):
        v = os.environ[k].strip()
        if v:
            _KEYS.append(v)

if not _KEYS:
    raise RuntimeError(
        "No Groq API keys found. "
        "Set GROQ_API_KEY (and optionally GROQ_API_KEY_7, etc.) in your .env file."
    )

print(f"[groq] {len(_KEYS)} API key(s) loaded for rotation")

# ─── Model names ──────────────────────────────────────────────────────────────
# llama-3.3-70b-versatile / llama-3.1-8b-instant were retired by Groq — replaced
# with the current openai/gpt-oss lineup (see Groq's models page for the live list).
HEAVY_MODEL = "openai/gpt-oss-120b"   # reasoning, profile generation, consultation
LIGHT_MODEL  = "openai/gpt-oss-20b"   # chat, quick responses
TRANSCRIBE_MODEL = "whisper-large-v3-turbo"

# ─── Round-robin state ────────────────────────────────────────────────────────
_idx  = 0
_lock = threading.Lock()


def _current_key() -> str:
    return _KEYS[_idx % len(_KEYS)]


def _rotate() -> str:
    """Advance to the next key and return it. Thread-safe."""
    global _idx
    with _lock:
        _idx = (_idx + 1) % len(_KEYS)
        new_key = _KEYS[_idx]
    print(f"[groq] Rotated to key index {_idx % len(_KEYS)} of {len(_KEYS)}")
    return new_key


# ─── Core call ────────────────────────────────────────────────────────────────
def call_groq(
    system_prompt: str,
    user_prompt: str,
    model: str = HEAVY_MODEL,
    json_mode: bool = False,
    max_tokens: int = 1500,
    temperature: float = 0.7,
) -> str:
    """
    Calls Groq with automatic key rotation on 429 errors.
    Tries every available key before giving up.
    Raises the last error if all keys are exhausted.

    Returns the raw string content of the first choice message.
    """
    last_error: Exception | None = None

    # Try each key in the rotation at most once, plus one extra attempt
    # in case the first key has recovered by the time we loop around.
    for attempt in range(len(_KEYS) + 1):
        try:
            client = Groq(api_key=_current_key())
            kwargs: dict = dict(
                model=model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user",   "content": user_prompt},
                ],
                temperature=temperature,
                max_tokens=max_tokens,
            )
            if json_mode:
                kwargs["response_format"] = {"type": "json_object"}

            response = client.chat.completions.create(**kwargs)
            return response.choices[0].message.content  # type: ignore

        except Exception as e:
            last_error = e
            err_str = str(e)
            is_rate_limit = (
                '429' in err_str
                or 'rate_limit' in err_str.lower()
                or 'rate limit' in err_str.lower()
            )

            if is_rate_limit and attempt < len(_KEYS):
                # Rotate and pause briefly before retrying
                _rotate()
                time.sleep(1)
                continue

            # Non-rate-limit errors bubble up immediately
            raise

    # All keys exhausted
    raise last_error or RuntimeError("All Groq API keys are rate-limited")


def call_groq_transcribe(
    file: tuple,
    model: str = TRANSCRIBE_MODEL,
    response_format: str = "json",
    language: str | None = None,
    fixture: str | None = None,
):
    """
    Calls Groq's Whisper transcription endpoint with the same key-rotation
    on 429 errors as the chat completion calls above.

    `file` is the (filename, bytes, content_type) tuple the Groq SDK expects.
    """
    if FIXTURE_MODE:
        if not fixture:
            raise RuntimeError(
                "[E2E_FIXTURE_MODE] call_groq_transcribe was called with no fixture key — "
                "every call site must pass one when fixture mode is on"
            )
        data = _load_json_fixture(fixture)
        return SimpleNamespace(text=data.get("text", ""))

    last_error: Exception | None = None

    for attempt in range(len(_KEYS) + 1):
        try:
            client = Groq(api_key=_current_key())
            kwargs: dict = dict(model=model, file=file, response_format=response_format)
            if language:
                kwargs["language"] = language
            return client.audio.transcriptions.create(**kwargs)

        except Exception as e:
            last_error = e
            err_str = str(e)
            is_rate_limit = (
                '429' in err_str
                or 'rate_limit' in err_str.lower()
                or 'rate limit' in err_str.lower()
            )

            if is_rate_limit and attempt < len(_KEYS):
                _rotate()
                time.sleep(1)
                continue

            raise

    raise last_error or RuntimeError("All Groq API keys are rate-limited")


def call_groq_json(
    system_prompt: str,
    user_prompt: str,
    model: str = HEAVY_MODEL,
    max_tokens: int = 1500,
    fixture: str | None = None,
) -> dict:
    """Convenience wrapper: calls Groq in JSON mode and returns parsed dict."""
    if FIXTURE_MODE:
        if not fixture:
            raise RuntimeError(
                "[E2E_FIXTURE_MODE] call_groq_json was called with no fixture key — "
                "every call site must pass one when fixture mode is on"
            )
        return _load_json_fixture(fixture)

    raw = call_groq(system_prompt, user_prompt, model=model,
                    json_mode=True, max_tokens=max_tokens)
    if not raw:
        # Reasoning models (e.g. openai/gpt-oss-*) can spend the token budget on
        # internal reasoning and return empty content if max_tokens is too tight.
        raise RuntimeError(
            "Groq returned empty content — likely truncated by max_tokens "
            f"(model={model}, max_tokens={max_tokens})"
        )
    try:
        return json.loads(raw)
    except json.JSONDecodeError as e:
        raise RuntimeError(f"Groq returned invalid JSON (likely truncated): {e}") from e


def _fixture_stream(text: str):
    """Yields fake chunks shaped like the Groq SDK's streaming response —
    chunk.choices[0].delta.content — so callers don't need to know they're
    in fixture mode."""
    for word in text.split(" "):
        yield SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=word + " "))])


def call_groq_stream(
    system_prompt: str,
    user_prompt: str,
    model: str = LIGHT_MODEL,
    max_tokens: int = 512,
    fixture: str | None = None,
    messages: list[dict] | None = None,
):
    """
    Returns a Groq streaming completion object.
    Rotates key on 429 like all other calls.

    `messages`, when given, is the full chat list (system prompt first) and
    replaces system_prompt/user_prompt — used by chat to send conversation
    history. Callers without history keep passing the two prompt strings.
    """
    if FIXTURE_MODE:
        if not fixture:
            raise RuntimeError(
                "[E2E_FIXTURE_MODE] call_groq_stream was called with no fixture key — "
                "every call site must pass one when fixture mode is on"
            )
        return _fixture_stream(_load_text_fixture(fixture))

    last_error: Exception | None = None
    for attempt in range(len(_KEYS) + 1):
        try:
            client = Groq(api_key=_current_key())
            return client.chat.completions.create(
                model=model,
                messages=messages or [
                    {"role": "system", "content": system_prompt},
                    {"role": "user",   "content": user_prompt},
                ],
                temperature=0.7,
                max_tokens=max_tokens,
                stream=True,
            )
        except Exception as e:
            last_error = e
            if ('429' in str(e) or 'rate_limit' in str(e).lower()) and attempt < len(_KEYS):
                _rotate()
                time.sleep(1)
                continue
            raise
    raise last_error or RuntimeError("All Groq API keys are rate-limited")