"""Rate limiting middleware for protecting FastMCP servers from abuse."""

import inspect
import time
from collections import defaultdict, deque
from collections.abc import Awaitable, Callable
from typing import Any, cast

import anyio
from mcp import McpError
from mcp.types import ErrorData

from .middleware import CallNext, Middleware, MiddlewareContext


class RateLimitError(McpError):
    """Error raised when rate limit is exceeded."""

    def __init__(self, message: str = "Rate limit exceeded"):
        super().__init__(ErrorData(code=-32000, message=message))


class TokenBucketRateLimiter:
    """Token bucket implementation for rate limiting."""

    def __init__(self, capacity: int, refill_rate: float):
        """Initialize token bucket.

        Args:
            capacity: Maximum number of tokens in the bucket
            refill_rate: Tokens added per second
        """
        self.capacity = capacity
        self.refill_rate = refill_rate
        self.tokens = capacity
        self.last_refill = time.time()
        self._lock = anyio.Lock()

    async def consume(self, tokens: int = 1) -> bool:
        """Try to consume tokens from the bucket.

        Args:
            tokens: Number of tokens to consume

        Returns:
            True if tokens were available and consumed, False otherwise
        """
        async with self._lock:
            now = time.time()
            elapsed = now - self.last_refill

            # Add tokens based on elapsed time
            self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
            self.last_refill = now

            if self.tokens >= tokens:
                self.tokens -= tokens
                return True
            return False


class SlidingWindowRateLimiter:
    """Sliding window rate limiter implementation."""

    def __init__(self, max_requests: int, window_seconds: int):
        """Initialize sliding window rate limiter.

        Args:
            max_requests: Maximum requests allowed in the time window
            window_seconds: Time window in seconds
        """
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self.requests = deque()
        self._lock = anyio.Lock()

    async def is_allowed(self) -> bool:
        """Check if a request is allowed."""
        async with self._lock:
            now = time.time()
            cutoff = now - self.window_seconds

            # Remove old requests outside the window
            while self.requests and self.requests[0] < cutoff:
                self.requests.popleft()

            if len(self.requests) < self.max_requests:
                self.requests.append(now)
                return True
            return False


class RateLimitingMiddleware(Middleware):
    """Middleware that implements rate limiting to prevent server abuse.

    Uses a token bucket algorithm by default, allowing for burst traffic
    while maintaining a sustainable long-term rate.

    Example:
        ```python
        from fastmcp.server.middleware.rate_limiting import RateLimitingMiddleware

        # Allow 10 requests per second with bursts up to 20
        rate_limiter = RateLimitingMiddleware(
            max_requests_per_second=10,
            burst_capacity=20
        )

        mcp = FastMCP("MyServer")
        mcp.add_middleware(rate_limiter)
        ```
    """

    def __init__(
        self,
        max_requests_per_second: float = 10.0,
        burst_capacity: int | None = None,
        get_client_id: Callable[[MiddlewareContext], str]
        | Callable[[MiddlewareContext], Awaitable[str]]
        | None = None,
        global_limit: bool = False,
    ):
        """Initialize rate limiting middleware.

        Args:
            max_requests_per_second: Sustained requests per second allowed
            burst_capacity: Maximum burst capacity. If None, defaults to 2x max_requests_per_second
            get_client_id: Function to extract client ID from context. Can be sync or async.
                If None, uses global limiting
            global_limit: If True, apply limit globally; if False, per-client
        """
        self.max_requests_per_second = max_requests_per_second
        self.burst_capacity = burst_capacity or int(max_requests_per_second * 2)
        self.get_client_id = get_client_id
        self.global_limit = global_limit

        # Storage for rate limiters per client
        self.limiters: dict[str, TokenBucketRateLimiter] = defaultdict(
            lambda: TokenBucketRateLimiter(
                self.burst_capacity, self.max_requests_per_second
            )
        )

        # Global rate limiter
        if self.global_limit:
            self.global_limiter = TokenBucketRateLimiter(
                self.burst_capacity, self.max_requests_per_second
            )

    async def _get_client_identifier(self, context: MiddlewareContext) -> str:
        """Get client identifier for rate limiting."""
        if self.get_client_id:
            client_id = self.get_client_id(context)
            if inspect.isawaitable(client_id):
                return cast(str, await client_id)
            return client_id
        return "global"

    async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
        """Apply rate limiting to requests."""
        if self.global_limit:
            # Global rate limiting
            allowed = await self.global_limiter.consume()
            if not allowed:
                raise RateLimitError("Global rate limit exceeded")
        else:
            # Per-client rate limiting
            client_id = await self._get_client_identifier(context)
            limiter = self.limiters[client_id]
            allowed = await limiter.consume()
            if not allowed:
                raise RateLimitError(f"Rate limit exceeded for client: {client_id}")

        return await call_next(context)


class SlidingWindowRateLimitingMiddleware(Middleware):
    """Middleware that implements sliding window rate limiting.

    Uses a sliding window approach which provides more precise rate limiting
    but uses more memory to track individual request timestamps.

    Example:
        ```python
        from fastmcp.server.middleware.rate_limiting import SlidingWindowRateLimitingMiddleware

        # Allow 100 requests per minute
        rate_limiter = SlidingWindowRateLimitingMiddleware(
            max_requests=100,
            window_minutes=1
        )

        mcp = FastMCP("MyServer")
        mcp.add_middleware(rate_limiter)
        ```
    """

    def __init__(
        self,
        max_requests: int,
        window_minutes: int = 1,
        get_client_id: Callable[[MiddlewareContext], str]
        | Callable[[MiddlewareContext], Awaitable[str]]
        | None = None,
    ):
        """Initialize sliding window rate limiting middleware.

        Args:
            max_requests: Maximum requests allowed in the time window
            window_minutes: Time window in minutes
            get_client_id: Function to extract client ID from context. Can be sync or async.
                If None, uses global limiting
        """
        self.max_requests = max_requests
        self.window_seconds = window_minutes * 60
        self.get_client_id = get_client_id

        # Storage for rate limiters per client
        self.limiters: dict[str, SlidingWindowRateLimiter] = defaultdict(
            lambda: SlidingWindowRateLimiter(self.max_requests, self.window_seconds)
        )

    async def _get_client_identifier(self, context: MiddlewareContext) -> str:
        """Get client identifier for rate limiting."""
        if self.get_client_id:
            client_id = self.get_client_id(context)
            if inspect.isawaitable(client_id):
                return cast(str, await client_id)
            return client_id
        return "global"

    async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
        """Apply sliding window rate limiting to requests."""
        client_id = await self._get_client_identifier(context)
        limiter = self.limiters[client_id]

        allowed = await limiter.is_allowed()
        if not allowed:
            raise RateLimitError(
                f"Rate limit exceeded: {self.max_requests} requests per "
                f"{self.window_seconds // 60} minutes for client: {client_id}"
            )

        return await call_next(context)
