"""Rate limiter distribuido para llamadas a OpenAI.

Usa Redis como backend para compartir el estado entre múltiples workers.
Implementa el algoritmo de token bucket con ventana deslizante para respetar
los límites de RPM (requests per minute) y TPM (tokens per minute) de OpenAI.
"""
from __future__ import annotations

import asyncio
import logging
import time
from typing import Optional

import redis.asyncio as aioredis

logger = logging.getLogger(__name__)

# Límites por defecto para OpenAI gpt-4o (tier 1).
# Sobreescribibles con variables de entorno vía config.
DEFAULT_RPM = 500       # requests per minute
DEFAULT_TPM = 800_000   # tokens per minute


class DistributedRateLimiter:
    """Token bucket distribuido respaldado por Redis.

    Un solo objeto se comparte entre todos los workers de un proceso.
    Múltiples réplicas del worker comparten el estado en Redis.

    Uso:
        limiter = DistributedRateLimiter(redis_url="redis://redis:6379/1")
        async with limiter.acquire(estimated_tokens=4000):
            response = await openai_client.chat.completions.create(...)
    """

    def __init__(
        self,
        redis_url: str,
        rpm: int = DEFAULT_RPM,
        tpm: int = DEFAULT_TPM,
        key_prefix: str = "ips:rl",
    ) -> None:
        self._redis_url = redis_url
        self.rpm = rpm
        self.tpm = tpm
        self._prefix = key_prefix
        self._client: Optional[aioredis.Redis] = None
        self._lock = asyncio.Lock()

    async def _get_client(self) -> aioredis.Redis:
        if self._client is None:
            async with self._lock:
                if self._client is None:
                    self._client = aioredis.from_url(
                        self._redis_url,
                        encoding="utf-8",
                        decode_responses=True,
                        socket_connect_timeout=5,
                    )
        return self._client

    async def _sliding_window_count(self, key: str, window_s: float = 60.0) -> int:
        """Cuenta eventos en la ventana deslizante usando un Sorted Set de Redis."""
        client = await self._get_client()
        now = time.monotonic()
        # Eliminar eventos expirados
        await client.zremrangebyscore(key, "-inf", now - window_s)
        return await client.zcard(key)

    async def _record_event(self, key: str, tokens: int = 0, window_s: float = 60.0) -> None:
        """Registra un evento en el sorted set. El score es el timestamp."""
        client = await self._get_client()
        now = time.monotonic()
        member = f"{now}:{tokens}"
        pipe = client.pipeline()
        pipe.zadd(key, {member: now})
        pipe.expire(key, int(window_s) + 5)
        await pipe.execute()

    async def _token_sum_in_window(self, key: str, window_s: float = 60.0) -> int:
        """Suma los tokens de todos los eventos en la ventana."""
        client = await self._get_client()
        now = time.monotonic()
        await client.zremrangebyscore(key, "-inf", now - window_s)
        members = await client.zrange(key, 0, -1)
        total = 0
        for m in members:
            parts = m.split(":", 1)
            if len(parts) == 2:
                try:
                    total += int(parts[1])
                except ValueError:
                    pass
        return total

    async def acquire(self, estimated_tokens: int = 2000, max_wait_s: float = 120.0) -> None:
        """Espera hasta que haya capacidad disponible en RPM y TPM.

        Args:
            estimated_tokens: Tokens que se espera consumir (input + output).
            max_wait_s: Tiempo máximo de espera antes de lanzar TimeoutError.
        """
        rpm_key = f"{self._prefix}:rpm"
        tpm_key = f"{self._prefix}:tpm"
        start = time.monotonic()
        sleep_s = 0.5

        while True:
            rpm_count = await self._sliding_window_count(rpm_key)
            tpm_used = await self._token_sum_in_window(tpm_key)

            rpm_ok = rpm_count < self.rpm
            tpm_ok = (tpm_used + estimated_tokens) <= self.tpm

            if rpm_ok and tpm_ok:
                await self._record_event(rpm_key, tokens=0)
                await self._record_event(tpm_key, tokens=estimated_tokens)
                return

            elapsed = time.monotonic() - start
            if elapsed >= max_wait_s:
                raise TimeoutError(
                    f"Rate limiter: espera maxima alcanzada ({max_wait_s}s). "
                    f"rpm={rpm_count}/{self.rpm} tpm={tpm_used}/{self.tpm}"
                )

            reason = "RPM" if not rpm_ok else "TPM"
            logger.debug(
                "Rate limiter: esperando por %s (rpm=%d/%d tpm=%d/%d) — sleep %.1fs",
                reason, rpm_count, self.rpm, tpm_used, self.tpm, sleep_s,
            )
            await asyncio.sleep(sleep_s)
            sleep_s = min(sleep_s * 1.5, 10.0)

    async def close(self) -> None:
        if self._client is not None:
            await self._client.aclose()
            self._client = None


# Instancia global compartida por todos los workers del proceso.
# Se inicializa de forma lazy en el primer uso.
_limiter: Optional[DistributedRateLimiter] = None


def get_rate_limiter() -> DistributedRateLimiter:
    """Devuelve el rate limiter global, inicializándolo si es necesario."""
    global _limiter
    if _limiter is None:
        from domain.config import settings
        _limiter = DistributedRateLimiter(
            redis_url=settings.redis_url,
            rpm=settings.openai_rpm_limit,
            tpm=settings.openai_tpm_limit,
        )
    return _limiter