"""
router.py — Intelligent provider routing for OneFirewall AI.

Features:
- Latency tracking (health check)
- Cost-based scoring
- Automatic fallback if preferred provider fails
- Support for OpenAI, Gemini, and Local (Ollama)
"""

import os
import time
import logging
import requests
import threading
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any

logger = logging.getLogger(__name__)

@dataclass
class Provider:
    name: str
    endpoint: str
    api_key_env: str
    cost_per_1k: float = 0.01  # USD
    big_model: str = "gpt-4o"
    small_model: str = "gpt-4o-mini"
    latency_ms: float = 9999.0
    healthy: bool = True
    request_count: int = 0
    error_count: int = 0
    avg_latency_ms: float = 9999.0

    @property
    def api_key(self) -> Optional[str]:
        # Check env first, then config (this will be injected by SmartRouter)
        return os.getenv(self.api_key_env) or self._api_key_override

    _api_key_override: Optional[str] = None

    @property
    def is_configured(self) -> bool:
        if "localhost" in self.endpoint or "127.0.0.1" in self.endpoint:
            return True
        return bool(self.api_key)

    def score(self, strategy: str = "balanced") -> float:
        if not self.healthy or not self.is_configured:
            return float("inf")
        
        # Lower score = better
        latency_score = self.avg_latency_ms / 1000.0
        cost_score = self.cost_per_1k * 100
        error_penalty = (self.error_count / (self.request_count + 1)) * 500
        
        if strategy == "speed":
            return latency_score + error_penalty
        elif strategy == "cost":
            return cost_score + error_penalty
        else:
            return (latency_score * 0.5) + (cost_score * 0.5) + error_penalty

class SmartRouter:
    """
    Intelligently routes requests to the best available LLM provider.
    Adapted from the upstream SmartRouter logic.
    """
    def __init__(self, strategy: str = "balanced"):
        self.strategy = strategy
        self.providers: List[Provider] = [
            Provider(
                name="onefirewall-gateway",
                endpoint="https://api.onefirewall.com", # Placeholder, updated from config
                api_key_env="ONEFIREWALL_API_KEY",
                cost_per_1k=0.005,
                big_model="openai/gpt-4o",
                small_model="openai/gpt-4o-mini",
            ),
            Provider(
                name="openai",
                endpoint="https://api.openai.com/v1",
                api_key_env="OPENAI_API_KEY",
                cost_per_1k=0.01,
                big_model="gpt-4o",
                small_model="gpt-4o-mini",
            ),
            Provider(
                name="gemini",
                endpoint="https://generativelanguage.googleapis.com/v1",
                api_key_env="GEMINI_API_KEY",
                cost_per_1k=0.002,
                big_model="gemini-1.5-pro",
                small_model="gemini-1.5-flash",
            ),
            Provider(
                name="ollama",
                endpoint="http://localhost:11434/v1",
                api_key_env="",
                cost_per_1k=0.0,
                big_model="llama3:8b",
                small_model="llama3:8b",
            )
        ]
        self._initialized = False

    def initialize(self, config_access: Dict[str, Any] = None):
        """Build initial health/latency data."""
        if config_access:
            # Inject primary gateway from config
            for p in self.providers:
                if p.name == "onefirewall-gateway":
                    p.endpoint = config_access.get("endpoint", p.endpoint)
                    p._api_key_override = config_access.get("api_key")
        
        # Ping all in parallel
        threads = []
        for p in self.providers:
            t = threading.Thread(target=self._ping_provider, args=(p,))
            t.start()
            threads.append(t)
        
        for t in threads:
            t.join(timeout=3.0)
        
        self._initialized = True

    def _ping_provider(self, provider: Provider):
        if not provider.is_configured:
            provider.healthy = False
            return

        start = time.monotonic()
        try:
            # Use models endpoint for broad compatibility
            url = provider.endpoint.rstrip("/") + "/models"
            headers = {}
            if provider.api_key:
                headers["Authorization"] = f"Bearer {provider.api_key}"
            
            resp = requests.get(url, headers=headers, timeout=5)
            # 200 is healthy, 401/403/400 often means reachable but key/param issue, 
            # still marks as "reachable" for router
            if resp.status_code in (200, 400, 401, 403):
                elapsed = (time.monotonic() - start) * 1000
                provider.healthy = True
                provider.latency_ms = elapsed
                provider.avg_latency_ms = elapsed
            else:
                provider.healthy = False
        except Exception:
            provider.healthy = False

    def get_best_provider(self, model_pref: str = None) -> Provider:
        """Select the best healthy/configured provider."""
        available = [p for p in self.providers if p.healthy and p.is_configured]
        if not available:
            # Fallback to gateway if everything else is dead
            return self.providers[0]
            
        return min(available, key=lambda p: p.score(self.strategy))

    def record_result(self, provider_name: str, success: bool, duration_ms: float):
        for p in self.providers:
            if p.name == provider_name:
                p.request_count += 1
                if success:
                    # Rolling average
                    alpha = 0.3
                    p.avg_latency_ms = (alpha * duration_ms) + ((1 - alpha) * p.avg_latency_ms)
                else:
                    p.error_count += 1
                    if p.error_count > 3:
                        p.healthy = False
                break
