# Copyright (c) 2026 MergeGuide, Inc. All rights reserved.
# SPDX-License-Identifier: Proprietary
"""Rate limiting decorator for Lambda handlers.

Provides a clean decorator pattern for adding rate limiting to Lambda
handlers. The decorator handles rate limit checking, response formatting,
and header injection automatically.

Example Usage:
    from mergeguide.rate_limiting import rate_limited

    @rate_limited(identifier_extractor=lambda e, c: e.get("org_id"))
    def handler(event, context):
        return {"statusCode": 200, "body": "OK"}

The decorator:
    1. Extracts the identifier from the event
    2. Checks rate limits against DynamoDB
    3. Returns 429 if limit exceeded
    4. Otherwise calls the handler and adds rate limit headers
"""

from __future__ import annotations

import functools
import json
import logging
from collections.abc import Callable
from typing import Any

from .config import DEFAULT_TIER
from .limiter import RateLimitResult, check_rate_limit

logger = logging.getLogger(__name__)


# Type alias for Lambda handlers
LambdaHandler = Callable[[dict[str, Any], Any], dict[str, Any]]
IdentifierExtractor = Callable[[dict[str, Any], Any], str | None]
TierExtractor = Callable[[dict[str, Any], Any], str]


def _rate_limit_exceeded_response(result: RateLimitResult) -> dict[str, Any]:
    """Build a 429 Too Many Requests response.

    Args:
        result: The rate limit result with retry information.

    Returns:
        API Gateway compatible response with 429 status.
    """
    return {
        "statusCode": 429,
        "body": json.dumps(
            {
                "error": "rate_limit_exceeded",
                "message": f"Rate limit exceeded. Retry after {result.retry_after} seconds.",
                "retry_after": result.retry_after,
                "limit": result.limit,
                "reset_at": result.reset_at,
            }
        ),
        "headers": {
            "Content-Type": "application/json",
            **result.to_headers(),
        },
    }


def _add_rate_limit_headers(response: dict[str, Any], result: RateLimitResult) -> dict[str, Any]:
    """Add rate limit headers to an existing response.

    Args:
        response: The handler's response dictionary.
        result: The rate limit result with header values.

    Returns:
        The response with rate limit headers added.
    """
    headers = response.get("headers", {})
    headers.update(result.to_headers())
    response["headers"] = headers
    return response


def rate_limited(
    identifier_extractor: IdentifierExtractor,
    tier_extractor: TierExtractor | None = None,
    default_tier: str = DEFAULT_TIER,
) -> Callable[[LambdaHandler], LambdaHandler]:
    """Decorator to add rate limiting to a Lambda handler.

    Args:
        identifier_extractor: Function to extract the rate limit identifier
            from the event. Should return a string identifier or None to skip
            rate limiting.
        tier_extractor: Optional function to extract the billing tier from
            the event. If not provided, uses default_tier.
        default_tier: Fallback tier when tier_extractor returns None.

    Returns:
        Decorator function that wraps the handler with rate limiting.

    Example:
        @rate_limited(
            identifier_extractor=lambda e, c: extract_org_id(e),
            tier_extractor=lambda e, c: get_org_tier(e),
        )
        def my_handler(event, context):
            return {"statusCode": 200, "body": "OK"}
    """

    def decorator(handler: LambdaHandler) -> LambdaHandler:
        @functools.wraps(handler)
        def wrapper(event: dict[str, Any], context: Any) -> dict[str, Any]:
            # Extract the rate limit identifier
            identifier = identifier_extractor(event, context)

            if identifier is None:
                # No identifier - skip rate limiting (e.g., unauthenticated)
                logger.debug("No rate limit identifier, skipping rate limit check")
                return handler(event, context)

            # Extract or use default tier
            tier = default_tier
            if tier_extractor is not None:
                extracted_tier = tier_extractor(event, context)
                if extracted_tier:
                    tier = extracted_tier

            # Check rate limit
            result = check_rate_limit(identifier, tier)

            if not result.allowed:
                logger.warning(
                    f"Rate limit exceeded for {identifier}: {result.current_count}/{result.limit}"
                )
                return _rate_limit_exceeded_response(result)

            # Call the actual handler
            response = handler(event, context)

            # Add rate limit headers to successful responses
            return _add_rate_limit_headers(response, result)

        return wrapper

    return decorator
