"""Redis-backed token-bucket rate limiting.

This module provides a reusable FastAPI dependency that rate-limits requests
per endpoint scope (optionally per-user) using an atomic Lua script on Redis.

If Redis is unavailable the limiter can either allow the request through
(fail-open) or reject it with a ``503`` (fail-closed), depending on the
policy configured for that limit. Chat send uses fail-closed by default
because flooding/abuse is worse than a temporary outage.

Example
-------
    from app.rate_limit import rate_limit, named_limit

    chat_send_limit = rate_limit(
        named_limit("chat_send", capacity=5, refill_rate=0.2, fail_open=False),
        user_id_dependency=Depends(get_current_user_id),
    )

    @router.post("/messages")
    async def send_message(
        payload: MessageCreate,
        current_user: CurrentUser,
        db: DbDep,
        _: None = Depends(chat_send_limit),
    ):
        ...
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any, Awaitable, Callable

from fastapi import Depends, HTTPException, Request, status
from redis.asyncio import Redis
from redis.exceptions import RedisError

from app.config import settings

# Atomic token-bucket implementation.
#
# KEYS[1] is the bucket key. ARGV[1] is the bucket capacity, ARGV[2] is the
# refill rate in tokens per second, ARGV[3] is the current Unix time in
# seconds (float), and ARGV[4] is the request cost (tokens to take).
#
# The bucket state is stored as a hash with "tokens" and "ts". The script
# refills the bucket up to capacity, then takes the requested tokens. It
# returns an array:
#   [0] allowed (1 or 0)
#   [1] tokens remaining after the take
#   [2] retry-after seconds (0 when allowed)
_TOKEN_BUCKET_LUA = """
local key = KEYS[1]
local capacity = tonumber(ARGV[1])
local refill_rate = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local cost = tonumber(ARGV[4])

local state = redis.call("HMGET", key, "tokens", "ts")
local tokens = capacity
local ts = now

if state[1] ~= false then
    tokens = tonumber(state[1])
    ts = tonumber(state[2]) or now
end

local elapsed = math.max(0, now - ts)
tokens = math.min(capacity, tokens + (elapsed * refill_rate))

local allowed = 0
local retry_after = 0

if tokens >= cost then
    tokens = tokens - cost
    allowed = 1
else
    -- Time until enough tokens for a single request have accrued.
    retry_after = math.ceil((cost - tokens) / refill_rate)
end

-- Keep the bucket state small and expire-able. Retaining the timestamp as a
-- float lets partial tokens accumulate smoothly between requests.
redis.call("HSET", key, "tokens", tostring(tokens), "ts", tostring(now))
redis.call("PEXPIRE", key, 600000)

