# Copyright (c) 2026 MergeGuide, Inc. All rights reserved.
# SPDX-License-Identifier: Proprietary
"""Core rate limiting implementation with DynamoDB atomic counters.

This module implements a sliding window rate limiter using DynamoDB as
the backing store. It uses atomic counters via UpdateItem ADD operations
to ensure thread-safe request counting across multiple Lambda invocations.

Key Features:
    - Atomic counter increments (no race conditions)
    - 1-minute sliding windows with automatic TTL cleanup
    - Fail-open behavior (allows requests if DynamoDB fails)
    - Standard rate limit headers for HTTP responses

Window Key Format:
    WINDOW#{YYYY-MM-DD-HH-mm}

    Example: WINDOW#2026-01-23-14-30 for 2:30 PM on Jan 23, 2026

DynamoDB Item Structure:
    {
        "pk": "RATELIMIT#{identifier}",
        "sk": "WINDOW#2026-01-23-14-30",
        "request_count": 45,
        "ttl": 1706054760  # Unix timestamp (5 minutes after window end)
    }
"""

from __future__ import annotations

import logging
import os
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING

from .config import get_tier_config

if TYPE_CHECKING:
    from mypy_boto3_dynamodb.service_resource import Table

logger = logging.getLogger(__name__)

# Environment configuration
RATE_LIMITS_TABLE = os.environ.get("RATE_LIMITS_TABLE", "mergeguide-rate-limits-dev")
AWS_REGION = os.environ.get("AWS_REGION", "us-east-1")

# TTL duration: how long to keep window records after they expire
# 5 minutes ensures we have history for debugging but don't accumulate cruft
TTL_SECONDS = 300

# Lazy-loaded DynamoDB table resource
_table: Table | None = None


@dataclass
class RateLimitResult:
    """Result of a rate limit check.

    Attributes:
        allowed: Whether the request is allowed.
        limit: The maximum requests allowed per window.
        remaining: Requests remaining in this window.
        reset_at: Unix timestamp when the current window resets.
        retry_after: Seconds until the window resets (for 429 responses).
        current_count: The current request count in this window.
    """

    allowed: bool
    limit: int
    remaining: int
    reset_at: int
    retry_after: int
    current_count: int = 0

    def to_headers(self) -> dict[str, str]:
        """Generate rate limit HTTP headers.

        Returns:
            Dictionary of headers to include in HTTP responses.
            Includes Retry-After only when rate limit is exceeded.
        """
        headers = {
            "X-RateLimit-Limit": str(self.limit),
            "X-RateLimit-Remaining": str(self.remaining),
            "X-RateLimit-Reset": str(self.reset_at),
        }
        if not self.allowed:
            headers["Retry-After"] = str(self.retry_after)
        return headers


def _get_table() -> Table:
    """Get or create the DynamoDB table resource.

    Returns:
        DynamoDB Table resource for rate limits.
    """
    global _table
    if _table is None:
        import boto3

        dynamodb = boto3.resource("dynamodb", region_name=AWS_REGION)
        _table = dynamodb.Table(RATE_LIMITS_TABLE)
    return _table


def _get_current_window_key() -> str:
    """Generate the sort key for the current minute window.

    Returns:
        Sort key in format WINDOW#YYYY-MM-DD-HH-mm
    """
    now = datetime.now(UTC)
    return f"WINDOW#{now.strftime('%Y-%m-%d-%H-%M')}"


def _get_window_end_timestamp() -> int:
    """Calculate when the current minute window ends.

    Returns:
        Unix timestamp of the next minute boundary.
    """
    now = datetime.now(UTC)
    # Round up to next minute
    seconds_until_next_minute = 60 - now.second
    return int(now.timestamp()) + seconds_until_next_minute


def check_rate_limit(identifier: str, tier: str) -> RateLimitResult:
    """Check and increment the rate limit counter for an identifier.

    This function performs an atomic increment on the request counter
    for the current minute window. It uses DynamoDB's ADD operation
    which is atomic and handles concurrent access correctly.

    Args:
        identifier: Unique identifier for rate limiting. Typically:
            - API key hash for authenticated requests
            - Installation ID for GitHub webhooks
            - Organization ID as fallback
        tier: Billing tier (free, trial, team, enterprise).

    Returns:
        RateLimitResult with allow/deny decision and metadata.

    Note:
        This function uses fail-open behavior: if DynamoDB fails,
        the request is allowed and remaining is set to 0.
    """
    window_key = _get_current_window_key()
    reset_at = _get_window_end_timestamp()
    now = int(time.time())

    # Get limit for this tier (using burst limit for more lenient enforcement)
    config = get_tier_config(tier)
    limit = config.burst_limit

    try:
        table = _get_table()

        # Atomic increment using ADD operation
        # This creates the item if it doesn't exist, or increments if it does
        response = table.update_item(
            Key={
                "pk": f"RATELIMIT#{identifier}",
                "sk": window_key,
            },
            UpdateExpression="ADD request_count :inc SET #ttl = :ttl",
            ExpressionAttributeNames={
                "#ttl": "ttl",  # ttl is a reserved word
            },
            ExpressionAttributeValues={
                ":inc": 1,
                ":ttl": now + TTL_SECONDS,
            },
            ReturnValues="UPDATED_NEW",
        )

        count = int(response["Attributes"]["request_count"])
        remaining = max(0, limit - count)
        allowed = count <= limit
        retry_after = reset_at - now if not allowed else 0

        return RateLimitResult(
            allowed=allowed,
            limit=limit,
            remaining=remaining,
            reset_at=reset_at,
            retry_after=retry_after,
            current_count=count,
        )

    except Exception as e:
        # Fail-open: allow the request if rate limiting fails
        # This prevents rate limit infrastructure issues from blocking users
        logger.warning(f"Rate limit check failed, allowing request: {e}")
        return RateLimitResult(
            allowed=True,
            limit=limit,
            remaining=0,  # Signal that we couldn't determine remaining
            reset_at=reset_at,
            retry_after=0,
            current_count=0,
        )


def get_current_usage(identifier: str) -> tuple[int, int]:
    """Get the current usage for an identifier without incrementing.

    Useful for dashboard displays and monitoring.

    Args:
        identifier: The rate limit identifier.

    Returns:
        Tuple of (current_count, reset_at_timestamp).
        Returns (0, reset_at) if no record exists.
    """
    window_key = _get_current_window_key()
    reset_at = _get_window_end_timestamp()

    try:
        table = _get_table()
        response = table.get_item(
            Key={
                "pk": f"RATELIMIT#{identifier}",
                "sk": window_key,
            },
            ProjectionExpression="request_count",
        )

        if "Item" in response:
            count = int(response["Item"]["request_count"])
            return count, reset_at
        return 0, reset_at

    except Exception as e:
        logger.warning(f"Failed to get current usage: {e}")
        return 0, reset_at


def build_rate_limit_headers(result: RateLimitResult) -> dict[str, str]:
    """Build HTTP headers from a rate limit result.

    Convenience function for adding rate limit headers to responses.

    Args:
        result: The RateLimitResult from check_rate_limit().

    Returns:
        Dictionary of header name to header value.
    """
    return result.to_headers()
