"""Learned Router - Trainable cost-aware model router.

Trains a simple classifier on synthetic traces to predict optimal model tier.
Can be swapped into ModelCascadeRouter._route_learned().
"""

import json
import pickle
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from collections import defaultdict


@dataclass
class RouterFeatures:
    task_type_encoded: int
    request_length: int
    num_words: int
    contains_code_keywords: bool
    contains_legal_keywords: bool
    contains_research_keywords: bool
    contains_tool_keywords: bool
    expected_difficulty: int
    has_prior_failures: bool
    prior_success_rate: float


class LearnedRouter:
    """Simple learned router that can be trained on trace data."""

    def __init__(self):
        self.weights: Dict[str, float] = {}
        self.task_type_bias: Dict[str, float] = {}
        self.trained = False

    def _extract_features(self, user_request: str, task_type: str, history: List[Dict]) -> Dict[str, float]:
        """Extract features from request and history."""
        req_lower = user_request.lower()
        
        features = {
            "request_length": len(user_request),
            "num_words": len(user_request.split()),
            "has_code": any(kw in req_lower for kw in ["python", "javascript", "code", "function", "bug", "debug", "refactor", "implement", "test"]),
            "has_legal": any(kw in req_lower for kw in ["contract", "legal", "compliance", "gdpr", "privacy", "policy", "regulatory"]),
            "has_research": any(kw in req_lower for kw in ["research", "find sources", "literature", "investigate", "compare", "analyze"]),
            "has_tools": any(kw in req_lower for kw in ["search", "fetch", "retrieve", "query", "api", "database", "scrape"]),
            "has_long_horizon": any(kw in req_lower for kw in ["plan", "project", "roadmap", "orchestrate", "multi-step"]),
        }
        
        # Historical success rate for this task type
        similar = [t for t in history if t.get("task_type") == task_type]
        if similar:
            success_count = sum(1 for t in similar if t.get("final_outcome") == "success")
            features["prior_success_rate"] = success_count / len(similar)
            features["has_prior_failures"] = any(t.get("final_outcome") == "failure" for t in similar[-5:])
        else:
            features["prior_success_rate"] = 0.5
            features["has_prior_failures"] = False
        
        return features

    def _score_tier(self, features: Dict[str, float], tier: int) -> float:
        """Score a tier given features. Higher is better."""
        if not self.trained:
            # Heuristic scoring before training
            base_score = {1: 0.3, 2: 0.5, 3: 0.7, 4: 0.85, 5: 0.9}.get(tier, 0.5)
            
            # Adjust by task complexity signals
            if features["has_legal"] and tier < 4:
                base_score -= 0.4
            if features["has_code"] and tier < 3:
                base_score -= 0.2
            if features["has_research"] and tier < 3:
                base_score -= 0.15
            if features["has_long_horizon"] and tier < 3:
                base_score -= 0.2
            
            # Adjust by history
            if features["has_prior_failures"] and tier < 3:
                base_score -= 0.3
            if features["prior_success_rate"] > 0.8 and tier > 2:
                base_score += 0.1
            
            return base_score
        
        # Trained scoring
        score = self.task_type_bias.get(str(tier), 0.5)
        for feat_name, feat_val in features.items():
            weight_key = f"{feat_name}_tier_{tier}"
            score += self.weights.get(weight_key, 0.0) * (1.0 if feat_val else 0.0)
        return score

    def predict_tier(self, user_request: str, task_type: str, history: List[Dict] = None) -> Tuple[int, float]:
        """Predict optimal model tier and confidence."""
        history = history or []
        features = self._extract_features(user_request, task_type, history)
        
        best_tier = 3
        best_score = -float("inf")
        
        for tier in [1, 2, 3, 4, 5]:
            score = self._score_tier(features, tier)
            if score > best_score:
                best_score = score
                best_tier = tier
        
        confidence = min(best_score, 1.0)
        return best_tier, confidence

    def train(self, traces: List[Dict]) -> None:
        """Train router on trace data using simple frequency-based weight learning."""
        # For each trace, determine optimal tier from outcome
        tier_counts = defaultdict(lambda: defaultdict(int))
        
        for trace in traces:
            task_type = trace.get("task_type", "unknown")
            outcome = trace.get("final_outcome", "failure")
            difficulty = trace.get("metadata", {}).get("difficulty", 3)
            actual_tier = trace.get("metadata", {}).get("actual_tier", 3)
            
            # Optimal tier is the minimum tier that would succeed
            if outcome == "success":
                optimal = actual_tier  # This tier succeeded
            else:
                optimal = min(actual_tier + 1, 5)  # Need higher tier
            
            # Extract features
            req = trace.get("user_request", "")
            features = self._extract_features(req, task_type, [])
            
            # Count successes per feature+tier combination
            for feat_name, feat_val in features.items():
                if isinstance(feat_val, bool) and feat_val:
                    tier_counts[feat_name][optimal] += 1
            
            tier_counts["_overall"][optimal] += 1
        
        # Compute bias per tier
        total = sum(tier_counts["_overall"].values())
        if total > 0:
            for tier in [1, 2, 3, 4, 5]:
                count = tier_counts["_overall"].get(tier, 0)
                self.task_type_bias[str(tier)] = count / total
        
        # Compute feature weights
        for feat_name, tier_dist in tier_counts.items():
            if feat_name == "_overall":
                continue
            total_feat = sum(tier_dist.values())
            if total_feat > 0:
                for tier in [1, 2, 3, 4, 5]:
                    weight_key = f"{feat_name}_tier_{tier}"
                    # Positive if this tier is common when feature is present
                    self.weights[weight_key] = (tier_dist.get(tier, 0) / total_feat) - self.task_type_bias.get(str(tier), 0.1)
        
        self.trained = True
        print(f"Trained router on {len(traces)} traces. Weights: {len(self.weights)}.")

    def save(self, path: str) -> None:
        with open(path, "wb") as f:
            pickle.dump({"weights": self.weights, "bias": self.task_type_bias, "trained": self.trained}, f)

    def load(self, path: str) -> None:
        with open(path, "rb") as f:
            data = pickle.load(f)
            self.weights = data["weights"]
            self.task_type_bias = data["bias"]
            self.trained = data["trained"]