return {allowed, tostring(tokens), retry_after}
"""


class RateLimitExceeded(HTTPException):
    """Raised when a rate-limit bucket has no tokens left."""

    def __init__(self, retry_after: int) -> None:
        super().__init__(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail="Rate limit exceeded",
            headers={"Retry-After": str(retry_after)},
        )


@dataclass(frozen=True)
class RateLimit:
    """Definition of a named token-bucket rate limit.

    ``capacity`` is the burst size and ``refill_rate`` is the steady-state
    token accumulation rate (tokens/second). For example,
    ``RateLimit(name="chat_send", capacity=10, refill_rate=1.0)`` allows a
    burst of 10 requests and then one request per second on average.
    """

    name: str
    capacity: int
    refill_rate: float
    # When Redis is unreachable: ``True`` allows the request (fail-open),
    # ``False`` rejects it with 503 (fail-closed).
    fail_open: bool = True

    def __post_init__(self) -> None:
        if not self.name:
            raise ValueError("name must be non-empty")
        if self.capacity <= 0:
            raise ValueError("capacity must be positive")
        if self.refill_rate <= 0:
            raise ValueError("refill_rate must be positive")


def default_user_scope(request: Request, user_id: int | None) -> str:
    """Scope a bucket per authenticated user, or per client IP if anonymous."""
    if user_id is not None:
        return f"user:{user_id}"
    client = request.client
    return f"ip:{client.host if client else 'unknown'}"


@dataclass(frozen=True)
class RateLimitConfig:
    """A named rate limit plus how it is applied to a request."""

    limit: RateLimit
    # Builds the Redis key fragment for the request. The resolved user id is
    # passed as the second argument (or ``None`` for anonymous limits).
    scope: Callable[[Request, int | None], str] = default_user_scope


@dataclass(frozen=True)
class RateLimitResult:
    """Outcome of a rate-limit check (used by tests)."""

    allowed: bool
    remaining: int
    retry_after: int


_redis_client: Redis | None = None


def get_redis() -> Redis:
    """Return a shared async Redis client for rate limiting.

    The client is created lazily (``Redis.from_url`` does not connect until a
    command is issued), so this dependency never blocks or fails just because
    Redis is down.
    """
    global _redis_client
    if _redis_client is None:
        _redis_client = Redis.from_url(
            settings.REDIS_URL,
            encoding="utf-8",
            decode_responses=True,
            socket_connect_timeout=settings.REDIS_TIMEOUT_SECONDS,
            socket_timeout=settings.REDIS_TIMEOUT_SECONDS,
        )
    return _redis_client


class RateLimiter:
    """Enforce a token bucket in Redis.

    Buckets are keyed as ``rate:{limit.name}:{scope_key}`` so multiple
    endpoints can share a named limit while remaining independently bucketed.
    """

    def __init__(self, config: RateLimitConfig, redis: Redis | None = None) -> None:
        self._config = config
        self._redis = redis

    def _redis_client(self) -> Redis:
        return self._redis or get_redis()

    def _bucket_key(self, request: Request, user_id: int | None) -> str:
        limit = self._config.limit
        scope_key = self._config.scope(request, user_id)
        return f"rate:{limit.name}:{scope_key}"

    async def check(self, request: Request, user_id: int | None = None) -> RateLimitResult:
        """Atomically refill/take a token and return the outcome.

        Redis failures are mapped to ``allowed=True`` (fail-open) or
        ``allowed=False`` (fail-closed, represented with ``retry_after=-1``)
        according to the limit's policy, without raising.
        """
        limit = self._config.limit
        key = self._bucket_key(request, user_id)

        try:
            result = await self._redis_client().eval(
                _TOKEN_BUCKET_LUA,
                1,
                key,
                str(limit.capacity),
                str(limit.refill_rate),
                str(time.time()),
                "1",  # each request consumes one token
            )
        except RedisError:
            if limit.fail_open:
                return RateLimitResult(allowed=True, remaining=0, retry_after=0)
            # Fail-closed: no reliable retry estimate is available.
            return RateLimitResult(allowed=False, remaining=0, retry_after=-1)

        # Redis may return ints or strings depending on decode_responses.
        allowed = bool(int(result[0]))
        remaining = int(result[1])
        retry_after = int(result[2])
        return RateLimitResult(allowed=allowed, remaining=remaining, retry_after=retry_after)

    async def __call__(self, request: Request, user_id: int | None = None) -> None:
        """Enforce the rate limit, raising 429 or 503 as appropriate."""
        result = await self.check(request, user_id)
        if result.allowed:
            return

        if result.retry_after < 0:
            # Redis was unavailable and this limit is fail-closed.
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="Rate limiting service unavailable",
            )
        raise RateLimitExceeded(retry_after=result.retry_after)


def named_limit(
    name: str,
    *,
    capacity: int,
    refill_rate: float,
    fail_open: bool = True,
) -> RateLimitConfig:
    """Create a named rate-limit config with the default per-user/IP scope.

    ``capacity`` is the burst size and ``refill_rate`` is tokens/second.
    """
    return RateLimitConfig(
        limit=RateLimit(
            name=name,
            capacity=capacity,
            refill_rate=refill_rate,
            fail_open=fail_open,
        )
    )


def rate_limit(
    config: RateLimitConfig,
    *,
    user_id_dependency: Any = None,
) -> Callable[[Request, Any], Awaitable[None]]:
    """Build a FastAPI dependency enforcing ``config``.

    ``user_id_dependency`` optionally injects the authenticated user id, e.g.
    ``Depends(get_current_user_id)``. If omitted the limiter buckets by client
    IP (the ``default_user_scope`` fallback), which is suitable for anonymous
    endpoints but should be paired with an auth resolver for per-user limits.

    The returned callable is intended for ``Depends(...)``.
    """

    async def _anonymous_user_id() -> int | None:
        return None

    resolver = user_id_dependency if user_id_dependency is not None else Depends(_anonymous_user_id)

    async def dependency(
        request: Request,
        user_id: int | None = resolver,
    ) -> None:
        limiter = RateLimiter(config)
        await limiter(request, user_id)

    return dependency