"""Model Router — smart routing and fallback for model requests."""

from __future__ import annotations

import logging
import threading
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from cyberbrain26.model_layer.model_registry import (
    ModelRecord,
    ModelRegistry,
    ModelStatus,
)
from cyberbrain26.model_layer.provider_registry import (
    ProviderRegistry,
    ProviderStatus,
)

logger = logging.getLogger(__name__)


class ModelRouterError(Exception):
    """Raised when no suitable model route can be found."""


class RouteStrategy(str, Enum):
    """Routing strategy for model selection."""

    COST_OPTIMIZED = "cost_optimized"
    PERFORMANCE = "performance"
    RELIABILITY = "reliability"
    LOCAL_FIRST = "local_first"
    CAPABILITY = "capability"
    ROUND_ROBIN = "round_robin"


@dataclass
class RouterConfig:
    """Configuration for Model Router."""

    default_strategy: RouteStrategy = RouteStrategy.COST_OPTIMIZED
    fallback_chain: list[str] | None = None
    prefer_local: bool = False
    prefer_free: bool = True
    max_fallback_depth: int = 5
    health_check_interval: float = 60.0
    circuit_breaker_threshold: int = 3
    circuit_breaker_timeout: float = 30.0


@dataclass
class RouteResult:
    """Result of a routing decision."""

    model_id: str
    provider: str
    strategy: RouteStrategy
    score: float
    fallback: bool = False
    reason: str = ""
    timestamp: float = field(default_factory=time.time)


