"""Model Cascade Router - Module 3.

Routes agent requests through a cascade of models:
tiny local → cheap small cloud → medium → frontier → specialist.

Supports:
A. always frontier
B. static routing
C. prompt-only router
D. trained cost-aware router
E. trained router + verifier fallback
"""

import random
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass

from .trace_schema import TaskType, Outcome
from .config import ACOConfig, ModelConfig
from .classifier import TaskPrediction


@dataclass
class RoutingDecision:
    model_id: str
    provider: str
    tier: int
    confidence: float
    reasoning: str
    fallback_model_id: Optional[str] = None
    use_verifier: bool = False
    max_tokens: int = 4096
    temperature: float = 0.7


class ModelCascadeRouter:
    """Routes tasks to the cheapest acceptable model."""

    TIER_ORDER = [1, 2, 3, 4, 5]  # tiny → cheap → medium → frontier → specialist

    def __init__(self, config: ACOConfig):
        self.config = config
        self.models_by_tier: Dict[int, List[ModelConfig]] = {t: [] for t in self.TIER_ORDER}
        self._build_tier_index()
        self.routing_stats: Dict[str, Dict] = {}
        self.decision_history: List[Dict] = []

    def _build_tier_index(self):
        for name, mc in self.config.models.items():
            self.models_by_tier.setdefault(mc.strength_tier, []).append(mc)

    def route(self, task_prediction: TaskPrediction, routing_mode: str = "cascade") -> RoutingDecision:
        """Select model based on task prediction and routing policy."""
        
        if routing_mode == "always_frontier":
            return self._route_always_frontier(task_prediction)
        elif routing_mode == "static":
            return self._route_static(task_prediction)
        elif routing_mode == "prompt_only":
            return self._route_prompt_only(task_prediction)
        elif routing_mode == "learned":
            return self._route_learned(task_prediction)
        elif routing_mode == "learned_verifier":
            return self._route_learned(task_prediction, verifier_fallback=True)
        else:
            return self._route_cascade(task_prediction)

    def _route_always_frontier(self, prediction: TaskPrediction) -> RoutingDecision:
        frontier = self.models_by_tier.get(4, [])
        if not frontier:
            frontier = self.models_by_tier.get(5, [])
        if not frontier:
            frontier = self.models_by_tier.get(3, [])
        model = frontier[0] if frontier else list(self.config.models.values())[0]
        return RoutingDecision(
            model_id=model.model_id,
            provider=model.provider,
            tier=4,
            confidence=1.0,
            reasoning="Always frontier policy",
            max_tokens=min(prediction.expected_cost * 50000, model.max_context),
        )

    def _route_static(self, prediction: TaskPrediction) -> RoutingDecision:
        # Static mapping: task type -> tier
        static_map = {
            TaskType.QUICK_ANSWER: 1,
            TaskType.UNKNOWN_AMBIGUOUS: 2,
            TaskType.TOOL_HEAVY: 2,
            TaskType.RETRIEVAL_HEAVY: 2,
            TaskType.DOCUMENT_DRAFTING: 3,
            TaskType.CODING: 3,
            TaskType.RESEARCH: 4,
            TaskType.LONG_HORIZON: 4,
            TaskType.LEGAL_REGULATED: 5,
        }
        tier = static_map.get(prediction.task_type, 3)
        models = self.models_by_tier.get(tier, self.models_by_tier[3])
        model = models[0] if models else list(self.config.models.values())[0]
        return RoutingDecision(
            model_id=model.model_id,
            provider=model.provider,
            tier=tier,
            confidence=0.6,
            reasoning=f"Static routing: {prediction.task_type.value} -> tier {tier}",
            fallback_model_id=self._next_tier_model(tier).model_id if tier < 5 else None,
        )

    def _route_prompt_only(self, prediction: TaskPrediction) -> RoutingDecision:
        """Use prompt heuristics to decide model tier."""
        # Heuristic: if risk > 0.7 or expected tier >= 4, use frontier
        if prediction.risk_of_failure > 0.7 or prediction.expected_model_tier >= 4:
            tier = 4
        elif prediction.expected_model_tier <= 2:
            tier = max(prediction.expected_model_tier, 1)
        else:
            # Start cheap, escalate on low confidence
            tier = max(prediction.expected_model_tier - 1, 1)
        
        models = self.models_by_tier.get(tier, self.models_by_tier[3])
        model = models[0] if models else list(self.config.models.values())[0]
        
        fallback = None
        if tier < 5 and prediction.risk_of_failure > 0.5:
            fallback = self._next_tier_model(tier)
            fallback_id = fallback.model_id if fallback else None
        else:
            fallback_id = None
        
        return RoutingDecision(
            model_id=model.model_id,
            provider=model.provider,
            tier=tier,
            confidence=1.0 - prediction.risk_of_failure,
            reasoning=f"Prompt heuristic: risk={prediction.risk_of_failure:.2f}, expected_tier={prediction.expected_model_tier}",
            fallback_model_id=fallback_id,
            use_verifier=prediction.verifier_required,
        )

    def _route_learned(self, prediction: TaskPrediction, verifier_fallback: bool = False) -> RoutingDecision:
        """Learned router with cost-quality tradeoff.
        
        In a full implementation, this would load a trained classifier.
        Here we use a heuristic calibrated from routing_stats.
        """
        # Check historical success rate per tier for this task type
        task_key = prediction.task_type.value
        best_tier = None
        best_score = -float("inf")
        
        for tier in self.TIER_ORDER:
            stats = self.routing_stats.get(f"{task_key}_tier_{tier}", {})
            success_rate = stats.get("success_rate", 0.5)
            avg_cost = stats.get("avg_cost", 0.01 * tier)
            
            # Score = success_weight * success_rate - cost_weight * cost
            score = 10 * success_rate - 100 * avg_cost
            
            # Penalize tiers below expected if risk is high
            if tier < prediction.expected_model_tier and prediction.risk_of_failure > 0.5:
                score -= 5
            
            if score > best_score:
                best_score = score
                best_tier = tier
        
        # Default to expected tier if no history
        if best_tier is None:
            best_tier = prediction.expected_model_tier
        
        models = self.models_by_tier.get(best_tier, self.models_by_tier[3])
        model = models[0] if models else list(self.config.models.values())[0]
        
        # Verifier fallback on uncertain predictions
        use_verifier = verifier_fallback and prediction.risk_of_failure > 0.5
        
        return RoutingDecision(
            model_id=model.model_id,
            provider=model.provider,
            tier=best_tier,
            confidence=min(best_score / 10 + 0.5, 1.0),
            reasoning=f"Learned router: tier {best_tier} scored {best_score:.3f} for {task_key}",
            fallback_model_id=self._next_tier_model(best_tier).model_id if best_tier < 5 else None,
            use_verifier=use_verifier,
        )

    def _route_cascade(self, prediction: TaskPrediction) -> RoutingDecision:
        """FrugalGPT-style cascade: try cheap first, escalate on low confidence."""
        start_tier = max(1, prediction.expected_model_tier - 2)
        
        # Don't start below tier 2 for risky tasks
        if prediction.risk_of_failure > 0.6:
            start_tier = max(start_tier, 2)
        
        models = self.models_by_tier.get(start_tier, [])
        if not models:
            models = self.models_by_tier.get(1, [])
        if not models:
            models = list(self.config.models.values())
        
        model = models[0]
        
        # Determine if we should pre-escalate (for critical tasks)
        pre_escalate = prediction.task_type == TaskType.LEGAL_REGULATED
        
        fallback = None
        if not pre_escalate and start_tier < prediction.expected_model_tier:
            fallback = self._next_tier_model(start_tier)
        
        return RoutingDecision(
            model_id=model.model_id,
            provider=model.provider,
            tier=start_tier,
            confidence=1.0 - prediction.risk_of_failure,
            reasoning=f"Cascade start at tier {start_tier}, expected tier {prediction.expected_model_tier}, risk={prediction.risk_of_failure:.2f}",
            fallback_model_id=fallback.model_id if fallback else None,
            use_verifier=prediction.verifier_required,
        )

    def _next_tier_model(self, current_tier: int) -> Optional[ModelConfig]:
        for tier in range(current_tier + 1, 6):
            models = self.models_by_tier.get(tier)
            if models:
                return models[0]
        return None

    def update_stats(self, task_type: TaskType, tier: int, cost: float, success: bool) -> None:
        key = f"{task_type.value}_tier_{tier}"
        stats = self.routing_stats.setdefault(key, {"count": 0, "successes": 0, "total_cost": 0.0})
        stats["count"] += 1
        if success:
            stats["successes"] += 1
        stats["total_cost"] += cost
        stats["success_rate"] = stats["successes"] / stats["count"]
        stats["avg_cost"] = stats["total_cost"] / stats["count"]

    def should_escalate(self, decision: RoutingDecision, step_outcome: Outcome, confidence: float) -> bool:
        """Decide whether to escalate to a stronger model after a step."""
        if decision.tier >= 5:
            return False
        if step_outcome == Outcome.FAILURE and confidence < 0.5:
            return True
        if step_outcome == Outcome.PARTIAL_SUCCESS and decision.tier < 4:
            return True
        return False