"""Lightweight router inference with cached artifacts.

This module provides a fast inference path that pre-caches all static artifacts:
- Model weights
- LLM embeddings
- Edge index template (static graph structure)
- LLM metadata (cost, latency)

Only query embeddings need to be computed per-request via OpenAI API.

Example:
    router = RouterInference.from_registry("cost_optimization_router_v1")

    # Route a query
    result = router.route(query_embedding)
    print(f"Best model: {result.best_model}")
    print(f"Confidence: {result.confidence:.2f}")
    print(f"All scores: {result.scores}")
"""

from __future__ import annotations

import json
import pickle
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import torch
import torch.nn.functional as F

if TYPE_CHECKING:
    from numpy.typing import NDArray


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

    best_model: str
    best_model_idx: int
    confidence: float  # Max probability
    entropy: float  # Normalized prediction entropy [0, 1]
    scores: dict[str, float]  # Raw logits per model
    probabilities: dict[str, float]  # Softmax probabilities per model


@dataclass
class CachedArtifacts:
    """Pre-cached inference artifacts."""

    model_state_dict: dict
    llm_embeddings: torch.Tensor  # (num_llms, embed_dim)
    llm_names: list[str]
    llm_metadata: dict[str, dict]  # cost, latency per LLM
    edge_index: torch.Tensor  # Static edge template
    config: dict


