"""Production inference pipeline for EmberGraphRouter.

This module provides a streamlined interface for model inference, including
zero-cost uncertainty and difficulty extraction from routing logits.

Usage:
    from embergraphrouter.inference import EmberGraphRouterInference

    # Load a trained model
    router = EmberGraphRouterInference.load(
        checkpoint_path="model_path/v7_gat4_optimal_best.pth",
        config_path="configs/v7_gat4_optimal.yaml",
    )

    # Route a query
    result = router.route(query_embedding, task_embedding)
    print(f"Best LLM: {result.chosen_llm} (confidence={result.confidence:.2f})")

    # Batch routing
    results = router.route_batch(query_embeddings, task_embeddings)
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Sequence

import torch
import torch.nn.functional as F
import yaml

from .data_processing.utils import loadjson, loadpkl
from .model.logit_signals import compute_routing_signals

if TYPE_CHECKING:
    from numpy.typing import NDArray

logger = logging.getLogger(__name__)


@dataclass
class RoutingResult:
    """Result of routing a single query through EmberGraphRouter.

    Attributes:
        chosen_llm: Name of the selected LLM.
        chosen_idx: Index of the selected LLM.
        confidence: Max softmax probability (how confident the router is).
        uncertainty: Normalized entropy of the distribution (0=certain, 1=uniform).
        difficulty: 1 - confidence (intrinsic difficulty of the query).
        all_scores: Dictionary mapping LLM names to raw logit scores.
        all_probs: Dictionary mapping LLM names to softmax probabilities.
    """

    chosen_llm: str
    chosen_idx: int
    confidence: float
    uncertainty: float
    difficulty: float
    all_scores: dict[str, float]
    all_probs: dict[str, float]


@dataclass
class BatchRoutingResult:
    """Result of routing a batch of queries.

    Attributes:
        chosen_llms: List of selected LLM names.
        chosen_indices: Tensor of selected LLM indices.
        confidences: Tensor of confidence scores.
        uncertainties: Tensor of uncertainty scores.
        difficulties: Tensor of difficulty scores.
        all_scores: Tensor of shape (batch_size, num_llms) with raw logits.
        all_probs: Tensor of shape (batch_size, num_llms) with softmax probs.
    """

    chosen_llms: list[str]
    chosen_indices: torch.Tensor
    confidences: torch.Tensor
    uncertainties: torch.Tensor
    difficulties: torch.Tensor
    all_scores: torch.Tensor
    all_probs: torch.Tensor


class EmberGraphRouterInference:
    """Production inference interface for EmberGraphRouter.

    This class provides a clean API for routing queries to LLMs using a
    trained EmberGraphRouter model. It includes zero-cost extraction of
    uncertainty and difficulty signals from the routing logits.

    Key features:
    - Single-query and batch routing
    - Zero-cost uncertainty/difficulty extraction
    - Temperature-calibrated probabilities
    - Thread-safe inference (no training state)
    """

    def __init__(
        self,
        model: torch.nn.Module,
        llm_names: Sequence[str],
        llm_embeddings: torch.Tensor,
        temperature: float = 1.0,
        device: str = "cpu",
    ) -> None:
        """Initialize the inference pipeline.

        Args:
            model: Trained GNN model (EncoderDecoderNet).
            llm_names: Ordered list of LLM names matching model output indices.
            llm_embeddings: LLM description embeddings (num_llms, embed_dim).
            temperature: Temperature for softmax calibration (from training).
            device: Device for inference ("cpu" or "cuda").
        """
        self.model = model.to(device)
        self.model.eval()
        self.llm_names = list(llm_names)
        self.llm_embeddings = llm_embeddings.to(device)
        self.temperature = temperature
        self.device = device
        self.num_llms = len(llm_names)

        # Pre-compute edge indices for single query (reused for efficiency)
        self._single_query_edge_index = torch.tensor(
            [[0] * self.num_llms, list(range(self.num_llms))],
            dtype=torch.long,
            device=device,
        )

    @classmethod
    def load(
        cls,
        checkpoint_path: str,
        config_path: str,
        device: str | None = None,
    ) -> "EmberGraphRouterInference":
        """Load a trained model from checkpoint.

        Args:
            checkpoint_path: Path to model checkpoint (.pth file).
            config_path: Path to config YAML file.
            device: Device for inference (defaults to cuda if available).

        Returns:
            Initialized EmberGraphRouterInference instance.
        """
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"

        checkpoint_path = Path(checkpoint_path)
        config_path = Path(config_path)

        # Load config
        with open(config_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)

        # Resolve relative paths
        base_dir = config_path.parent
        for key in ("llm_description_path", "llm_embedding_path"):
            if key in config and not Path(config[key]).is_absolute():
                config[key] = str(base_dir / config[key])

        # Load LLM metadata
        llm_descriptions = loadjson(config["llm_description_path"])
        llm_names = list(llm_descriptions.keys())
        llm_embeddings = torch.tensor(loadpkl(config["llm_embedding_path"]), device=device)

        # Reconstruct model architecture
        from .model.graph_nn import EncoderDecoderNet

        query_dim = 1536  # Default for sentence-transformers/all-mpnet-base-v2
        llm_dim = llm_embeddings.shape[1]

        model = EncoderDecoderNet(
            query_feature_dim=query_dim,
            llm_feature_dim=llm_dim,
            channels=config.get("embedding_dim", 256),
            in_edges=config.get("edge_dim", 6),
            num_llm=len(llm_names),
            num_gnn_layers=config.get("num_gnn_layers", 4),
            num_attention_heads=config.get("num_attention_heads", 4),
            attention_dropout=config.get("attention_dropout", 0.15),
            dropout=config.get("dropout", 0.25),
            query_mlp_layers=config.get("query_mlp_layers", 1),
            llm_mlp_layers=config.get("llm_mlp_layers", 1),
            use_residual=config.get("use_residual", True),
            use_swiglu=config.get("use_swiglu", False),
            use_variational=config.get("use_variational", False),
            use_contrastive=config.get("use_contrastive", True),
            contrastive_temperature=config.get("contrastive_temperature", 0.1),
            contrastive_weight=config.get("contrastive_weight", 0.1),
            use_difficulty_head=config.get("use_difficulty_head", False),
            use_gated_skip=config.get("use_gated_skip", False),
        )

        # Load checkpoint weights
        state_dict = torch.load(checkpoint_path, map_location=device, weights_only=True)
        model.load_state_dict(state_dict)
        model.eval()

        logger.info(
            "Loaded EmberGraphRouter from %s (%d LLMs, %d params)",
            checkpoint_path,
            len(llm_names),
            sum(p.numel() for p in model.parameters()),
        )

        # Load learned temperature if available
        temperature = 1.0
        temp_path = checkpoint_path.with_suffix(".temperature")
        if temp_path.exists():
            temperature = float(temp_path.read_text().strip())
            logger.info("Loaded calibrated temperature: %.3f", temperature)

        return cls(
            model=model,
            llm_names=llm_names,
            llm_embeddings=llm_embeddings,
            temperature=temperature,
            device=device,
        )

    def _prepare_edge_features(
        self,
        num_queries: int,
        preference_weights: tuple[float, float, float] = (1.0, 0.0, 0.0),
    ) -> torch.Tensor:
        """Prepare edge features with preference encoding.

        Args:
            num_queries: Number of queries in the batch.
            preference_weights: (perf, cost, latency) weights.

        Returns:
            Edge features tensor of shape (num_queries * num_llms, edge_dim).
        """
        wp, wc, wl = preference_weights
        total_edges = num_queries * self.num_llms

        # Default edge features: [cost, perf, latency, wp, wc, wl]
        # Since we don't have actual cost/perf/latency during inference,
        # we use placeholder values and rely on the learned embeddings
        edge_features = torch.zeros(total_edges, 6, device=self.device)
        edge_features[:, 3] = wp  # perf weight
        edge_features[:, 4] = wc  # cost weight
        edge_features[:, 5] = wl  # latency weight

        return edge_features

    def route(
        self,
        query_embedding: torch.Tensor | "NDArray",
        task_embedding: torch.Tensor | "NDArray",
        preference_weights: tuple[float, float, float] = (1.0, 0.0, 0.0),
    ) -> RoutingResult:
        """Route a single query to the best LLM.

        Args:
            query_embedding: Query embedding vector (embed_dim,).
            task_embedding: Task description embedding vector (embed_dim,).
            preference_weights: (perf, cost, latency) weights for scoring.

        Returns:
            RoutingResult with selected LLM and diagnostic signals.
        """
        # Convert to tensors
        if not isinstance(query_embedding, torch.Tensor):
            query_embedding = torch.tensor(query_embedding, dtype=torch.float32)
        if not isinstance(task_embedding, torch.Tensor):
            task_embedding = torch.tensor(task_embedding, dtype=torch.float32)

        query_embedding = query_embedding.to(self.device).unsqueeze(0)  # (1, embed_dim)
        task_embedding = task_embedding.to(self.device).unsqueeze(0)  # (1, embed_dim)

        # Prepare inputs
        edge_features = self._prepare_edge_features(1, preference_weights)
        edge_mask = torch.ones(self.num_llms, dtype=torch.bool, device=self.device)
        edge_can_see = torch.zeros(self.num_llms, dtype=torch.bool, device=self.device)

        with torch.no_grad():
            logits = self.model(
                task_features=task_embedding,
                query_features=query_embedding,
                llm_features=self.llm_embeddings,
                edge_index=self._single_query_edge_index,
                edge_mask=edge_mask,
                edge_can_see=edge_can_see,
                edge_weight=edge_features,
            )

        # Reshape to (1, num_llms) for signal extraction
        logits_grouped = logits.reshape(1, self.num_llms)
        signals = compute_routing_signals(logits_grouped, temperature=self.temperature)

        # Compute probabilities
        probs = F.softmax(logits_grouped / self.temperature, dim=-1).squeeze(0)

        chosen_idx = int(signals["chosen_idx"].item())

        return RoutingResult(
            chosen_llm=self.llm_names[chosen_idx],
            chosen_idx=chosen_idx,
            confidence=float(signals["confidence"].item()),
            uncertainty=float(signals["uncertainty"].item()),
            difficulty=float(signals["difficulty"].item()),
            all_scores={name: float(logits[i].item()) for i, name in enumerate(self.llm_names)},
            all_probs={name: float(probs[i].item()) for i, name in enumerate(self.llm_names)},
        )

    def route_batch(
        self,
        query_embeddings: torch.Tensor | "NDArray",
        task_embeddings: torch.Tensor | "NDArray",
        preference_weights: tuple[float, float, float] = (1.0, 0.0, 0.0),
    ) -> BatchRoutingResult:
        """Route a batch of queries to the best LLMs.

        Args:
            query_embeddings: Query embeddings (batch_size, embed_dim).
            task_embeddings: Task embeddings (batch_size, embed_dim).
            preference_weights: (perf, cost, latency) weights for scoring.

        Returns:
            BatchRoutingResult with selected LLMs and diagnostic signals.
        """
        # Convert to tensors
        if not isinstance(query_embeddings, torch.Tensor):
            query_embeddings = torch.tensor(query_embeddings, dtype=torch.float32)
        if not isinstance(task_embeddings, torch.Tensor):
            task_embeddings = torch.tensor(task_embeddings, dtype=torch.float32)

        query_embeddings = query_embeddings.to(self.device)
        task_embeddings = task_embeddings.to(self.device)

        batch_size = query_embeddings.shape[0]
        total_edges = batch_size * self.num_llms

        # Build batch edge index
        source_nodes = torch.arange(batch_size, device=self.device).repeat_interleave(self.num_llms)
        dest_nodes = torch.arange(self.num_llms, device=self.device).repeat(batch_size)
        edge_index = torch.stack([source_nodes, dest_nodes])

        # Prepare edge features
        edge_features = self._prepare_edge_features(batch_size, preference_weights)
        edge_mask = torch.ones(total_edges, dtype=torch.bool, device=self.device)
        edge_can_see = torch.zeros(total_edges, dtype=torch.bool, device=self.device)

        with torch.no_grad():
            logits = self.model(
                task_features=task_embeddings,
                query_features=query_embeddings,
                llm_features=self.llm_embeddings,
                edge_index=edge_index,
                edge_mask=edge_mask,
                edge_can_see=edge_can_see,
                edge_weight=edge_features,
            )

        # Reshape to (batch_size, num_llms) for signal extraction
        logits_grouped = logits.reshape(batch_size, self.num_llms)
        signals = compute_routing_signals(logits_grouped, temperature=self.temperature)

        # Compute probabilities
        probs = F.softmax(logits_grouped / self.temperature, dim=-1)

        chosen_indices = signals["chosen_idx"]
        chosen_llms = [self.llm_names[int(idx.item())] for idx in chosen_indices]

        return BatchRoutingResult(
            chosen_llms=chosen_llms,
            chosen_indices=chosen_indices,
            confidences=signals["confidence"],
            uncertainties=signals["uncertainty"],
            difficulties=signals["difficulty"],
            all_scores=logits_grouped,
            all_probs=probs,
        )

    def get_routing_explanation(self, result: RoutingResult) -> str:
        """Generate a human-readable explanation of the routing decision.

        Args:
            result: RoutingResult from route() method.

        Returns:
            Formatted explanation string.
        """
        lines = [
            f"Selected LLM: {result.chosen_llm}",
            f"Confidence: {result.confidence:.1%}",
            f"Query difficulty: {result.difficulty:.2f} ({'easy' if result.difficulty < 0.3 else 'medium' if result.difficulty < 0.6 else 'hard'})",
            "",
            "Score breakdown:",
        ]

        # Sort LLMs by probability
        sorted_llms = sorted(
            result.all_probs.items(),
            key=lambda x: x[1],
            reverse=True,
        )

        for llm_name, prob in sorted_llms:
            marker = "→ " if llm_name == result.chosen_llm else "  "
            lines.append(f"  {marker}{llm_name}: {prob:.1%}")

        if result.uncertainty > 0.8:
            lines.append("")
            lines.append("Note: High uncertainty - consider using top-k sampling or ensemble.")

        return "\n".join(lines)
