"""Graph-aware EB1 Router inference with full heterogeneous graph context.

This module provides inference that leverages the full training graph for
GNN message passing, matching the evaluation protocol used during training.

During training evaluation, the GNN uses:
- All query nodes (train + valid + test)
- All LLM nodes (7 models)
- edge_can_see = train_mask | valid_mask (for message passing)
- edge_mask = test edges only (for prediction)

For production inference, we:
1. Load the cached training graph (15,745 queries, 7 LLMs)
2. Append the new query as an additional node
3. Run GNN with train+valid edges visible
4. Predict on edges from the new query to all LLMs

This approach achieves the 68% frac_oracle performance reported in training.
"""

from __future__ import annotations

import hashlib
import pickle
import time
from dataclasses import dataclass
from functools import lru_cache
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 GraphRoutingResult:
    """Result of graph-aware routing decision."""

    best_model: str
    best_model_idx: int
    confidence: float
    entropy: float
    scores: dict[str, float]
    probabilities: dict[str, float]
    embedding_latency_ms: float = 0.0
    inference_latency_ms: float = 0.0
    graph_nodes: int = 0  # Total nodes in graph (queries + LLMs)
    graph_edges: int = 0  # Total edges used for message passing

    @property
    def total_latency_ms(self) -> float:
        return self.embedding_latency_ms + self.inference_latency_ms


@dataclass
class GraphRouterConfig:
    """Configuration for graph-aware router."""

    embed_dim: int = 512
    hidden_dim: int = 384
    num_gat_layers: int = 3
    num_gat_heads: int = 4
    dropout: float = 0.4
    edge_feature_dim: int = 6
    num_llms: int = 8

    # Inference preferences
    perf_weight: float = 0.7
    cost_weight: float = 0.2
    latency_weight: float = 0.1

    # OpenAI embedding config
    openai_model: str = "text-embedding-3-large"
    openai_dimensions: int = 512


