import functools
import inspect
from typing import Callable, Optional, TypeVar

from fastapi import Request

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

F = TypeVar("F", bound=Callable)


def rate_limit(
    rule_str: str,
    key: str = "ip",
    config: Optional[RateLimiterConfig] = None,
    identity_resolver: Optional[IdentityResolver] = None,
) -> Callable[[F], F]:
    """
    Decorator for rate limiting FastAPI endpoints.

    Args:
        rule_str: Rate limit rule string (e.g., "5/minute")
        key: Identity key type ("ip", "user_id", "api_key")
        config: Optional rate limiter configuration
        identity_resolver: Optional custom identity resolver

    Returns:
        Decorated function

    Example:
        @app.get("/login")
        @rate_limit("5/minute", key="ip")
        async def login(request: Request):
            return {"message": "Login"}
    """
    rule = RateLimitRule.from_string(rule_str, key=key)
    limiter_config = config or RateLimiterConfig.from_env()
    limiter = RateLimiter(config=limiter_config, rules=[rule])
    resolver = identity_resolver or get_identity_resolver(
        key, trust_proxy=limiter_config.trust_proxy_headers
    )

    def decorator(func: F) -> F:
        # Check if function already has Request parameter
        sig = inspect.signature(func)
        has_request_param = any(
            param.annotation == Request or param.name == "request"
            for param in sig.parameters.values()
        )

        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            # Find Request object in args/kwargs or get from FastAPI context
            request = None
            
            # Try to find in args/kwargs first
            for arg in args:
                if isinstance(arg, Request):
                    request = arg
                    break
            if request is None:
                request = kwargs.get("request")
            
            # If not found and function doesn't have request param, inject it
            if request is None and not has_request_param:
                # Try to get from FastAPI's dependency context
                try:
                    from starlette.requests import Request as StarletteRequest
                    # Check if we can get it from the current context
                    # This is a fallback - ideally the endpoint should have Request param
                    for arg in args:
                        if isinstance(arg, StarletteRequest):
                            request = arg
                            break
                except Exception:
                    pass
            
            # If still not found, try to inject Request dependency
            if request is None:
                # FastAPI should inject Request automatically, but if it's not in signature,
                # we need to get it from the context. For now, raise a helpful error.
                raise ValueError(
                    "Request object not found. Add 'request: Request' parameter to your endpoint function."
                )

            # Resolve identity
            identity = await resolver.resolve(request)

            # Get endpoint path
            endpoint = request.url.path if hasattr(request, "url") else ""

            # Check rate limits and get info in one call
            rate_limit_info, rate_limit_exception = await 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
                from fastapi import HTTPException

                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),
                }
                raise HTTPException(
                    status_code=429,
                    detail=f"Rate limit exceeded. Retry after {e.retry_after} seconds.",
                    headers=headers,
                )

            # Call the original function
            result = await func(*args, **kwargs)

            # Add rate limit headers to response
            if rate_limit_info:
                from fastapi import Response
                from starlette.responses import JSONResponse

                headers = rate_limit_info.to_headers()

                # If result is already a Response object, add headers to it
                if isinstance(result, Response):
                    for key, value in headers.items():
                        result.headers[key] = value
                    return result
                else:
                    # Wrap non-Response return values in JSONResponse with headers
                    # FastAPI will use this Response object instead of auto-wrapping
                    return JSONResponse(content=result, headers=headers)

            return result

        return wrapper  # type: ignore

    return decorator
