"""Token-bucket rate limiter for sensitive endpoints.

Used by the privacy routes to prevent a misconfigured client (or a runaway
script) from spamming right-to-erasure calls. Lightweight in-process
implementation with no Redis dependency — this is a local personal-agent
tool, not a multi-tenant service.

Buckets are keyed by (scope, caller) where `caller` is the client IP for
HTTP requests. The bucket holds `capacity` tokens and refills at `refill`
tokens per second up to capacity. Every successful `acquire()` consumes one
token; a request arriving when the bucket is empty returns False.

SQLite persistence (v4): bucket state can optionally be persisted to the
main database so restarts don't reset counters. Without persistence, a
restart loop would let a caller burst past the hourly cap. The `attach_db`
method rehydrates buckets at startup and writes each `acquire()` through
to the `rate_limit_buckets` table.
"""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import aiosqlite

log = logging.getLogger(__name__)


def _now_iso() -> str:
    from datetime import UTC, datetime
    return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")


@dataclass
class _Bucket:
    capacity: float
    refill: float  # tokens per second
    tokens: float = 0.0
    last_refill: float = field(default_factory=time.monotonic)


class TokenBucketRateLimiter:
    """Async-safe token-bucket rate limiter.

    Example:
        limiter = TokenBucketRateLimiter()
        limiter.configure("privacy_forget", capacity=5, refill=1 / 60)  # 5 per minute
        if not await limiter.acquire("privacy_forget", caller=client_ip):
            raise HTTPException(429, "Rate limit exceeded")
    """

    _SCHEMA = """
    CREATE TABLE IF NOT EXISTS rate_limit_buckets (
        scope        TEXT NOT NULL,
        caller       TEXT NOT NULL,
        tokens       REAL NOT NULL,
        last_refill  REAL NOT NULL,
        updated_at   TEXT NOT NULL,
        PRIMARY KEY (scope, caller)
    );
    """

    def __init__(self) -> None:
        self._lock = asyncio.Lock()
        self._buckets: dict[tuple[str, str], _Bucket] = {}
        self._policies: dict[str, tuple[float, float]] = {}  # scope → (capacity, refill)
        self._db: aiosqlite.Connection | None = None
        self._persist_initialized = False

    def configure(self, scope: str, *, capacity: float, refill: float) -> None:
        """Set (or update) the policy for a scope.

        capacity — bucket size (max burst).
        refill   — tokens added per second.
        """
        self._policies[scope] = (capacity, refill)

    async def attach_db(self, db: aiosqlite.Connection) -> None:
        """Wire a SQLite connection for persistence.

        Creates the table (idempotent), rehydrates any previously-persisted
        buckets into memory, and marks the limiter as persistent. Safe to
        call once at startup. A limiter without a db is still fully
        functional — it just resets on restart.
        """
        self._db = db
        try:
            await self._db.executescript(self._SCHEMA)
            await self._db.commit()
            cur = await self._db.execute(
                "SELECT scope, caller, tokens, last_refill FROM rate_limit_buckets"
            )
            rows = await cur.fetchall()
        except Exception:
            log.exception("Failed to rehydrate rate_limit_buckets table")
            return

        # Restore buckets into memory. `last_refill` was stored as wall clock
        # seconds-since-epoch (time.time), but bucket internals use monotonic
        # time, so we compute an offset that compensates for the restart.
        now_wall = time.time()
        now_mono = time.monotonic()
        offset = now_mono - now_wall

        for scope, caller, tokens, wall_last_refill in rows:
            policy = self._policies.get(scope)
            if policy is None:
                # Drop buckets whose scope is no longer configured
                continue
            capacity, refill = policy
            mono_last_refill = float(wall_last_refill) + offset
            self._buckets[(scope, caller)] = _Bucket(
                capacity=capacity,
                refill=refill,
                tokens=float(tokens),
                last_refill=mono_last_refill,
            )
        self._persist_initialized = True
        log.info("Rate limiter rehydrated %d buckets from SQLite", len(rows))

    async def _persist(self, scope: str, caller: str, bucket: _Bucket) -> None:
        """Write-through a single bucket to SQLite. Non-fatal on failure."""
        if self._db is None:
            return
        # Convert bucket's monotonic last_refill back to wall-clock seconds
        # for persistence, so restarts can compute an offset against the new
        # monotonic clock.
        now_wall = time.time()
        now_mono = time.monotonic()
        wall_last_refill = now_wall - (now_mono - bucket.last_refill)
        try:
            await self._db.execute(
                "INSERT INTO rate_limit_buckets "
                "(scope, caller, tokens, last_refill, updated_at) "
                "VALUES (?, ?, ?, ?, ?) "
                "ON CONFLICT(scope, caller) DO UPDATE SET "
                "tokens=excluded.tokens, last_refill=excluded.last_refill, "
                "updated_at=excluded.updated_at",
                (scope, caller, bucket.tokens, wall_last_refill,
                 _now_iso()),
            )
            await self._db.commit()
        except Exception:
            log.debug("Rate limiter persist failed", exc_info=True)

    async def acquire(self, scope: str, caller: str = "default") -> bool:
        """Consume one token for (scope, caller). Returns False if empty."""
        policy = self._policies.get(scope)
        if policy is None:
            return True  # unlimited when not configured
        capacity, refill = policy

        async with self._lock:
            key = (scope, caller)
            now = time.monotonic()
            bucket = self._buckets.get(key)
            if bucket is None:
                bucket = _Bucket(capacity=capacity, refill=refill, tokens=capacity, last_refill=now)
                self._buckets[key] = bucket
            else:
                elapsed = now - bucket.last_refill
                bucket.tokens = min(capacity, bucket.tokens + elapsed * refill)
                bucket.last_refill = now

            if bucket.tokens < 1:
                # Persist the REFRESHED state even on rejection so a tight
                # retry loop doesn't keep resetting the clock to "just now"
                # through bucket creation.
                await self._persist(scope, caller, bucket)
                return False
            bucket.tokens -= 1
            await self._persist(scope, caller, bucket)
            return True

    def reset(self, scope: str | None = None, caller: str | None = None) -> None:
        """Clear tracked buckets — used by tests."""
        if scope is None and caller is None:
            self._buckets.clear()
            return
        to_remove = [
            k for k in self._buckets
            if (scope is None or k[0] == scope) and (caller is None or k[1] == caller)
        ]
        for k in to_remove:
            self._buckets.pop(k, None)

    def snapshot(self, scope: str) -> dict:
        """Debug / health info for a scope."""
        return {
            "scope": scope,
            "policy": self._policies.get(scope),
            "active_callers": sum(1 for k in self._buckets if k[0] == scope),
        }


