models/grid_router.py

"""
backend/models/grid_router.py — Grid Orchestration Router (S766-GRID-1)

Rotazione intelligente dei token tra profili A, B, C, D con health tracking,
rate-limit detection e active balancing.

Architettura:
- HealthTracker: Monitora TTFT, errori, rate-limit per ogni provider/profilo
- GridRouter: Seleziona il miglior profilo basato su metriche recenti
- RateLimitDetector: Identifica quando un profilo è saturo e lo esclude temporaneamente
"""

import os
import time
import asyncio
import logging
from typing import Optional, Dict, List, Tuple
from dataclasses import dataclass, field
from enum import Enum
import json

_logger = logging.getLogger("grid_router")

# ── Configurazione ─────────────────────────────────────────────────────────
GRID_HEALTH_WINDOW_S = 300  # Finestra di 5 minuti per calcolare metriche
GRID_RATE_LIMIT_COOLDOWN_S = 60  # Escludere un profilo per 60s se rate-limited
GRID_ERROR_THRESHOLD = 5  # Escludere se 5+ errori negli ultimi GRID_HEALTH_WINDOW_S
GRID_ENABLE = os.getenv("GRID_ENABLE", "true").lower() == "true"


class ProviderStatus(Enum):
    """Stato di un provider nella Grid."""
    HEALTHY = "healthy"
    DEGRADED = "degraded"
    RATE_LIMITED = "rate_limited"
    OFFLINE = "offline"


@dataclass
class ProviderMetric:
    """Metrica di salute per un singolo provider/profilo."""
    provider_name: str
    profile: str  # "A", "B", "C", "D"
    ttft_ms: float = 0.0  # Time To First Token
    success_count: int = 0
    error_count: int = 0
    rate_limit_count: int = 0
    last_error: Optional[str] = None
    last_rate_limit_time: float = 0.0
    status: ProviderStatus = ProviderStatus.HEALTHY
    last_updated: float = field(default_factory=time.time)

    def is_rate_limited(self) -> bool:
        """Verifica se il provider è attualmente rate-limited."""
        if self.rate_limit_count == 0:
            return False
        elapsed = time.time() - self.last_rate_limit_time
        return elapsed < GRID_RATE_LIMIT_COOLDOWN_S

    def health_score(self) -> float:
        """Calcola un score di salute (0-100)."""
        if self.is_rate_limited():
            return 0.0
        if self.status == ProviderStatus.OFFLINE:
            return 0.0
        if self.status == ProviderStatus.RATE_LIMITED:
            return 10.0

        total_requests = self.success_count + self.error_count
        if total_requests == 0:
            return 50.0  # Neutrale se non testato

        success_rate = self.success_count / total_requests
        # Penalità per TTFT alto (>1000ms)
        ttft_penalty = min(self.ttft_ms / 1000.0, 1.0) * 20.0
        score = (success_rate * 100.0) - ttft_penalty
        return max(0.0, min(100.0, score))


class HealthTracker:
    """Traccia la salute di tutti i provider/profili nella Grid."""

    def __init__(self):
        self.metrics: Dict[str, ProviderMetric] = {}
        self._lock = asyncio.Lock()

    async def record_success(
        self,
        provider_name: str,
        profile: str,
        ttft_ms: float,
    ):
        """Registra un successo."""
        async with self._lock:
            key = f"{provider_name}:{profile}"
            if key not in self.metrics:
                self.metrics[key] = ProviderMetric(provider_name, profile)

            metric = self.metrics[key]
            metric.success_count += 1
            metric.ttft_ms = (metric.ttft_ms * 0.7) + (ttft_ms * 0.3)  # EMA
            metric.error_count = max(0, metric.error_count - 1)  # Decadimento errori
            metric.last_updated = time.time()

            if metric.error_count == 0:
                metric.status = ProviderStatus.HEALTHY

    async def record_error(
        self,
        provider_name: str,
        profile: str,
        error: str,
    ):
        """Registra un errore."""
        async with self._lock:
            key = f"{provider_name}:{profile}"
            if key not in self.metrics:
                self.metrics[key] = ProviderMetric(provider_name, profile)

            metric = self.metrics[key]
            metric.error_count += 1
            metric.last_error = error
            metric.last_updated = time.time()

            if metric.error_count >= GRID_ERROR_THRESHOLD:
                metric.status = ProviderStatus.OFFLINE

    async def record_rate_limit(
        self,
        provider_name: str,
        profile: str,
    ):
        """Registra un rate-limit."""
        async with self._lock:
            key = f"{provider_name}:{profile}"
            if key not in self.metrics:
                self.metrics[key] = ProviderMetric(provider_name, profile)

            metric = self.metrics[key]
            metric.rate_limit_count += 1
            metric.last_rate_limit_time = time.time()
            metric.status = ProviderStatus.RATE_LIMITED

    async def get_best_provider(
        self,
        provider_name: str,
        profiles: List[str] = None,
    ) -> Optional[str]:
        """
        Restituisce il profilo migliore per un provider.
        Profili: ["A", "B", "C", "D"]
        """
        if profiles is None:
            profiles = ["A", "B", "C", "D"]

        async with self._lock:
            best_profile = None
            best_score = -1.0

            for profile in profiles:
                key = f"{provider_name}:{profile}"
                metric = self.metrics.get(key)

                if metric is None:
                    # Non testato — assegna score neutrale
                    score = 50.0
                else:
                    score = metric.health_score()

                if score > best_score:
                    best_score = score
                    best_profile = profile

            return best_profile

    async def get_metrics_summary(self) -> Dict:
        """Restituisce un riepilogo delle metriche per il monitoraggio."""
        async with self._lock:
            summary = {}
            for key, metric in self.metrics.items():
                summary[key] = {
                    "status": metric.status.value,
                    "health_score": metric.health_score(),
                    "ttft_ms": round(metric.ttft_ms, 2),
                    "success_count": metric.success_count,
                    "error_count": metric.error_count,
                    "rate_limit_count": metric.rate_limit_count,
                    "last_error": metric.last_error,
                }
            return summary


