#!/usr/bin/env python3
"""Script to add rate limiting decorators to API endpoints."""

import re

# Read the API file
with open("apps/api/main.py") as f:
    content = f.read()

# Pattern to find endpoint definitions
endpoint_pattern = r"(@app\.(get|post|put|delete)\([^\)]+\))\ndef\s+(\w+)\("


def add_request_param(match):
    """Add Request parameter to function signature if not present."""
    decorator = match.group(1)
    func_name = match.group(3)

    # Skip if already has rate limiting
    if "@limiter.limit" in content[max(0, match.start() - 100) : match.start()]:
        return match.group(0)

    # Determine rate limit based on endpoint type
    if "/admin/" in decorator:
        rate_limit = "20/minute"  # Lower limit for admin endpoints
    elif "predictions" in decorator:
        rate_limit = "30/minute"  # Moderate limit for compute-heavy endpoints
    else:
        rate_limit = "100/minute"  # Higher limit for read endpoints

    # Add rate limiter decorator and Request parameter
    result = f'{decorator}\n@limiter.limit("{rate_limit}")\ndef {func_name}(request: Request, '

    return result


# Apply the transformation
new_content = re.sub(endpoint_pattern, add_request_param, content)

# Write back
with open("apps/api/main.py", "w") as f:
    f.write(new_content)

print("Rate limiting decorators added to all endpoints!")