class RouterInference:
    """Fast inference router with cached artifacts.

    This class provides a lightweight inference path that avoids loading
    the full training dataset. All static artifacts are pre-cached.
    """

    def __init__(
        self,
        artifacts: CachedArtifacts,
        device: str = "cpu",
    ) -> None:
        """Initialize router with cached artifacts.

        Args:
            artifacts: Pre-cached inference artifacts.
            device: Device to run inference on ("cpu" or "cuda").
        """
        self.device = device
        self.artifacts = artifacts
        self.llm_names = artifacts.llm_names
        self.num_llms = len(self.llm_names)

        # Move embeddings to device
        self.llm_embeddings = artifacts.llm_embeddings.to(device)
        self.edge_index = artifacts.edge_index.to(device)

        # Initialize model architecture
        self._init_model(artifacts.config)

        # Load weights
        self.model.load_state_dict(artifacts.model_state_dict)
        self.model.eval()
        self.model.to(device)

    def _init_model(self, config: dict) -> None:
        """Initialize model architecture from config."""
        # Import here to avoid circular dependencies
        from ..model.graph_nn import HeteroGATRouter

        embed_dim = config.get("embed_dim", 512)
        hidden_dim = config.get("hidden_dim", 384)
        num_layers = config.get("num_gat_layers", 3)
        num_heads = config.get("num_gat_heads", 4)
        dropout = config.get("dropout", 0.4)

        self.model = HeteroGATRouter(
            query_dim=embed_dim,
            llm_dim=embed_dim,
            task_dim=embed_dim,  # Task dim same as query for single-query inference
            hidden_dim=hidden_dim,
            num_layers=num_layers,
            num_heads=num_heads,
            dropout=dropout,
            edge_dim=config.get("edge_feature_dim", 6),
            use_difficulty_head=False,  # Not needed for inference
            use_uncertainty_head=config.get("use_uncertainty_head", False),
        )

    @classmethod
    def from_registry(
        cls,
        registry_name: str,
        registry_root: str | Path | None = None,
        device: str = "cpu",
    ) -> RouterInference:
        """Load router from registry.

        Args:
            registry_name: Name of the registry entry (e.g., "cost_optimization_router_v1").
            registry_root: Root directory of registry. Defaults to embergraphrouter/registry/.
            device: Device to run inference on.

        Returns:
            Initialized RouterInference instance.
        """
        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 entry not found: {registry_path}")

        artifacts = cls._load_artifacts(registry_path)
        return cls(artifacts, device=device)

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint_path: str | Path,
        llm_embedding_path: str | Path,
        llm_descriptions_path: str | Path,
        config: dict | None = None,
        device: str = "cpu",
    ) -> RouterInference:
        """Load router from individual checkpoint files.

        Args:
            checkpoint_path: Path to model.pth weights.
            llm_embedding_path: Path to LLM embeddings pickle.
            llm_descriptions_path: Path to LLM descriptions JSON.
            config: Model configuration dict. If None, uses defaults.
            device: Device to run inference on.

        Returns:
            Initialized RouterInference instance.
        """
        checkpoint_path = Path(checkpoint_path)
        llm_embedding_path = Path(llm_embedding_path)
        llm_descriptions_path = Path(llm_descriptions_path)

        # Load model weights
        state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)

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

        # Load LLM descriptions
        with open(llm_descriptions_path, "r") as f:
            llm_descriptions = json.load(f)

        llm_names = list(llm_descriptions.keys())
        num_llms = len(llm_names)

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

        # Build static edge index template
        # For a single query: edges from query (node 0) to each LLM (nodes 1..num_llms)
        edge_index = cls._build_edge_template(num_llms)

        # Default config
        if config is None:
            config = {
                "embed_dim": llm_embeddings.shape[1],
                "hidden_dim": 384,
                "num_gat_layers": 3,
                "num_gat_heads": 4,
                "dropout": 0.4,
                "edge_feature_dim": 6,
            }

        artifacts = CachedArtifacts(
            model_state_dict=state_dict,
            llm_embeddings=llm_embeddings,
            llm_names=llm_names,
            llm_metadata=llm_metadata,
            edge_index=edge_index,
            config=config,
        )

        return cls(artifacts, device=device)

    @staticmethod
    def _build_edge_template(num_llms: int) -> torch.Tensor:
        """Build static edge index for single-query inference.

        Graph structure:
        - Node 0: Query node
        - Nodes 1..num_llms: LLM nodes
        - Edges: Bidirectional between query and each LLM

        Args:
            num_llms: Number of LLM nodes.

        Returns:
            Edge index tensor of shape (2, num_edges).
        """
        # Query -> LLM edges
        src = torch.zeros(num_llms, dtype=torch.long)
        dst = torch.arange(1, num_llms + 1, dtype=torch.long)

        # Bidirectional
        edge_index = torch.stack([
            torch.cat([src, dst]),
            torch.cat([dst, src]),
        ])

        return edge_index

    @classmethod
    def _load_artifacts(cls, registry_path: Path) -> CachedArtifacts:
        """Load artifacts from registry directory."""
        # Load model weights
        model_path = registry_path / "model.pth"
        state_dict = torch.load(model_path, map_location="cpu", weights_only=True)

        # 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, "r") as f:
            llm_descriptions = json.load(f)

        llm_names = list(llm_descriptions.keys())
        num_llms = len(llm_names)

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

        # Load config
        config_path = registry_path / "config.yaml"
        if config_path.exists():
            import yaml
            with open(config_path, "r") as f:
                config = yaml.safe_load(f)
        else:
            config = {
                "embed_dim": llm_embeddings.shape[1],
                "hidden_dim": 384,
                "num_gat_layers": 3,
                "num_gat_heads": 4,
                "dropout": 0.4,
                "edge_feature_dim": 6,
            }

        # Build edge template
        edge_index = cls._build_edge_template(num_llms)

        return CachedArtifacts(
            model_state_dict=state_dict,
            llm_embeddings=llm_embeddings,
            llm_names=llm_names,
            llm_metadata=llm_metadata,
            edge_index=edge_index,
            config=config,
        )

    def route(
        self,
        query_embedding: NDArray[np.float32] | torch.Tensor,
        task_embedding: NDArray[np.float32] | torch.Tensor | None = None,
    ) -> RoutingResult:
        """Route a query to the best LLM.

        Args:
            query_embedding: Query embedding vector (embed_dim,).
            task_embedding: Optional task embedding. If None, uses query embedding.

        Returns:
            RoutingResult with best model, confidence, and all scores.
        """
        # Convert to tensor
        if isinstance(query_embedding, np.ndarray):
            query_embedding = torch.from_numpy(query_embedding).float()
        query_embedding = query_embedding.to(self.device)

        if task_embedding is None:
            task_embedding = query_embedding
        elif isinstance(task_embedding, np.ndarray):
            task_embedding = torch.from_numpy(task_embedding).float()
        task_embedding = task_embedding.to(self.device)

        # Ensure correct shape
        if query_embedding.dim() == 1:
            query_embedding = query_embedding.unsqueeze(0)
        if task_embedding.dim() == 1:
            task_embedding = task_embedding.unsqueeze(0)

        # Run inference
        with torch.no_grad():
            logits = self._forward(query_embedding, task_embedding)

        # Convert to probabilities
        probs = F.softmax(logits, dim=-1).squeeze(0).cpu().numpy()
        logits_np = logits.squeeze(0).cpu().numpy()

        # Compute entropy
        entropy = -np.sum(probs * np.log(probs + 1e-8))
        max_entropy = np.log(self.num_llms)
        normalized_entropy = entropy / max_entropy

        # Find best model
        best_idx = int(np.argmax(probs))
        best_model = self.llm_names[best_idx]
        confidence = float(probs[best_idx])

        # Build result dicts
        scores = {name: float(logits_np[i]) for i, name in enumerate(self.llm_names)}
        probabilities = {name: float(probs[i]) for i, name in enumerate(self.llm_names)}

        return RoutingResult(
            best_model=best_model,
            best_model_idx=best_idx,
            confidence=confidence,
            entropy=normalized_entropy,
            scores=scores,
            probabilities=probabilities,
        )

    def route_with_cost_optimization(
        self,
        query_embedding: NDArray[np.float32] | torch.Tensor,
        entropy_threshold: float = 0.6,
        cheap_model: str | None = None,
        task_embedding: NDArray[np.float32] | torch.Tensor | None = None,
    ) -> RoutingResult:
        """Route with cost optimization based on prediction entropy.

        When prediction entropy is high (routing doesn't matter), use cheap model.
        When entropy is low (router found a winner), trust the prediction.

        Args:
            query_embedding: Query embedding vector.
            entropy_threshold: Entropy threshold for cost optimization.
                Values above this use cheap_model.
            cheap_model: Name of cheap fallback model. Defaults to cheapest in registry.
            task_embedding: Optional task embedding.

        Returns:
            RoutingResult with potentially overridden best_model.
        """
        result = self.route(query_embedding, task_embedding)

        if result.entropy > entropy_threshold:
            # High entropy: use cheap model
            if cheap_model is None:
                # Find cheapest model
                cheap_model = min(
                    self.artifacts.llm_metadata.items(),
                    key=lambda x: x[1].get("cost", float("inf")),
                )[0]

            if cheap_model in self.llm_names:
                result = RoutingResult(
                    best_model=cheap_model,
                    best_model_idx=self.llm_names.index(cheap_model),
                    confidence=result.probabilities.get(cheap_model, 0.0),
                    entropy=result.entropy,
                    scores=result.scores,
                    probabilities=result.probabilities,
                )

        return result

    def _forward(
        self,
        query_embedding: torch.Tensor,
        task_embedding: torch.Tensor,
    ) -> torch.Tensor:
        """Run model forward pass.

        Note: This is a simplified forward pass for single-query inference.
        The full heterogeneous graph structure is not needed here.

        Args:
            query_embedding: Query embedding (1, embed_dim).
            task_embedding: Task embedding (1, embed_dim).

        Returns:
            Logits tensor (1, num_llms).
        """
        # Build minimal graph for single query
        # Node features: [query, llm1, llm2, ..., llmN]
        node_features = torch.cat([
            query_embedding,
            self.llm_embeddings,
        ], dim=0)

        # Create edge mask (all visible for inference)
        num_edges = self.edge_index.shape[1]
        edge_mask = torch.ones(num_edges, dtype=torch.bool, device=self.device)

        # Create edge features (simplified - just preference weights)
        # [perf, cost, latency, pref_perf, pref_cost, pref_latency]
        edge_features = torch.zeros(num_edges, 6, device=self.device)
        edge_features[:, 0] = 1.0  # perf weight
        edge_features[:, 3] = 0.7  # pref_perf
        edge_features[:, 4] = 0.2  # pref_cost
        edge_features[:, 5] = 0.1  # pref_latency

        # Forward through model
        # Note: The model expects task_features, query_features, llm_features separately
        # but for inference we can simplify by passing the combined graph
        logits = self.model(
            task_features=task_embedding,
            query_features=query_embedding,
            llm_features=self.llm_embeddings.unsqueeze(0) if self.llm_embeddings.dim() == 2 else self.llm_embeddings,
            edge_index=self.edge_index,
            edge_mask=edge_mask,
            edge_can_see=edge_mask,
            edge_weight=edge_features,
        )

        return logits


