"""Lightweight rate limiting for the Space's live-solve path.

The Space serves precomputed QAOA results instantly, but it also runs the
pure-Python classical solvers (brute-force, greedy) live so visitors can confirm
the served numbers are real. Those solves are cheap (milliseconds), but a public
endpoint still warrants a guard against rapid-fire hammering. Two layers:

1. **Per-IP rate limit** - at most one live solve per ``window_seconds`` per
   visitor IP (read from ``x-forwarded-for``, set by HF's reverse proxy).
2. **Global daily cap** - a generous ceiling on live solves per UTC day across
   all visitors, persisted to JSON when a writable dir is available.

There is no IBM-quota guard here (this Space has no hardware path). Verdicts are
returned as a dataclass so the Gradio layer stays thin and the logic is trivially
unit-testable.
"""

from __future__ import annotations

import json
import threading
from dataclasses import dataclass
from datetime import UTC, date, datetime
from pathlib import Path
from typing import Literal

SafetyReason = Literal["rate_limited", "daily_cap"]


@dataclass(frozen=True)
class SafetyVerdict:
    """Outcome of a :meth:`RateLimiter.check_and_register` call."""

    allowed: bool
    reason: SafetyReason | None
    detail: str
    daily_remaining: int
    daily_cap: int


class RateLimiter:
    """Per-IP + daily gate for the live-solve path. Thread-safe via one lock."""

    def __init__(
        self,
        *,
        window_seconds: int = 3,
        daily_cap: int = 5000,
        persist_path: Path | None = None,
    ) -> None:
        if window_seconds < 0:
            raise ValueError("window_seconds must be >= 0")
        if daily_cap < 0:
            raise ValueError("daily_cap must be >= 0")
        self._window = window_seconds
        self._cap = daily_cap
        self._persist_path = persist_path
        self._lock = threading.Lock()
        self._last_ip: dict[str, float] = {}
        self._day: date | None = None
        self._count: int = 0
        self._load_persisted()

    def check_and_register(self, *, ip: str, now: datetime) -> SafetyVerdict:
        """Decide whether ``ip`` may run a live solve at ``now`` (commit-on-allow)."""
        with self._lock:
            self._roll_day_if_needed(now)
            self._evict_stale_ips(now.timestamp())

            if self._count >= self._cap:
                return SafetyVerdict(
                    allowed=False,
                    reason="daily_cap",
                    detail=f"Daily limit of {self._cap} live solves reached. Resets midnight UTC.",
                    daily_remaining=0,
                    daily_cap=self._cap,
                )

            last = self._last_ip.get(ip)
            if last is not None and (now.timestamp() - last) < self._window:
                wait = self._window - int(now.timestamp() - last)
                return SafetyVerdict(
                    allowed=False,
                    reason="rate_limited",
                    detail=f"Please wait {max(wait, 1)}s between live solves.",
                    daily_remaining=max(0, self._cap - self._count),
                    daily_cap=self._cap,
                )

            self._last_ip[ip] = now.timestamp()
            self._count += 1
            self._persist()
            return SafetyVerdict(
                allowed=True,
                reason=None,
                detail="ok",
                daily_remaining=max(0, self._cap - self._count),
                daily_cap=self._cap,
            )

    def daily_remaining(self, now: datetime) -> int:
        with self._lock:
            self._roll_day_if_needed(now)
            return max(0, self._cap - self._count)

    def daily_cap(self) -> int:
        return self._cap

    def _roll_day_if_needed(self, now: datetime) -> None:
        today = now.astimezone(UTC).date()
        if self._day != today:
            self._day = today
            self._count = 0

    def _evict_stale_ips(self, now_ts: float) -> None:
        if self._window <= 0:
            self._last_ip.clear()
            return
        cutoff = now_ts - self._window
        self._last_ip = {ip: ts for ip, ts in self._last_ip.items() if ts > cutoff}

    def _load_persisted(self) -> None:
        if self._persist_path is None:
            return
        try:
            raw = json.loads(self._persist_path.read_text(encoding="utf-8"))
        except (FileNotFoundError, json.JSONDecodeError):
            return
        if not isinstance(raw, dict):
            return
        try:
            self._day = date.fromisoformat(str(raw.get("date")))
            self._count = int(raw.get("count", 0))
        except (TypeError, ValueError):
            self._day = None
            self._count = 0

    def _persist(self) -> None:
        if self._persist_path is None or self._day is None:
            return
        try:
            self._persist_path.parent.mkdir(parents=True, exist_ok=True)
            self._persist_path.write_text(
                json.dumps({"date": self._day.isoformat(), "count": self._count}),
                encoding="utf-8",
            )
        except OSError:
            self._persist_path = None


def default_persist_path() -> Path | None:
    """Return ``/data/qagent_quota.json`` when HF Persistent Storage is mounted."""
    candidate = Path("/data")
    if candidate.is_dir():
        try:
            test = candidate / ".qa_write_test"
            test.write_text("x")
            test.unlink()
        except OSError:
            return None
        return candidate / "qagent_quota.json"
    return None