# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

"""Simple in-process token-bucket rate limiter middleware.

Each distinct (client-IP, route-prefix) pair gets its own bucket.
Limits are configured via RateLimitConfig.  The defaults below are
intentionally generous; tighten them through the config file or CLI.

Endpoints covered:
  /v1/chat/completions      — expensive LLM inference
  /v1/images/               — image generation
  /v1/audio/                — TTS / STT / audio generation
  /v1/video/                — video generation
  /v1/embeddings            — embedding
  /v1/completions           — legacy completions
"""

import time
import threading
from collections import defaultdict
from typing import Dict, Tuple

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


# Per-route-prefix defaults: (max_requests, window_seconds)
_DEFAULT_LIMITS: Dict[str, Tuple[int, int]] = {
    "/v1/chat/completions": (60, 60),
    "/v1/completions":      (60, 60),
    "/v1/images/":          (30, 60),
    "/v1/audio/":           (60, 60),
    "/v1/video/":           (10, 60),
    "/v1/embeddings":       (120, 60),
}

# API prefixes that count against the request queue
_QUEUED_PREFIXES = ("/v1/",)

# Global toggle — set to False to disable rate limiting entirely.
RATE_LIMITING_ENABLED = True


class _Bucket:
    """Fixed-window counter."""
    __slots__ = ("count", "window_start")

    def __init__(self, now: float):
        self.count = 0
        self.window_start = now


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Apply per-IP, per-route-prefix rate limiting to API endpoints."""

    def __init__(self, app, limits: Dict[str, Tuple[int, int]] = None):
        super().__init__(app)
        self._limits = limits or _DEFAULT_LIMITS
        # (client_ip, prefix) → _Bucket
        self._buckets: Dict[Tuple[str, str], _Bucket] = defaultdict(lambda: _Bucket(time.monotonic()))
        self._lock = threading.Lock()

    def _get_prefix(self, path: str) -> str:
        for prefix in self._limits:
            if path.startswith(prefix):
                return prefix
        return ""

    async def dispatch(self, request: Request, call_next):
        if not RATE_LIMITING_ENABLED:
            return await call_next(request)

        path = request.url.path

        # Queue-size enforcement for authenticated API requests
        if any(path.startswith(p) for p in _QUEUED_PREFIXES):
            from codai.queue.manager import queue_manager
            if await queue_manager.is_full():
                return JSONResponse(
                    status_code=429,
                    content={
                        "error": {
                            "message": "Server queue is full. Please retry later.",
                            "type": "rate_limit_error",
                            "code": 429,
                        }
                    },
                    headers={"Retry-After": "5"},
                )

        prefix = self._get_prefix(path)
        if not prefix:
            return await call_next(request)

        max_req, window = self._limits[prefix]
        client_ip = (
            request.headers.get("x-forwarded-for", "").split(",")[0].strip()
            or (request.client.host if request.client else "unknown")
        )
        key = (client_ip, prefix)
        now = time.monotonic()

        with self._lock:
            bucket = self._buckets[key]
            if now - bucket.window_start >= window:
                bucket.count = 0
                bucket.window_start = now
            bucket.count += 1
            count = bucket.count

        remaining = max(0, max_req - count)
        reset_at = int(time.time() + (window - (now - self._buckets[key].window_start)))

        if count > max_req:
            return JSONResponse(
                status_code=429,
                content={
                    "error": {
                        "message": "Rate limit exceeded. Please slow down.",
                        "type": "rate_limit_error",
                        "code": 429,
                    }
                },
                headers={
                    "X-RateLimit-Limit": str(max_req),
                    "X-RateLimit-Remaining": "0",
                    "X-RateLimit-Reset": str(reset_at),
                    "Retry-After": str(window),
                },
            )

        response = await call_next(request)
        response.headers["X-RateLimit-Limit"] = str(max_req)
        response.headers["X-RateLimit-Remaining"] = str(remaining)
        response.headers["X-RateLimit-Reset"] = str(reset_at)
        return response