class ModelRouter:
    """Smart model router with fallback and load balancing.

    Routes model requests based on capability requirements, provider
    health, cost, latency, and trust. Implements circuit breaker pattern
    and fallback chains for resilience.
    """

    def __init__(
        self,
        model_registry: ModelRegistry | None = None,
        provider_registry: ProviderRegistry | None = None,
        config: RouterConfig | None = None,
        model_intelligence: Any = None,
    ) -> None:
        self.model_registry = model_registry or ModelRegistry()
        self.provider_registry = provider_registry or ProviderRegistry()
        self.config = config or RouterConfig()
        self._model_intelligence = model_intelligence

        self._circuit_breakers: dict[str, dict[str, Any]] = {}
        self._rr_counters: dict[str, int] = {}
        self._rr_lock = threading.Lock()
        self._route_history: list[RouteResult] = []

    def _score_model(self, model: ModelRecord) -> float:
        """Score a model. Uses model_intelligence if available, else basic scoring."""
        if self._model_intelligence is not None:
            return self._model_intelligence.score_model(model)
        return self._basic_score(model)

    @staticmethod
    def _basic_score(model: ModelRecord) -> float:
        score = 0.0
        if model.free_status.value in ("free", "free_for_test", "free_for_study"):
            score += 25
        if model.trust.value in ("official", "verified"):
            score += 25
        if model.context_window is not None:
            if model.context_window >= 128000:
                score += 20
            elif model.context_window >= 32000:
                score += 15
            elif model.context_window >= 8000:
                score += 10
            elif model.context_window >= 2000:
                score += 5
        if model.status == ModelStatus.ACTIVE:
            score += 10
        return min(score, 100.0) / 100.0

    # ------------------------------------------------------------------
    # Circuit Breaker
    # ------------------------------------------------------------------

    def _check_circuit_breaker(self, provider_name: str) -> bool:
        cb = self._circuit_breakers.get(provider_name)
        if cb is None:
            return True
        if not cb.get("open", False):
            return True
        if time.time() - cb.get("last_failure", 0) > self.config.circuit_breaker_timeout:
            cb["open"] = False
            cb["failures"] = 0
            logger.info("Circuit breaker closed for %s", provider_name)
            return True
        return False

    def _record_failure(self, provider_name: str) -> None:
        cb = self._circuit_breakers.setdefault(
            provider_name,
            {"failures": 0, "last_failure": 0.0, "open": False},
        )
        cb["failures"] += 1
        cb["last_failure"] = time.time()
        if cb["failures"] >= self.config.circuit_breaker_threshold:
            cb["open"] = True
            logger.warning("Circuit breaker OPEN for %s", provider_name)

    def _record_success(self, provider_name: str) -> None:
        cb = self._circuit_breakers.get(provider_name)
        if cb:
            cb["failures"] = 0
            cb["open"] = False

    # ------------------------------------------------------------------
    # Routing
    # ------------------------------------------------------------------

    @staticmethod
    def _is_free_eligible(model: ModelRecord) -> bool:
        """Check if a model is eligible for free-only usage."""
        if model is None:
            return False
        # Check free status (FREE, FREE_FOR_TEST, FREE_FOR_STUDY)
        if model.free_status.value in ("free", "free_for_test", "free_for_study"):
            return True
        # Also check if model is marked as free in registry
        if hasattr(model, 'free_status') and model.free_status.value in ("free", "free_for_test", "free_for_study"):
            return True
        return False

    def route(
        self,
        *,
        model_id: str | None = None,
        provider: str | None = None,
        capabilities: dict[str, bool] | None = None,
        prefer_local: bool | None = None,
        prefer_free: bool | None = None,
        strategy: RouteStrategy | None = None,
        free_only: bool = False,
    ) -> RouteResult:
        strategy = strategy or self.config.default_strategy
        prefer_local = prefer_local if prefer_local is not None else self.config.prefer_local
        prefer_free = prefer_free if prefer_free is not None else self.config.prefer_free

        requested_model = None
        if model_id:
            requested_model = self.model_registry.get(model_id)
            route = self._route_specific_model(
                model_id, provider, capabilities, strategy, free_only,
            )
            if route:
                self._route_history.append(route)
                return route

        candidates = self._get_candidates(provider, capabilities, free_only)

        if not candidates:
            raise ModelRouterError("No suitable models found for the given criteria")

        if strategy == RouteStrategy.COST_OPTIMIZED:
            route = self._route_cost_optimized(candidates, prefer_local, prefer_free)
        elif strategy == RouteStrategy.PERFORMANCE:
            route = self._route_performance(candidates)
        elif strategy == RouteStrategy.RELIABILITY:
            route = self._route_reliability(candidates)
        elif strategy == RouteStrategy.LOCAL_FIRST:
            route = self._route_local_first(candidates)
        elif strategy == RouteStrategy.CAPABILITY:
            route = self._route_by_capability(candidates, capabilities)
        elif strategy == RouteStrategy.ROUND_ROBIN:
            route = self._route_round_robin(candidates)
        else:
            route = self._route_cost_optimized(candidates, prefer_local, prefer_free)

        # Honest reason: if a specific model was requested but not used due to free_only,
        # record the substitution instead of the generic strategy reason.
        if requested_model is not None and free_only and not self._is_free_eligible(requested_model):
            if route.model_id != requested_model.model_id:
                route.reason = (
                    f"specific model {requested_model.model_id!r} requested with free_only=True "
                    f"but not free-eligible; substituted with {route.model_id!r} ({route.reason})"
                )

        self._route_history.append(route)
        return route

    def route_with_fallback(
        self,
        *,
        model_id: str | None = None,
        capabilities: dict[str, bool] | None = None,
        max_depth: int | None = None,
        free_only: bool = False,
    ) -> list[RouteResult]:
        max_depth = max_depth or self.config.max_fallback_depth
        routes: list[RouteResult] = []

        try:
            primary = self.route(
                model_id=model_id,
                capabilities=capabilities,
                free_only=free_only,
            )
            routes.append(primary)
            self._record_success(primary.provider)
            return routes
        except ModelRouterError:
            pass

        fallback_models = self._get_fallback_candidates(capabilities, free_only)
        for model in fallback_models[:max_depth]:
            try:
                route = self.route(
                    model_id=model.model_id,
                    capabilities=capabilities,
                    free_only=free_only,
                )
                routes.append(route)
                self._record_success(route.provider)
                return routes
            except ModelRouterError:
                self._record_failure(model.provider)
                continue

        if not routes:
            raise ModelRouterError(
                "All fallback routes exhausted for the given criteria"
            )

        return routes

    # ------------------------------------------------------------------
    # Strategy implementations
    # ------------------------------------------------------------------

    def _route_specific_model(
        self,
        model_id: str,
        provider: str | None,
        capabilities: dict[str, bool] | None,
        strategy: RouteStrategy,
        free_only: bool = False,
    ) -> RouteResult | None:
        model = self.model_registry.get(model_id)
        if model is None:
            return None
        if model.status != ModelStatus.ACTIVE:
            return None
        if free_only and not self._is_free_eligible(model):
            logger.warning(
                "refusing specific model %r: not free-eligible under free_only",
                model_id,
            )
            return None

        provider_name = provider or model.provider
        provider_rec = self.provider_registry.get(provider_name)
        if provider_rec is None:
            return None
        if provider_rec.status == ProviderStatus.BLOCKED:
            return None
        if not self._check_circuit_breaker(provider_name):
            return None

        return RouteResult(
            model_id=model_id,
            provider=provider_name,
            strategy=strategy,
            score=self._score_model(model),
            reason="specific model requested",
        )

    def _get_candidates(
        self,
        provider: str | None,
        capabilities: dict[str, bool] | None,
        free_only: bool = False,
    ) -> list[ModelRecord]:
        if provider:
            candidates = self.model_registry.list_by_provider(provider)
        else:
            candidates = self.model_registry.list_active()

        if free_only:
            candidates = [m for m in candidates if self._is_free_eligible(m)]

        if capabilities:
            filtered = []
            for m in candidates:
                match = True
                for cap, required in capabilities.items():
                    if required:
                        actual = getattr(m.capabilities, cap, False)
                        if not actual:
                            match = False
                            break
                if match:
                    filtered.append(m)
            candidates = filtered

        candidates = [
            m for m in candidates
            if self._check_circuit_breaker(m.provider)
        ]

        return candidates

    def _get_fallback_candidates(
        self,
        capabilities: dict[str, bool] | None,
        free_only: bool = False,
    ) -> list[ModelRecord]:
        candidates = self.model_registry.list_active()

        if free_only:
            candidates = [m for m in candidates if self._is_free_eligible(m)]

        if capabilities:
            filtered = []
            for m in candidates:
                match = True
                for cap, required in capabilities.items():
                    if required:
                        actual = getattr(m.capabilities, cap, False)
                        if not actual:
                            match = False
                            break
                if match:
                    filtered.append(m)
            candidates = filtered

        candidates.sort(key=lambda m: self._score_model(m), reverse=True)
        return candidates

    def _route_cost_optimized(
        self,
        candidates: list[ModelRecord],
        prefer_local: bool,
        prefer_free: bool,
    ) -> RouteResult:
        scored = []
        for m in candidates:
            score = self._score_model(m)
            if prefer_local and m.local_available:
                score += 0.15
            if prefer_free and m.is_free:
                score += 0.1
            scored.append((m, score))

        scored.sort(key=lambda x: x[1], reverse=True)
        best_model, best_score = scored[0]

        return RouteResult(
            model_id=best_model.model_id,
            provider=best_model.provider,
            strategy=RouteStrategy.COST_OPTIMIZED,
            score=min(best_score, 1.0),
            reason="cost optimized",
        )

    def _route_performance(self, candidates: list[ModelRecord]) -> RouteResult:
        scored = [(m, self._score_model(m)) for m in candidates]
        scored.sort(key=lambda x: x[1], reverse=True)
        best_model, best_score = scored[0]

        return RouteResult(
            model_id=best_model.model_id,
            provider=best_model.provider,
            strategy=RouteStrategy.PERFORMANCE,
            score=best_score,
            reason="performance optimized",
        )

    def _route_reliability(self, candidates: list[ModelRecord]) -> RouteResult:
        scored = []
        for m in candidates:
            score = self._score_model(m)
            if m.trust.value in ("official", "verified"):
                score += 0.2
            scored.append((m, score))
        scored.sort(key=lambda x: x[1], reverse=True)
        best_model, best_score = scored[0]

        return RouteResult(
            model_id=best_model.model_id,
            provider=best_model.provider,
            strategy=RouteStrategy.RELIABILITY,
            score=min(best_score, 1.0),
            reason="reliability optimized",
        )

    def _route_local_first(self, candidates: list[ModelRecord]) -> RouteResult:
        local = [m for m in candidates if m.local_available]
        pool = local if local else candidates

        scored = [(m, self._score_model(m)) for m in pool]
        scored.sort(key=lambda x: x[1], reverse=True)
        best_model, best_score = scored[0]

        return RouteResult(
            model_id=best_model.model_id,
            provider=best_model.provider,
            strategy=RouteStrategy.LOCAL_FIRST,
            score=best_score,
            fallback=not bool(local),
            reason="local first" if not local else "local model available",
        )

    def _route_by_capability(
        self,
        candidates: list[ModelRecord],
        capabilities: dict[str, bool] | None,
    ) -> RouteResult:
        if not capabilities:
            return self._route_performance(candidates)

        scored = []
        for m in candidates:
            score = self._score_model(m)
            match_count = 0
            total_count = 0
            for cap, required in capabilities.items():
                total_count += 1
                actual = getattr(m.capabilities, cap, False)
                if required and actual:
                    match_count += 1
            if total_count > 0:
                match_ratio = match_count / total_count
                score = score * 0.5 + match_ratio * 0.5
            scored.append((m, score))

        scored.sort(key=lambda x: x[1], reverse=True)
        best_model, best_score = scored[0]

        return RouteResult(
            model_id=best_model.model_id,
            provider=best_model.provider,
            strategy=RouteStrategy.CAPABILITY,
            score=best_score,
            reason="capability optimized",
        )

    def _route_round_robin(self, candidates: list[ModelRecord]) -> RouteResult:
        if not candidates:
            raise ModelRouterError("No candidates for round-robin routing")

        provider = candidates[0].provider
        # One router instance is shared by the whole gateway, so the read and
        # the increment have to be a single critical section; two threads
        # otherwise land on the same candidate and lose a count.
        with self._rr_lock:
            idx = self._rr_counters.get(provider, 0)
            selected = candidates[idx % len(candidates)]
            self._rr_counters[provider] = idx + 1

        return RouteResult(
            model_id=selected.model_id,
            provider=selected.provider,
            strategy=RouteStrategy.ROUND_ROBIN,
            score=self._score_model(selected),
            reason="round-robin",
        )

    # ------------------------------------------------------------------
    # Query helpers
    # ------------------------------------------------------------------

    def get_route_history(self, limit: int | None = None) -> list[RouteResult]:
        if limit:
            return self._route_history[-limit:]
        return list(self._route_history)

    def get_circuit_breaker_status(self, provider_name: str | None = None) -> dict[str, Any]:
        if provider_name:
            cb = self._circuit_breakers.get(provider_name, {})
            return {provider_name: cb}
        return dict(self._circuit_breakers)

    def reset_circuit_breaker(self, provider_name: str) -> bool:
        if provider_name in self._circuit_breakers:
            self._circuit_breakers[provider_name] = {
                "failures": 0,
                "last_failure": 0.0,
                "open": False,
            }
            return True
        return False