from __future__ import annotations

import time
from dataclasses import dataclass


@dataclass
class TokenBucket:
    capacity: int
    refill_rate: float  # tokens per second
    tokens: float
    last_refill: float

    def allow(self, cost: int = 1) -> bool:
        now = time.time()
        elapsed = now - self.last_refill
        self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
        self.last_refill = now
        if self.tokens >= cost:
            self.tokens -= cost
            return True
        return False


class RateLimiter:
    """Simple per-org token bucket rate limiter.

    Defaults to unlimited (capacity=None). Production deployments should
    explicitly configure limits via config. Rate limiting is opt-in.
    """

    def __init__(self, capacity: int | None = None, refill_rate: float | None = None):
        self.capacity = capacity
        self.refill_rate = refill_rate
        self._buckets: dict[str, TokenBucket] = {}

    def check(self, org_id: str, cost: int = 1) -> bool:
        bucket = self._buckets.get(org_id)
        if not bucket:
            if self.capacity is None or self.refill_rate is None:
                # Unlimited until a specific per-org limit is configured.
                return True
            bucket = TokenBucket(
                self.capacity,
                self.refill_rate,
                float(self.capacity),
                time.time(),
            )
            self._buckets[org_id] = bucket
        return bucket.allow(cost)

    def set_limit(self, org_id: str, rpm: int, burst: int | None = None) -> None:
        """Set per-org request limits using RPM semantics.

        Args:
            org_id: Organization identifier.
            rpm: Requests per minute steady-state.
            burst: Optional burst capacity. Defaults to rpm if not provided.
        """
        rpm = max(0, int(rpm))
        if rpm <= 0:
            # Remove explicit limits so the org falls back to the default bucket.
            self._buckets.pop(org_id, None)
            return

        capacity = max(1, int(burst if burst is not None else rpm))
        refill = float(rpm) / 60.0
        now = time.time()
        bucket = self._buckets.get(org_id)
        if bucket is None:
            self._buckets[org_id] = TokenBucket(
                capacity=capacity, refill_rate=refill, tokens=capacity, last_refill=now
            )
            return
        # Update existing bucket atomically: clamp tokens to new capacity
        bucket.capacity = capacity
        bucket.refill_rate = refill
        bucket.tokens = min(bucket.tokens, float(capacity))

    def snapshot(self, org_id: str) -> dict[str, float | int | None]:
        """Return current token bucket state for diagnostics."""

        bucket = self._buckets.get(org_id)
        if bucket is None:
            return {
                "capacity": self.capacity,
                "tokens": None if self.capacity is None else float(self.capacity),
                "refill_rate": self.refill_rate,
            }
        return {
            "capacity": bucket.capacity,
            "tokens": round(bucket.tokens, 4),
            "refill_rate": bucket.refill_rate,
        }
