from __future__ import annotations

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

logger = logging.getLogger(__name__)


@dataclass
class TokenBucket:
    rate: float
    burst: int

    _tokens: float = field(default=0, init=False)
    _last_refill: float = field(default_factory=time.monotonic, init=False)
    _lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False)

    def __post_init__(self):
        self._tokens = float(self.burst)

    def _refill(self) -> None:
        now = time.monotonic()
        elapsed = now - self._last_refill
        self._tokens = min(self._tokens + elapsed * self.rate, float(self.burst))
        self._last_refill = now

    async def acquire(self, tokens: float = 1.0) -> None:
        while True:
            async with self._lock:
                self._refill()
                if self._tokens >= tokens:
                    self._tokens -= tokens
                    return

            wait = (tokens - self._tokens) / self.rate if self.rate > 0 else 0.1
            await asyncio.sleep(max(wait, 0.01))

    def try_acquire(self, tokens: float = 1.0) -> bool:
        self._refill()
        if self._tokens >= tokens:
            self._tokens -= tokens
            return True
        return False


@dataclass
class RouteRateLimiter:
    defaults: dict[str, tuple[float, int]] = field(
        default_factory=lambda: {
            "send_message": (5.0, 10),
            "send_media": (1.0, 3),
            "upload_media": (0.5, 2),
            "default": (10.0, 20),
        }
    )

    _buckets: dict[str, TokenBucket] = field(default_factory=dict, init=False)
    _lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False)

    async def acquire(self, route: str = "default", tokens: float = 1.0) -> None:
        bucket = await self._get_or_create_bucket(route)
        await bucket.acquire(tokens)

    def try_acquire(self, route: str = "default", tokens: float = 1.0) -> bool:
        bucket = self._get_or_create_bucket_sync(route)
        return bucket.try_acquire(tokens)

    async def _get_or_create_bucket(self, route: str) -> TokenBucket:
        async with self._lock:
            if route not in self._buckets:
                rate, burst = self.defaults.get(route, self.defaults["default"])
                self._buckets[route] = TokenBucket(rate=rate, burst=burst)
            return self._buckets[route]

    def _get_or_create_bucket_sync(self, route: str) -> TokenBucket:
        if route not in self._buckets:
            rate, burst = self.defaults.get(route, self.defaults["default"])
            self._buckets[route] = TokenBucket(rate=rate, burst=burst)
        return self._buckets[route]

    def get_stats(self) -> dict[str, dict]:
        return {
            route: {"tokens": bucket._tokens, "rate": bucket.rate, "burst": bucket.burst}
            for route, bucket in self._buckets.items()
        }