"""Budget-constrained tiered routing.

This module provides a router that selects optimal thresholds based on
a budget constraint. At runtime, given a cost or latency budget, it loads
the pre-computed Pareto frontier and selects the thresholds that maximize
quality within the budget.

Example:
    from embergraphrouter.inference.budget_router import BudgetAwareRouter

    # Load router with budget constraint
    router = BudgetAwareRouter.from_registry(
        "eb1-gnn-preview-12-19",
        cost_budget_per_1k=2.0,  # $2 per 1K queries
    )

    # Route queries - thresholds are automatically set for the budget
    result = router.route("What is 2+2?")
    print(f"Tier: {result.tier.name}, Model: {result.selected_model}")

Dynamic budget adjustment:
    # Change budget at runtime without reloading
    router.set_budget(cost_budget_per_1k=5.0)
    result = router.route("Prove Fermat's Last Theorem")
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    pass

from .tiered_router import TieredRouter, TieredRoutingResult, TieredThresholds

logger = logging.getLogger(__name__)


class BudgetAwareRouter(TieredRouter):
    """TieredRouter with budget-constrained threshold selection.

    This router extends TieredRouter with the ability to dynamically
    adjust thresholds based on cost or latency budgets. It loads a
    pre-computed Pareto frontier from the registry and selects optimal
    thresholds at runtime.
    """

    def __init__(
        self,
        gnn_model,
        llm_embeddings,
        llm_names: list[str],
        llm_metadata: dict[str, dict],
        pareto_frontier: dict,
        initial_thresholds: TieredThresholds | None = None,
        tier_configs=None,
        device: str = "cpu",
        embed_dim: int = 512,
        graph_cache: dict | None = None,
        inference_weights: tuple[float, float, float] = (0.7, 0.2, 0.1),
    ):
        """Initialize budget-aware router.

        Args:
            gnn_model: Trained GNN model.
            llm_embeddings: LLM description embeddings.
            llm_names: List of LLM names.
            llm_metadata: LLM cost/latency metadata.
            pareto_frontier: Pre-computed Pareto frontier data.
            initial_thresholds: Initial thresholds (optional).
            tier_configs: Tier configurations.
            device: Inference device.
            embed_dim: Embedding dimension.
            graph_cache: Training graph tensors for full GNN inference.
            inference_weights: Tuple of (perf, cost, latency) weights.
        """
        super().__init__(
            gnn_model=gnn_model,
            llm_embeddings=llm_embeddings,
            llm_names=llm_names,
            llm_metadata=llm_metadata,
            thresholds=initial_thresholds,
            tier_configs=tier_configs,
            device=device,
            embed_dim=embed_dim,
            graph_cache=graph_cache,
            inference_weights=inference_weights,
        )

        self._pareto_frontier = pareto_frontier
        self._current_budget = None

        logger.info(
            "BudgetAwareRouter initialized with %d Pareto points",
            len(pareto_frontier.get("points", [])),
        )

    @classmethod
    def from_registry(
        cls,
        registry_name: str,
        registry_root: str | Path | None = None,
        cost_budget_per_1k: float | None = None,
        latency_budget_ms: float | None = None,
        device: str = "cpu",
    ) -> "BudgetAwareRouter":
        """Load budget-aware router from registry.

        Args:
            registry_name: Name of registry entry.
            registry_root: Root directory for registry.
            cost_budget_per_1k: Cost budget per 1K queries (USD).
            latency_budget_ms: Latency budget per query (ms).
            device: Device for inference.

        Returns:
            Initialized BudgetAwareRouter.
        """
        import pickle

        import torch
        import yaml

        if registry_root is None:
            registry_root = Path(__file__).parent.parent / "registry"
        registry_root = Path(registry_root)
        registry_path = registry_root / registry_name

        if not registry_path.exists():
            raise FileNotFoundError(f"Registry not found: {registry_path}")

        # Load training graph cache (REQUIRED for full GNN inference)
        cache_path = registry_path / "training_graph_cache.pkl"
        if not cache_path.exists():
            raise FileNotFoundError(
                f"Training graph cache not found: {cache_path}. "
                "The full GNN forward pass requires the training graph "
                "for message passing context. Run "
                "package_inference_cache.py to create it."
            )
        with open(cache_path, "rb") as f:
            graph_cache = pickle.load(f)

        # Load config
        config_path = registry_path / "config.yaml"
        with open(config_path) as f:
            config_dict = yaml.safe_load(f)

        embed_dim = config_dict.get("embed_dim", 512)
        hidden_dim = config_dict.get("hidden_dim", 384)
        num_gat_layers = config_dict.get("num_gat_layers", 3)
        num_gat_heads = config_dict.get("num_gat_heads", 4)
        dropout = config_dict.get("dropout", 0.4)
        edge_feature_dim = config_dict.get("edge_feature_dim", 6)

        # Load LLM embeddings
        emb_path = registry_path / "llm_embeddings.pkl"
        with open(emb_path, "rb") as f:
            llm_embeddings_np = pickle.load(f)
        llm_embeddings = torch.from_numpy(llm_embeddings_np).float()

        # Load LLM descriptions
        desc_path = registry_path / "llm_descriptions.json"
        with open(desc_path) as f:
            llm_descriptions = json.load(f)

        # CRITICAL: Use llm_order from training graph as authoritative source.
        # This ensures GNN output indices map to correct model names.
        if "llm_order" not in graph_cache:
            raise ValueError(
                f"Training graph cache missing 'llm_order' key: {cache_path}. "
                "This is required to ensure GNN indices map to correct models."
            )
        llm_names = graph_cache["llm_order"]

        # Validate consistency with llm_descriptions.json
        json_names = set(llm_descriptions.keys())
        cache_names = set(llm_names)
        if json_names != cache_names:
            missing_in_json = cache_names - json_names
            missing_in_cache = json_names - cache_names
            raise ValueError(
                f"LLM name mismatch between training cache and descriptions. "
                f"In cache but not JSON: {missing_in_json}. "
                f"In JSON but not cache: {missing_in_cache}."
            )

        llm_metadata = {
            name: {
                "cost": desc.get("cost", 0.0),
                "latency": desc.get("latency", 0.0),
            }
            for name, desc in llm_descriptions.items()
        }

        # Load Pareto frontier (budget cache)
        budget_cache_path = registry_path / "budget_cache.json"
        if budget_cache_path.exists():
            with open(budget_cache_path) as f:
                pareto_frontier = json.load(f).get("pareto_frontier", {"points": []})
        else:
            logger.warning(
                "Budget cache not found at %s. Run calibration with --budget-optimize "
                "to generate it. Using default thresholds.",
                budget_cache_path,
            )
            pareto_frontier = {"points": []}

        # Initialize model
        try:
            from embergraphrouter.model.graph_nn import GATEncoderDecoderNet
        except ModuleNotFoundError:
            from model.graph_nn import GATEncoderDecoderNet

        # Load weights first to detect architecture from state_dict
        model_path = registry_path / "model.pth"
        state_dict = torch.load(model_path, map_location=device, weights_only=True)
        use_score_head = "score_head.weight" in state_dict

        model = GATEncoderDecoderNet(
            query_feature_dim=embed_dim,
            llm_feature_dim=embed_dim,
            hidden_features=hidden_dim,
            in_edges=edge_feature_dim,
            num_gnn_layers=num_gat_layers,
            num_attention_heads=num_gat_heads,
            dropout=dropout,
            hidden_multiplier=2,
            use_contrastive=True,
            use_score_head=use_score_head,
            use_difficulty_head=True,
            use_uncertainty_head=False,
        )

        model.load_state_dict(state_dict, strict=False)
        model.to(device)

        # Load inference weights from config
        iw = config_dict.get("inference_weights", {})
        inference_weights = (
            iw.get("perf", 0.7),
            iw.get("cost", 0.2),
            iw.get("latency", 0.1),
        )

        # Create router
        router = cls(
            gnn_model=model,
            llm_embeddings=llm_embeddings,
            llm_names=llm_names,
            llm_metadata=llm_metadata,
            pareto_frontier=pareto_frontier,
            device=device,
            embed_dim=embed_dim,
            graph_cache=graph_cache,
            inference_weights=inference_weights,
        )

        # Set initial budget if provided
        if cost_budget_per_1k is not None:
            router.set_budget(cost_budget_per_1k=cost_budget_per_1k)
        elif latency_budget_ms is not None:
            router.set_budget(latency_budget_ms=latency_budget_ms)

        return router

    def set_budget(
        self,
        cost_budget_per_1k: float | None = None,
        latency_budget_ms: float | None = None,
    ) -> None:
        """Set budget constraint and update thresholds.

        Args:
            cost_budget_per_1k: Maximum cost per 1K queries (USD).
            latency_budget_ms: Maximum latency per query (ms).

        Raises:
            ValueError: If no feasible configuration exists for the budget.
        """
        if not self._pareto_frontier.get("points"):
            logger.warning(
                "No Pareto frontier available. Using default thresholds."
            )
            return

        points = self._pareto_frontier["points"]

        # Filter to feasible points
        feasible = points
        if cost_budget_per_1k is not None:
            feasible = [p for p in feasible if p["cost_per_1k"] <= cost_budget_per_1k]
            self._current_budget = f"${cost_budget_per_1k}/1K"

        if latency_budget_ms is not None:
            feasible = [p for p in feasible if p["latency_ms"] <= latency_budget_ms]
            if self._current_budget:
                self._current_budget += f", {latency_budget_ms}ms"
            else:
                self._current_budget = f"{latency_budget_ms}ms"

        if not feasible:
            min_cost = min(p["cost_per_1k"] for p in points)
            raise ValueError(
                f"No feasible configuration for budget. "
                f"Minimum achievable cost is ${min_cost:.2f}/1K."
            )

        # Select highest quality feasible point
        best = max(feasible, key=lambda p: p["quality"])
        thresholds_dict = best["thresholds"]

        self.thresholds = TieredThresholds(
            trivial_mean_logit=thresholds_dict["trivial_mean_logit"],
            uncertain_mean_logit=thresholds_dict["uncertain_mean_logit"],
            uncertain_logit_margin=thresholds_dict["uncertain_logit_margin"],
        )

        logger.info(
            "Budget set to %s. Expected quality: %.1f%%, cost: $%.2f/1K",
            self._current_budget,
            best["quality"] * 100,
            best["cost_per_1k"],
        )

    def get_budget_info(self) -> dict:
        """Get current budget configuration info.

        Returns:
            Dict with current budget, thresholds, and expected metrics.
        """
        return {
            "budget": self._current_budget,
            "thresholds": {
                "trivial_mean_logit": self.thresholds.trivial_mean_logit,
                "uncertain_mean_logit": self.thresholds.uncertain_mean_logit,
                "uncertain_logit_margin": self.thresholds.uncertain_logit_margin,
            },
            "pareto_points": len(self._pareto_frontier.get("points", [])),
        }

    def route(self, query: str) -> TieredRoutingResult:
        """Route query using budget-optimized thresholds.

        Same as parent class but uses dynamically set thresholds.

        Args:
            query: Query text to route.

        Returns:
            TieredRoutingResult with routing decision.
        """
        return self.classify(query)


def create_budget_router(
    cost_budget_per_1k: float,
    registry_name: str = "eb1-gnn-preview-latest",
    device: str = "cpu",
) -> BudgetAwareRouter:
    """Convenience function to create a budget-aware router.

    Args:
        cost_budget_per_1k: Maximum cost per 1K queries (USD).
        registry_name: Name of registry entry.
        device: Inference device.

    Returns:
        BudgetAwareRouter configured for the given budget.
    """
    return BudgetAwareRouter.from_registry(
        registry_name,
        cost_budget_per_1k=cost_budget_per_1k,
        device=device,
    )