# ── Process-wide singleton (wired from main.py) ──
_limiter: TokenBucketRateLimiter | None = None


def get_limiter() -> TokenBucketRateLimiter:
    """Return the process-wide limiter, creating it on first use."""
    global _limiter
    if _limiter is None:
        _limiter = TokenBucketRateLimiter()
        # Default policies for sensitive endpoints
        _limiter.configure("privacy_forget", capacity=5, refill=5 / 3600)   # 5/hour
        _limiter.configure("privacy_export", capacity=3, refill=3 / 3600)   # 3/hour
        _limiter.configure("decisions_prune", capacity=10, refill=10 / 3600)  # 10/hour

        # MCP SSE — per-token (caller is the token_id, so a noisy token
        # can't starve others). Three scopes:
        #   mcp_session_open      — opening a new SSE stream. Cheap to
        #                           cap (a session-open storm doesn't
        #                           come from a normal client).
        #   mcp_message_full      — per-message limit for full-scope
        #                           tokens. Tighter because these can
        #                           write (send email, run pipelines).
        #   mcp_message_read_only — per-message limit for read-only
        #                           tokens. More generous since the
        #                           blast radius is just CPU + token
        #                           cost, not data mutation.
        _limiter.configure("mcp_session_open",      capacity=10,  refill=10 / 60)  # 10/min burst, 10/min sustained
        _limiter.configure("mcp_message_full",      capacity=60,  refill=1.0)      # 60 burst, 60/sec sustained
        _limiter.configure("mcp_message_read_only", capacity=300, refill=5.0)      # 300 burst, 5/sec sustained
    return _limiter
