from typing import Callable, List, Optional

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import ASGIApp

from rate_limiter.config import RateLimiterConfig
from rate_limiter.core.identity import IdentityResolver, get_identity_resolver
from rate_limiter.core.limiter import RateLimiter
from rate_limiter.core.rules import RateLimitRule
from rate_limiter.exceptions import RateLimitExceeded


class RateLimiterMiddleware(BaseHTTPMiddleware):
    """Starlette middleware for rate limiting."""

    def __init__(
        self,
        app: ASGIApp,
        rules: Optional[List[RateLimitRule]] = None,
        config: Optional[RateLimiterConfig] = None,
        identity_resolver: Optional[IdentityResolver] = None,
    ):
        """
        Initialize rate limiter middleware.

        Args:
            app: ASGI application
            rules: List of rate limit rules
            config: Rate limiter configuration
            identity_resolver: Optional custom identity resolver
        """
        super().__init__(app)
        self.config = config or RateLimiterConfig.from_env()
        self.limiter = RateLimiter(config=self.config, rules=rules)
        self.identity_resolver = identity_resolver or get_identity_resolver(
            "ip", trust_proxy=self.config.trust_proxy_headers
        )

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        """
        Process request through rate limiter.

        Args:
            request: Starlette request
            call_next: Next middleware/handler

        Returns:
            Response
        """
        # Resolve identity
        identity = await self.identity_resolver.resolve(request)

        # Get endpoint path for scoped rules
        endpoint = request.url.path

        # Check rate limits and get info in one call
        rate_limit_info, rate_limit_exception = await self.limiter.check_with_info(
            identity=identity, endpoint=endpoint
        )

        if rate_limit_exception:
            # Return 429 Too Many Requests with rate limit headers
            e = rate_limit_exception
            import time

            reset_timestamp = int(time.time()) + e.retry_after
            headers = {
                "Retry-After": str(e.retry_after),
                "X-RateLimit-Limit": str(e.limit),
                "X-RateLimit-Remaining": "0",
                "X-RateLimit-Reset": str(reset_timestamp),
            }
            return Response(
                content=f"Rate limit exceeded. Retry after {e.retry_after} seconds.",
                status_code=429,
                headers=headers,
            )

        # Process request
        response = await call_next(request)

        # Add rate limit headers to successful response
        if rate_limit_info:
            headers = rate_limit_info.to_headers()
            for key, value in headers.items():
                response.headers[key] = value

        return response
