#!/usr/bin/env python3
"""
Per-user rate limiting for the FhirFlame demo.
Protects shared AI resources from excessive usage.
"""

import os
import hmac
from dataclasses import dataclass
from typing import Optional, Any

from database import db_manager


@dataclass
class RateLimitResult:
    allowed: bool
    message: str
    remaining_hour: int = 0
    remaining_day: int = 0


class RateLimiter:
    """Tracks and enforces per-user request limits."""

    def __init__(self):
        self.enabled = os.getenv("RATE_LIMIT_ENABLED", "true").lower() == "true"
        self.per_hour = int(os.getenv("RATE_LIMIT_PER_HOUR", "12"))
        self.per_day = int(os.getenv("RATE_LIMIT_PER_DAY", "30"))
        self.batch_per_hour = int(os.getenv("RATE_LIMIT_BATCH_PER_HOUR", "2"))
        self.access_password = os.getenv("DEMO_ACCESS_PASSWORD", "")
        self.password_request_url = self._safe_request_url(
            os.getenv("DEMO_PASSWORD_REQUEST_URL", "https://bartoszlenart.com")
        )

    def check(
        self,
        user_id: str,
        action: str = "process",
        access_password: Optional[str] = None,
    ) -> RateLimitResult:
        """Check whether a user may perform an action without recording it."""
        if not self.enabled:
            return RateLimitResult(True, "", self.per_hour, self.per_day)

        user_id = self._normalize_user_id(user_id)
        hour_count = db_manager.count_rate_limit_events(user_id, action=None, since_seconds=3600)
        day_count = db_manager.count_rate_limit_events(user_id, action=None, since_seconds=86400)

        if action == "batch":
            batch_hour_count = db_manager.count_rate_limit_events(
                user_id, action="batch", since_seconds=3600
            )
            if batch_hour_count >= self.batch_per_hour:
                if self._valid_access_password(access_password):
                    return RateLimitResult(True, "", 0, max(0, self.per_day - day_count))
                return RateLimitResult(
                    allowed=False,
                    message=self._batch_limit_message(self.batch_per_hour, self.password_request_url),
                    remaining_hour=max(0, self.batch_per_hour - batch_hour_count),
                    remaining_day=max(0, self.per_day - day_count),
                )

        if hour_count >= self.per_hour:
            if self._valid_access_password(access_password):
                return RateLimitResult(True, "", 0, max(0, self.per_day - day_count))
            return RateLimitResult(
                allowed=False,
                message=self._hour_limit_message(self.per_hour, self.password_request_url),
                remaining_hour=0,
                remaining_day=max(0, self.per_day - day_count),
            )

        if day_count >= self.per_day:
            if self._valid_access_password(access_password):
                return RateLimitResult(True, "", max(0, self.per_hour - hour_count), 0)
            return RateLimitResult(
                allowed=False,
                message=self._day_limit_message(self.per_day, self.password_request_url),
                remaining_hour=max(0, self.per_hour - hour_count),
                remaining_day=0,
            )

        return RateLimitResult(
            allowed=True,
            message="",
            remaining_hour=max(0, self.per_hour - hour_count - 1),
            remaining_day=max(0, self.per_day - day_count - 1),
        )

    def record(self, user_id: str, action: str = "process") -> None:
        """Record a completed rate-limited action."""
        if not self.enabled:
            return
        user_id = self._normalize_user_id(user_id)
        db_manager.record_rate_limit_event(user_id, action)
        db_manager.cleanup_rate_limit_events(older_than_hours=48)

    def check_and_record(
        self,
        user_id: str,
        action: str = "process",
        access_password: Optional[str] = None,
    ) -> RateLimitResult:
        """Atomically check limits and record the request if allowed."""
        result = self.check(user_id, action, access_password)
        if result.allowed:
            self.record(user_id, action)
        return result

    def _valid_access_password(self, candidate: Optional[str]) -> bool:
        """Validate a configured limit-override password in constant time."""
        if not self.access_password or not candidate:
            return False
        return hmac.compare_digest(str(candidate), self.access_password)

    @staticmethod
    def _safe_request_url(url: str) -> str:
        """Allow only web links in the rate-limit message."""
        normalized = str(url).strip()
        if normalized.startswith(("https://", "http://")):
            return normalized
        return "https://bartoszlenart.com"

    @staticmethod
    def _normalize_user_id(user_id: Optional[str]) -> str:
        if not user_id or not str(user_id).strip():
            return "anonymous"
        return str(user_id).strip()[:255]

    @staticmethod
    def _hour_limit_message(limit: int, request_url: str) -> str:
        return (
            f"🔒 **Rate limit reached**\n\n"
            f"This demo allows **{limit} processing requests per hour** per user "
            f"to protect shared AI resources.\n\n"
            f"Enter the access password below to continue, or "
            f"[request a password]({request_url})."
        )

    @staticmethod
    def _day_limit_message(limit: int, request_url: str) -> str:
        return (
            f"🔒 **Daily limit reached**\n\n"
            f"This demo allows **{limit} processing requests per day** per user.\n\n"
            f"Enter the access password below to continue, or "
            f"[request a password]({request_url})."
        )

    @staticmethod
    def _batch_limit_message(limit: int, request_url: str) -> str:
        return (
            f"🔒 **Batch demo limit reached**\n\n"
            f"Live batch processing is limited to **{limit} runs per hour** per user.\n\n"
            f"Enter the access password below to continue, or "
            f"[request a password]({request_url})."
        )


rate_limiter = RateLimiter()


def get_user_id_from_request(request: Any) -> str:
    """Extract a stable per-user identifier from a Gradio or HTTP request."""
    if request is None:
        return "anonymous"

    headers = getattr(request, "headers", None) or {}
    if hasattr(headers, "get"):
        forwarded = headers.get("x-forwarded-for") or headers.get("X-Forwarded-For")
        if forwarded:
            return forwarded.split(",")[0].strip()

    client = getattr(request, "client", None)
    if client and getattr(client, "host", None):
        return client.host

    return "anonymous"


def enforce_rate_limit(
    user_id: Optional[str],
    action: str = "process",
    access_password: Optional[str] = None,
) -> Optional[str]:
    """
    Check rate limits for a user action.
    Returns an error message if blocked, otherwise None.
    """
    result = rate_limiter.check_and_record(user_id, action, access_password)
    return None if result.allowed else result.message