class GraphAwareRouter:
    """Graph-aware EB1 Router with full heterogeneous graph context.

    This router maintains the full training graph in memory and appends
    new queries as nodes for inference, enabling proper GNN message passing.
    """

    def __init__(
        self,
        model: torch.nn.Module,
        graph_cache: dict,
        config: GraphRouterConfig,
        device: str = "cpu",
    ):
        self.model = model
        self.config = config
        self.device = device

        # Extract graph data
        self.query_embeddings = torch.from_numpy(
            graph_cache["query_embeddings"]
        ).float().to(device)
        self.task_embeddings = torch.from_numpy(
            graph_cache["task_embeddings"]
        ).float().to(device)
        self.llm_embeddings = torch.from_numpy(
            graph_cache["llm_embeddings"]
        ).float().to(device)
        self.combined_edge = torch.from_numpy(
            graph_cache["combined_edge"]
        ).float().to(device)
        self.source_nodes = torch.from_numpy(
            graph_cache["source_nodes"]
        ).long().to(device)
        self.dest_nodes = torch.from_numpy(
            graph_cache["dest_nodes"]
        ).long().to(device)
        self.train_mask = torch.from_numpy(
            graph_cache["train_mask"]
        ).bool().to(device)
        self.valid_mask = torch.from_numpy(
            graph_cache["valid_mask"]
        ).bool().to(device)

        self.num_queries = graph_cache["num_queries"]
        self.num_llms = graph_cache["num_llms"]

        # CRITICAL: Validate llm_order exists to ensure consistent key mapping
        if "llm_order" not in graph_cache:
            raise ValueError(
                "Training graph cache missing 'llm_order' key. "
                "This is required to ensure GNN indices map to correct models."
            )
        self.llm_names = graph_cache["llm_order"]

        # Pre-compute edge_can_see for training graph
        self.edge_can_see_base = self.train_mask | self.valid_mask

        # Initialize embedding client lazily
        self._embedding_client = None

        self.model.eval()

    def _get_embedding_client(self):
        """Lazy-load embedding client."""
        if self._embedding_client is None:
            from .eb1_router import EmbeddingClient
            self._embedding_client = EmbeddingClient(
                model=self.config.openai_model,
                dimensions=self.config.openai_dimensions,
            )
        return self._embedding_client

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

        Args:
            registry_name: Name of registry entry (e.g., "eb1-gnn-preview-12-19").
            registry_root: Root directory. Defaults to embergraphrouter/registry/.
            device: Device for inference ("cpu" or "cuda").

        Returns:
            Initialized GraphAwareRouter with full training graph.
        """
        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 config
        config_path = registry_path / "config.yaml"
        with open(config_path) as f:
            config_dict = yaml.safe_load(f)

        config = GraphRouterConfig(
            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),
            num_llms=config_dict.get("llm_num", 8),
        )

        # Load training graph cache
        cache_path = registry_path / "training_graph_cache.pkl"
        if not cache_path.exists():
            raise FileNotFoundError(
                f"Training graph cache not found: {cache_path}. "
                "Run scripts/package_inference_cache.py to create it."
            )
        with open(cache_path, "rb") as f:
            graph_cache = pickle.load(f)

        # Initialize model architecture
        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=config.embed_dim,
            llm_feature_dim=config.embed_dim,
            hidden_features=config.hidden_dim,
            in_edges=config.edge_feature_dim,
            num_gnn_layers=config.num_gat_layers,
            num_attention_heads=config.num_gat_heads,
            dropout=config.dropout,
            hidden_multiplier=2,
            use_contrastive=True,
            use_score_head=use_score_head,
            use_difficulty_head=True,
            use_uncertainty_head=False,
        )

        # Load weights (strict=False to ignore deprecated keys)
        model.load_state_dict(state_dict, strict=False)
        model.to(device)

        return cls(
            model=model,
            graph_cache=graph_cache,
            config=config,
            device=device,
        )

    def route(self, query: str) -> GraphRoutingResult:
        """Route a query using full graph context.

        Args:
            query: The query text to route.

        Returns:
            GraphRoutingResult with best model and confidence.
        """
        client = self._get_embedding_client()
        embedding, embed_latency = client.embed(query)

        start = time.perf_counter()
        result = self._route_from_embedding(embedding)
        inference_latency = (time.perf_counter() - start) * 1000

        result.embedding_latency_ms = embed_latency
        result.inference_latency_ms = inference_latency

        return result

    def route_from_embedding(
        self, embedding: NDArray[np.float32] | torch.Tensor
    ) -> GraphRoutingResult:
        """Route from pre-computed embedding."""
        return self._route_from_embedding(embedding)

    def _route_from_embedding(
        self, embedding: NDArray[np.float32] | torch.Tensor
    ) -> GraphRoutingResult:
        """Internal routing using full graph context."""
        if isinstance(embedding, np.ndarray):
            embedding = torch.from_numpy(embedding).float()
        embedding = embedding.to(self.device)

        if embedding.dim() == 1:
            embedding = embedding.unsqueeze(0)

        with torch.no_grad():
            logits = self._forward_with_graph(embedding)

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

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

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

        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 GraphRoutingResult(
            best_model=best_model,
            best_model_idx=best_idx,
            confidence=confidence,
            entropy=normalized_entropy,
            scores=scores,
            probabilities=probabilities,
            graph_nodes=self.num_queries + 1 + self.num_llms,
            graph_edges=int(self.edge_can_see_base.sum()) + self.num_llms,
        )

    def _forward_with_graph(self, query_embedding: torch.Tensor) -> torch.Tensor:
        """Run GNN forward pass with full graph context.

        This method:
        1. Appends the new query to existing query embeddings
        2. Builds extended edge index with new query's edges
        3. Runs GNN with train+valid edges visible for message passing
        4. Predicts only on the new query's edges
        """
        # New query index (appended after existing queries)
        new_query_idx = self.num_queries

        # Extend query and task embeddings with the new query
        # Use query embedding for both (task context = query for inference)
        extended_query_emb = torch.cat([self.query_embeddings, query_embedding], dim=0)
        extended_task_emb = torch.cat([self.task_embeddings, query_embedding], dim=0)

        # Build new edges: new_query -> all LLMs
        new_source = torch.full(
            (self.num_llms,), new_query_idx, dtype=torch.long, device=self.device
        )
        new_dest = torch.arange(self.num_llms, dtype=torch.long, device=self.device)

        # Extended edge index.
        # Nodes: [query_0..query_N, new_query, llm_0..llm_K]
        # LLM nodes start at index num_queries_extended, so offset dest.
        num_queries_extended = new_query_idx + 1
        extended_source = torch.cat([self.source_nodes, new_source])
        extended_dest = torch.cat([
            self.dest_nodes + num_queries_extended,
            new_dest + num_queries_extended,
        ])
        extended_edge_index = torch.stack([extended_source, extended_dest])

        # Extended edge features (use preference weights for new edges)
        new_edge_features = torch.zeros(self.num_llms, 3, device=self.device)
        # Set default costs/perfs/latencies (will be combined with preferences in augmentation)
        extended_combined_edge = torch.cat([self.combined_edge, new_edge_features], dim=0)

        # Augment edge features with preferences: [cost, perf, latency, wp, wc, wl]
        wp, wc, wl = self.config.perf_weight, self.config.cost_weight, self.config.latency_weight
        pref_features = torch.tensor(
            [[wp, wc, wl]], device=self.device
        ).expand(len(extended_combined_edge), 3)
        augmented_edge = torch.cat([extended_combined_edge, pref_features], dim=1)

        # Build masks
        total_edges = len(extended_source)
        num_new_edges = self.num_llms

        # edge_can_see: train + valid edges (existing), NO new edges visible
        edge_can_see = torch.cat([
            self.edge_can_see_base,
            torch.zeros(num_new_edges, dtype=torch.bool, device=self.device),
        ])

        # edge_mask: only predict on new query's edges
        edge_mask = torch.cat([
            torch.zeros(total_edges - num_new_edges, dtype=torch.bool, device=self.device),
            torch.ones(num_new_edges, dtype=torch.bool, device=self.device),
        ])

        # Forward through model
        logits = self.model(
            task_features=extended_task_emb,
            query_features=extended_query_emb,
            llm_features=self.llm_embeddings,
            edge_index=extended_edge_index,
            edge_mask=edge_mask,
            edge_can_see=edge_can_see,
            edge_weight=augmented_edge,
        )

        # logits contains predictions for masked edges only (new query's edges)
        # Shape: (num_llms,) -> reshape to (1, num_llms)
        if logits.dim() == 1:
            logits = logits.unsqueeze(0)

        return logits

    def benchmark(self, num_queries: int = 100) -> dict:
        """Benchmark graph-aware inference latency.

        Args:
            num_queries: Number of queries to benchmark.

        Returns:
            Dict with latency statistics.
        """
        embeddings = [
            np.random.randn(self.config.embed_dim).astype(np.float32)
            for _ in range(num_queries)
        ]

        latencies = []
        for emb in embeddings:
            start = time.perf_counter()
            _ = self._route_from_embedding(emb)
            latencies.append((time.perf_counter() - start) * 1000)

        latencies = np.array(latencies)

        return {
            "num_queries": num_queries,
            "mean_ms": float(np.mean(latencies)),
            "std_ms": float(np.std(latencies)),
            "p50_ms": float(np.percentile(latencies, 50)),
            "p95_ms": float(np.percentile(latencies, 95)),
            "p99_ms": float(np.percentile(latencies, 99)),
            "throughput_qps": 1000 / np.mean(latencies),
            "graph_queries": self.num_queries,
            "graph_edges_visible": int(self.edge_can_see_base.sum()),
        }
