# 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


class BearerAuthMiddleware(BaseHTTPMiddleware):
    """Reject /v1/ API requests that lack a valid Bearer token or active web session."""

    _EXEMPT_PATHS = {"/v1/images/progress"}

    async def dispatch(self, request: Request, call_next):
        path = request.url.path
        if not path.startswith("/v1/") or path in self._EXEMPT_PATHS:
            return await call_next(request)

        # Requests from the ASGI broker bridge are in-process and have no real
        # Bearer token. Identify them by the sentinel server tuple set in asgi_bridge.py.
        if request.scope.get("server") == ("internal", 80):
            return await call_next(request)

        from codai.admin import routes as _admin_routes
        sm = _admin_routes.session_manager
        if sm is None:
            return await call_next(request)

        # Accept a valid Bearer token
        auth_header = request.headers.get("authorization", "")
        if auth_header.lower().startswith("bearer "):
            token = auth_header[7:].strip()
            if sm.verify_token(token):
                return await call_next(request)

        # Accept a valid web session cookie (logged-in browser user)
        cookie = request.cookies.get("session", "")
        if cookie.endswith(".MUST_CHANGE"):
            cookie = cookie[:-12]
        if cookie and sm.validate_session(cookie):
            return await call_next(request)

        return JSONResponse(
            status_code=401,
            content={
                "error": {
                    "message": "Invalid API key. Provide a valid Bearer token.",
                    "type": "invalid_request_error",
                    "code": "invalid_api_key",
                }
            },
        )


# 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 ""

    # Lightweight polling endpoints that must never be rate-limited
    _EXEMPT_PATHS = {"/v1/images/progress"}

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

        path = request.url.path

        if path in self._EXEMPT_PATHS:
            return await call_next(request)

        # Queue-size enforcement for authenticated API requests (not for status polls)
        if path not in self._EXEMPT_PATHS and 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
