# ======================================================
# key_pool.py — Thread-safe API key pool with rotation
# ======================================================
#
# Reads comma-separated keys from an env var (e.g. POOL_OPENAI_API_KEYS).
# Falls back to a single-key env var (e.g. OPENAI_API_KEY) when the pool
# is not configured.
#
# On rate-limit (429), the caller marks the key; the pool puts it on
# cooldown and rotates to the next available key automatically.
#

import os
import re
import time
import threading
from datetime import datetime, timezone


def _parse_duration(s: str) -> float | None:
    """Parse OpenAI-style duration strings like '6s', '1m30s', '2m' into seconds."""
    m = re.fullmatch(r"(?:(\d+)m)?(?:(\d+(?:\.\d+)?)s)?", s.strip())
    if not m or not any(m.groups()):
        return None
    minutes = int(m.group(1)) if m.group(1) else 0
    seconds = float(m.group(2)) if m.group(2) else 0.0
    return minutes * 60 + seconds


class KeyPool:
    """Round-robin key pool with per-key rate-limit cooldown."""

    def __init__(
        self,
        pool_env_var: str,
        fallback_env_var: str,
        cooldown_seconds: float = 60.0,
    ):
        self._lock = threading.Lock()
        self._cooldown_seconds = cooldown_seconds

        # 1. Try the pool env var (comma-separated)
        pool_raw = os.environ.get(pool_env_var, "")
        keys = [k.strip() for k in pool_raw.split(",") if k.strip()] if pool_raw else []

        # 2. Fallback to single-key env var
        if not keys:
            single = os.environ.get(fallback_env_var, "")
            if single:
                keys = [single]

        self._keys = keys
        self._index = 0
        self._cooldowns: dict[str, float] = {}  # key -> monotonic expiry

        # Startup log (mask keys)
        if keys:
            masked = [f"...{k[-4:]}" for k in keys]
            source = pool_env_var if len(keys) > 1 or pool_raw else fallback_env_var
            print(f"KeyPool({pool_env_var}): {len(keys)} key(s) from {source} [{', '.join(masked)}]")
        else:
            print(f"KeyPool({pool_env_var}): no keys found")

    # ------------------------------------------------------------------
    @property
    def size(self) -> int:
        return len(self._keys)

    @property
    def available(self) -> bool:
        return len(self._keys) > 0

    # ------------------------------------------------------------------
    def get_key(self) -> str | None:
        """Return next non-cooldown key (round-robin).

        If every key is on cooldown, sleeps until the soonest one expires
        and returns it.  Returns None only when the pool is empty.
        """
        if not self._keys:
            return None

        with self._lock:
            now = time.monotonic()
            n = len(self._keys)

            # Try each key starting from the current index
            for _ in range(n):
                idx = self._index % n
                self._index = (self._index + 1) % n
                key = self._keys[idx]
                if now >= self._cooldowns.get(key, 0):
                    return key

            # All on cooldown — find the one that expires soonest
            soonest_key = min(self._keys, key=lambda k: self._cooldowns.get(k, 0))
            wait = self._cooldowns[soonest_key] - now

        # Sleep *outside* the lock so other threads aren't blocked
        if wait > 0:
            print(f"KeyPool: all keys on cooldown, waiting {wait:.1f}s for ...{soonest_key[-4:]}")
            time.sleep(wait)

        return soonest_key

    # ------------------------------------------------------------------
    def mark_rate_limited(self, key: str, headers: dict | None = None) -> None:
        """Put *key* on cooldown after a 429 response.

        If the server sent rate-limit headers, use them to determine the
        actual cooldown and log when the key will be available again.
        """
        headers = headers or {}
        cooldown = self._parse_cooldown(headers)
        available_at = datetime.now(timezone.utc).timestamp() + cooldown
        available_str = datetime.fromtimestamp(available_at, tz=timezone.utc).strftime("%H:%M:%S UTC")

        with self._lock:
            self._cooldowns[key] = time.monotonic() + cooldown

        header_info = f" (headers: {headers})" if headers else ""
        print(
            f"KeyPool: rate-limited ...{key[-4:]}, "
            f"cooldown {cooldown:.0f}s, available ~{available_str}{header_info}"
        )

    # ------------------------------------------------------------------
    @staticmethod
    def _parse_cooldown(headers: dict) -> float:
        """Extract cooldown seconds from rate-limit headers.

        Checks (in priority order):
          1. retry-after  (seconds or HTTP-date)
          2. x-ratelimit-reset-requests  (e.g. "6s", "1m30s", "2m")
          3. x-ratelimit-reset-tokens    (same format)
        Falls back to the instance default when no header is usable.
        """
        # retry-after: either integer seconds or an HTTP-date
        if "retry-after" in headers:
            val = headers["retry-after"].strip()
            try:
                return float(val)
            except ValueError:
                pass

        # OpenAI duration strings like "6s", "1m30s", "2m"
        for hdr in ("x-ratelimit-reset-requests", "x-ratelimit-reset-tokens"):
            if hdr in headers:
                seconds = _parse_duration(headers[hdr])
                if seconds is not None:
                    return seconds

        return 60.0  # fallback

    def clear_cooldown(self, key: str) -> None:
        """Remove cooldown for *key* (call after a successful request)."""
        with self._lock:
            self._cooldowns.pop(key, None)