"""Round-robin Gemini API key pool with failover on quota/rate-limit errors."""

from __future__ import annotations

import logging
import threading
from datetime import date
from typing import Callable, TypeVar

logger = logging.getLogger("docverify.gemini.keys")

T = TypeVar("T")

_lock = threading.Lock()
_index = 0
_daily_counts: dict[str, int] = {}
_daily_date: date | None = None

# gemini-3-flash free tier: 1000 RPD per key; warn at 90% usage
DAILY_WARN_THRESHOLD = 900

_RETRYABLE_MARKERS = (
    "429",
    "resource exhausted",
    "quota",
    "rate limit",
    "too many requests",
    "exceeded",
)


def get_api_keys() -> list[str]:
    from config import GEMINI_API_KEY, GEMINI_API_KEYS

    keys = list(GEMINI_API_KEYS)
    if GEMINI_API_KEY and GEMINI_API_KEY not in keys:
        keys.insert(0, GEMINI_API_KEY)
    return [k for k in keys if k]


def _mask_key(key: str) -> str:
    if len(key) <= 8:
        return "***"
    return f"{key[:4]}...{key[-4:]}"


def _bump_usage(key: str) -> None:
    global _daily_date
    today = date.today()
    with _lock:
        if _daily_date != today:
            _daily_counts.clear()
            _daily_date = today
        _daily_counts[key] = _daily_counts.get(key, 0) + 1
        count = _daily_counts[key]
        if count >= DAILY_WARN_THRESHOLD:
            logger.warning("Gemini key %s at %d requests today (free tier limit ~20/day)", _mask_key(key), count)


def _is_retryable(exc: Exception) -> bool:
    msg = str(exc).lower()
    return any(m in msg for m in _RETRYABLE_MARKERS)


def call_with_failover(fn: Callable[[str], T]) -> tuple[T | None, int | None, Exception | None]:
    """Try each API key in round-robin order. Returns (result, key_index, last_error).

    On quota/rate-limit errors: cycles through ALL remaining keys before failing.
    On non-retryable errors: fails immediately (wrong API key, invalid request, etc.).
    """
    keys = get_api_keys()
    if not keys:
        return None, None, ValueError("No Gemini API keys configured")

    global _index
    with _lock:
        start = _index % len(keys)
        _index = (_index + 1) % len(keys)

    last_error: Exception | None = None
    tried: list[int] = []

    for offset in range(len(keys)):
        key_idx = (start + offset) % len(keys)
        key = keys[key_idx]
        tried.append(key_idx)
        try:
            result = fn(key)
            _bump_usage(key)
            logger.debug("Gemini call succeeded with key %d/%d", key_idx + 1, len(keys))
            return result, key_idx, None
        except Exception as exc:
            last_error = exc
            if _is_retryable(exc):
                logger.warning(
                    "Gemini key %d/%d quota/rate-limit (%s), trying next",
                    key_idx + 1, len(keys), _mask_key(key)
                )
                continue
            # Non-retryable (bad key, invalid request, etc.) — fail fast
            logger.warning("Gemini key %d/%d non-retryable error: %s", key_idx + 1, len(keys), exc)
            return None, key_idx, exc

    logger.error("All %d Gemini keys exhausted. Last error: %s", len(keys), last_error)
    return None, None, last_error