def create_inference_cache(
    checkpoint_path: str | Path,
    llm_embedding_path: str | Path,
    llm_descriptions_path: str | Path,
    output_dir: str | Path,
    config: dict | None = None,
) -> Path:
    """Create a cached inference bundle for deployment.

    This bundles all artifacts needed for fast inference into a single directory.

    Args:
        checkpoint_path: Path to model weights.
        llm_embedding_path: Path to LLM embeddings.
        llm_descriptions_path: Path to LLM descriptions.
        output_dir: Output directory for cached bundle.
        config: Model config to include.

    Returns:
        Path to the created cache directory.
    """
    import shutil
    import yaml

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Copy artifacts
    shutil.copy(checkpoint_path, output_dir / "model.pth")
    shutil.copy(llm_embedding_path, output_dir / "llm_embeddings.pkl")
    shutil.copy(llm_descriptions_path, output_dir / "llm_descriptions.json")

    # Save config
    if config is not None:
        with open(output_dir / "config.yaml", "w") as f:
            yaml.dump(config, f)

    # Create metadata
    metadata = {
        "version": "1.0",
        "created_by": "create_inference_cache",
        "artifacts": [
            "model.pth",
            "llm_embeddings.pkl",
            "llm_descriptions.json",
            "config.yaml",
        ],
    }
    with open(output_dir / "METADATA.yaml", "w") as f:
        yaml.dump(metadata, f)

    return output_dir
