"""In-memory fixed-window rate limiter for public API endpoints."""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Literal

from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse

from hive_compute.coordinator.auth import resolve_valid_api_key_hash, resolve_valid_worker_id
from hive_compute.coordinator.manager_auth import resolve_valid_manager_session_hash

MAX_BUCKETS = 10_000  # Evict oldest when exceeded to prevent unbounded growth


@dataclass
class _Bucket:
    window_start: float
    requests: int
    capacity: int

    def consume(self) -> bool:
        now = time.monotonic()
        if now - self.window_start >= 60:
            self.window_start = now
            self.requests = 0
        if self.requests < self.capacity:
            self.requests += 1
            return True
        return False

    @property
    def retry_after(self) -> int:
        if self.requests < self.capacity:
            return 0
        return max(1, int(60 - (time.monotonic() - self.window_start)))

    @property
    def remaining(self) -> int:
        return max(0, self.capacity - self.requests)


@dataclass(frozen=True)
class RateLimitRule:
    """Per-route rate limit configuration."""

    path_prefix: str
    requests_per_minute: int
    scope: Literal["api_key", "worker", "global", "manager_session"]
    methods: frozenset[str] | None = None

    def matches(self, method: str, path: str) -> bool:
        if self.methods and method not in self.methods:
            return False
        return path.startswith(self.path_prefix)

    @property
    def bucket_name(self) -> str:
        methods = ",".join(sorted(self.methods)) if self.methods else "ALL"
        return f"{methods}:{self.path_prefix}"


DEFAULT_RULES: tuple[RateLimitRule, ...] = (
    RateLimitRule(
        path_prefix="/api/v1/workers/register",
        requests_per_minute=10,
        scope="global",
        methods=frozenset({"POST"}),
    ),
    RateLimitRule(
        path_prefix="/api/v1/workers/poll",
        requests_per_minute=15,
        scope="worker",
        methods=frozenset({"GET"}),
    ),
    RateLimitRule(
        path_prefix="/manager/api/",
        requests_per_minute=30,
        scope="manager_session",
    ),
    RateLimitRule(
        path_prefix="/api/v1/jobs",
        requests_per_minute=60,
        scope="api_key",
    ),
)

# Paths exempt from rate limiting. Worker job lifecycle endpoints are exempt
# because they're already bounded by job ownership rather than request volume.
EXEMPT_PATHS = {"/health", "/docs", "/openapi.json", "/static"}
EXEMPT_SUFFIXES = ["/heartbeat", "/result", "/fail"]


class RateLimitMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, rules: tuple[RateLimitRule, ...] | None = None):
        super().__init__(app)
        self.rules = tuple(sorted(rules or DEFAULT_RULES, key=lambda rule: len(rule.path_prefix), reverse=True))
        self._buckets: dict[str, _Bucket] = {}

    async def dispatch(self, request: Request, call_next) -> Response:
        path = request.url.path

        if any(path.startswith(prefix) for prefix in EXEMPT_PATHS):
            return await call_next(request)
        if request.method == "POST" and any(path.endswith(suffix) for suffix in EXEMPT_SUFFIXES):
            return await call_next(request)

        rule = self._match_rule(request.method, path)
        if not rule:
            return await call_next(request)

        identity = _resolve_identity(request, rule)
        bucket = self._get_bucket(f"{rule.bucket_name}:{identity}", rule)

        if not bucket.consume():
            return JSONResponse(
                status_code=429,
                content={"detail": "Rate limit exceeded"},
                headers={
                    "Retry-After": str(bucket.retry_after),
                    "X-RateLimit-Limit": str(rule.requests_per_minute),
                    "X-RateLimit-Remaining": "0",
                },
            )

        response = await call_next(request)
        response.headers["X-RateLimit-Limit"] = str(rule.requests_per_minute)
        response.headers["X-RateLimit-Remaining"] = str(bucket.remaining)
        return response

    def _match_rule(self, method: str, path: str) -> RateLimitRule | None:
        for rule in self.rules:
            if rule.matches(method, path):
                return rule
        return None

    def _get_bucket(self, key: str, rule: RateLimitRule) -> _Bucket:
        if key not in self._buckets:
            if len(self._buckets) >= MAX_BUCKETS:
                oldest_key = min(self._buckets, key=lambda bucket_key: self._buckets[bucket_key].window_start)
                del self._buckets[oldest_key]

            self._buckets[key] = _Bucket(
                window_start=time.monotonic(),
                requests=0,
                capacity=rule.requests_per_minute,
            )
        return self._buckets[key]


def _resolve_identity(request: Request, rule: RateLimitRule) -> str:
    client_ip = _get_client_ip(request)

    if rule.scope == "global":
        return "global"
    if rule.scope == "api_key":
        api_key_hash = _resolve_valid_api_key_hash(request)
        return f"api-key:{api_key_hash}" if api_key_hash else f"ip:{client_ip}"
    if rule.scope == "worker":
        worker_id = _resolve_valid_worker_id(request)
        return f"worker:{worker_id}" if worker_id else f"ip:{client_ip}"
    if rule.scope == "manager_session":
        session_hash = _resolve_valid_manager_session_hash(request)
        return f"manager-session:{session_hash}" if session_hash else f"ip:{client_ip}"
    return f"ip:{client_ip}"


def _resolve_valid_api_key_hash(request: Request) -> str | None:
    return resolve_valid_api_key_hash(request, request.headers.get("authorization"))


def _resolve_valid_worker_id(request: Request) -> str | None:
    return resolve_valid_worker_id(
        request,
        authorization=request.headers.get("authorization"),
        x_worker_id=request.headers.get("x-worker-id"),
    )


def _resolve_valid_manager_session_hash(request: Request) -> str | None:
    return resolve_valid_manager_session_hash(request)


def _get_client_ip(request: Request) -> str:
    """Extract client IP using Fly-Client-IP (trusted, set by Fly's edge proxy).

    Falls back to request.client.host for direct connections (local dev).
    Does NOT trust X-Forwarded-For since it's spoofable by external callers.
    """
    fly_ip = request.headers.get("fly-client-ip")
    if fly_ip:
        return fly_ip.strip()
    return request.client.host if request.client else "unknown"