class GridRouter:
    """Router intelligente per la Grid di profili multi-account."""

    def __init__(self):
        self.health_tracker = HealthTracker()
        self._enabled = GRID_ENABLE

    async def select_provider_config(
        self,
        base_provider: str,  # "groq", "nvidia", "gemini", "openrouter"
    ) -> Tuple[str, str]:
        """
        Seleziona il miglior profilo per un provider.
        Restituisce: (provider_name_con_profilo, api_key_env_var)

        Esempio:
        - Input: "groq"
        - Output: ("groq-b", "GROQ_API_KEY_B")
        """
        if not self._enabled:
            # Fallback: usa il profilo A (default)
            return (base_provider, f"{base_provider.upper()}_API_KEY")

        best_profile = await self.health_tracker.get_best_provider(base_provider)

        if best_profile is None or best_profile == "A":
            # Profilo A è il default
            return (base_provider, f"{base_provider.upper()}_API_KEY")

        # Profili B, C, D
        suffix = f"_{best_profile}"
        provider_name = f"{base_provider}{suffix.lower()}"
        env_var = f"{base_provider.upper()}_API_KEY{suffix}"

        _logger.info(f"Grid Router: {base_provider} → {provider_name} (score-based selection)")
        return (provider_name, env_var)

    async def wrap_provider_call(
        self,
        provider_name: str,
        profile: str,
        coro,
    ):
        """
        Wrapper per le chiamate ai provider che registra metriche.
        S766-GRID-2: Diagnosi avanzata Rate Limit (429/402/Quota).
        """
        start_time = time.time()
        try:
            result = await coro
            ttft_ms = (time.time() - start_time) * 1000
            await self.health_tracker.record_success(provider_name, profile, ttft_ms)
            return result
        except Exception as exc:
            exc_str = str(exc)
            # Diagnosi specifica per blocchi di quota/rate-limit
            is_quota_issue = any(x in exc_str.lower() for x in ["429", "402", "rate_limit", "quota", "depleted", "insufficient_balance"])
            
            if is_quota_issue:
                _logger.warning(f"Grid Router: Profilo {profile} ({provider_name}) in RATE LIMIT/QUOTA. Escludo per {GRID_RATE_LIMIT_COOLDOWN_S}s.")
                await self.health_tracker.record_rate_limit(provider_name, profile)
                # Solleviamo un'eccezione specifica per permettere all'agente di capire che deve ruotare
                raise Exception(f"GRID_RATE_LIMIT:{profile}:{provider_name}") from exc
            else:
                await self.health_tracker.record_error(provider_name, profile, exc_str)
            raise

    async def get_health_status(self) -> Dict:
        """Restituisce lo stato di salute della Grid."""
        return await self.health_tracker.get_metrics_summary()


# ── Singleton globale ──────────────────────────────────────────────────────
_grid_router_instance: Optional[GridRouter] = None


def get_grid_router() -> GridRouter:
    """Restituisce l'istanza globale del GridRouter."""
    global _grid_router_instance
    if _grid_router_instance is None:
        _grid_router_instance = GridRouter()
    return _grid_router_instance