class RouterDatasetBuilder:
    """Builds training datasets from traces for learned router training."""

    @staticmethod
    def from_traces(traces: List[Dict]) -> List[Dict]:
        """Convert traces to (features, optimal_tier) training examples."""
        examples = []
        for trace in traces:
            difficulty = trace.get("metadata", {}).get("difficulty", 3)
            actual_tier = trace.get("metadata", {}).get("actual_tier", 3)
            outcome = trace.get("final_outcome", "failure")
            
            # Optimal tier
            if outcome == "success":
                optimal = actual_tier
            else:
                optimal = min(actual_tier + 1, 5)
            
            # Simple feature extraction
            req = trace.get("user_request", "").lower()
            features = {
                "length": len(req),
                "has_code": any(kw in req for kw in ["python", "code", "function", "bug", "debug"]),
                "has_legal": any(kw in req for kw in ["contract", "legal", "compliance", "gdpr"]),
                "has_research": any(kw in req for kw in ["research", "find sources", "literature"]),
                "task_type": trace.get("task_type", "unknown"),
                "difficulty": difficulty,
            }
            
            examples.append({"features": features, "optimal_tier": optimal, "outcome": outcome})
        
        return examples

    @staticmethod
    def compute_oracle_savings(traces: List[Dict]) -> Dict[str, float]:
        """Compute what an oracle router (perfect tier selection) would save."""
        total_cost = 0.0
        oracle_cost = 0.0
        
        for trace in traces:
            cost = trace.get("total_cost", 0.0)
            difficulty = trace.get("metadata", {}).get("difficulty", 3)
            
            # Oracle uses cheapest tier that can handle difficulty
            oracle_tier = max(1, difficulty - 1)  # Slightly optimistic
            oracle_mult = {1: 0.05, 2: 0.15, 3: 0.75, 4: 1.0, 5: 1.5}.get(oracle_tier, 0.75)
            actual_mult = {1: 0.05, 2: 0.15, 3: 0.75, 4: 1.0, 5: 1.5}.get(
                trace.get("metadata", {}).get("actual_tier", 3), 0.75)
            
            oracle_cost += cost * (oracle_mult / actual_mult) if actual_mult > 0 else cost
            total_cost += cost
        
        return {
            "total_cost": total_cost,
            "oracle_cost": oracle_cost,
            "savings": (total_cost - oracle_cost) / max(total_cost, 1),
            "per_trace_savings": (total_cost - oracle_cost) / len(traces) if traces else 0,
        }