"""Compatibility-based query routing inference pipeline.

This module provides the inference pipeline for the compatibility scorer,
which models query-to-LLM affinity for routing decisions.
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from pathlib import Path

import torch

from partner_inference_harness.embedding.embedder import TextEmbedder
from partner_inference_harness.model.hetero_graph import (
    CompatibilityConfig,
    CompatibilityScorer,
)


@dataclass
class GraphRoutingResult:
    """Result of a compatibility-based routing decision.

    Attributes:
        selected_model: Identifier of the selected model.
        confidence: Confidence score for the selection (0-1).
        latency_ms: Inference latency in milliseconds.
        model_scores: Dictionary mapping model IDs to their affinity scores.
    """

    selected_model: str
    confidence: float
    latency_ms: float
    model_scores: dict[str, float] = field(default_factory=dict)


# Default LLM descriptions for embedding
DEFAULT_LLM_DESCRIPTIONS = [
    "GPT-4: Advanced reasoning, complex problem-solving, nuanced text generation",
    "Claude: Thoughtful analysis, safety-focused, long-context understanding",
    "Gemini: Multimodal understanding, code generation, factual knowledge",
    "Llama: Open-source, efficient inference, instruction following",
    "Mixtral: Mixture of experts, fast inference, multilingual",
    "Codex: Code generation, programming assistance, technical documentation",
    "PaLM: Large-scale language understanding, summarization, translation",
]


class GraphRouter:
    """Compatibility-based query routing inference pipeline.

    Uses a dual-encoder architecture to model query-LLM compatibility.
    The model learns to score:
    - Query characteristics (from query embeddings)
    - LLM capabilities (from LLM description embeddings)
    - Query-LLM affinity (from learned compatibility scoring)

    Attributes:
        device: Compute device (cuda/cpu).
        model_list: List of model identifiers.
        llm_descriptions: Text descriptions of each LLM's capabilities.
    """

    def __init__(
        self,
        artifacts_dir: Path | str | None = None,
        device: str | None = None,
        model_list: list[str] | None = None,
        llm_descriptions: list[str] | None = None,
    ) -> None:
        """Initialize the graph router.

        Args:
            artifacts_dir: Directory containing model artifacts.
            device: Compute device. Defaults to CUDA if available.
            model_list: List of model identifiers.
            llm_descriptions: Text descriptions of each LLM's capabilities.
                Used to generate LLM node embeddings.
        """
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model_list = model_list or [
            "gpt-4",
            "claude",
            "gemini",
            "llama",
            "mixtral",
            "codex",
            "palm",
        ]
        self.llm_descriptions = llm_descriptions or DEFAULT_LLM_DESCRIPTIONS

        if len(self.llm_descriptions) != len(self.model_list):
            raise ValueError(
                f"Number of LLM descriptions ({len(self.llm_descriptions)}) "
                f"must match number of models ({len(self.model_list)})"
            )

        if artifacts_dir is not None:
            self._artifacts_dir = Path(artifacts_dir)
        else:
            self._artifacts_dir = (
                Path(__file__).parent.parent / "artifacts" / "graph_router"
            )

        self._embedder: TextEmbedder | None = None
        self._router_model: CompatibilityScorer | None = None
        self._llm_embeddings: torch.Tensor | None = None
        self._loaded = False

    def _ensure_loaded(self) -> None:
        """Lazily load models on first use."""
        if self._loaded:
            return

        # Load text embedder
        self._embedder = TextEmbedder(device=self.device)

        # Load or create compatibility scorer
        model_path = self._artifacts_dir / "graph_model.pth"
        if model_path.exists():
            from partner_inference_harness.model.hetero_graph import load_model
            self._router_model = load_model(str(model_path), device=self.device)
        else:
            # Create with random init for testing
            config = CompatibilityConfig(num_llms=len(self.model_list))
            self._router_model = CompatibilityScorer(config)
            self._router_model.to(self.device)
            self._router_model.eval()

        # Embed LLM descriptions
        with torch.no_grad():
            self._llm_embeddings = self._embedder.embed_batch(self.llm_descriptions)

        # Load model list if available
        model_list_path = self._artifacts_dir / "llm_list.json"
        if model_list_path.exists():
            with open(model_list_path) as f:
                data = json.load(f)
                if "models" in data:
                    self.model_list = data["models"]
                if "descriptions" in data:
                    self.llm_descriptions = data["descriptions"]
                    # Re-embed with new descriptions
                    self._llm_embeddings = self._embedder.embed_batch(
                        self.llm_descriptions
                    )

        self._loaded = True

    def route(self, query: str) -> GraphRoutingResult:
        """Route a single query to a model.

        Args:
            query: User query text.

        Returns:
            GraphRoutingResult with selected model and metadata.
        """
        return self.route_batch([query])[0]

    def route_batch(self, queries: list[str]) -> list[GraphRoutingResult]:
        """Route a batch of queries to models.

        Args:
            queries: List of user query texts.

        Returns:
            List of GraphRoutingResult objects, one per query.
        """
        self._ensure_loaded()
        assert self._embedder is not None
        assert self._router_model is not None
        assert self._llm_embeddings is not None

        start_time = time.perf_counter()

        # Embed queries
        query_embeddings = self._embedder.embed_batch(queries)

        # Route through compatibility scorer
        with torch.no_grad():
            logits = self._router_model(query_embeddings, self._llm_embeddings)
            probabilities = torch.softmax(logits, dim=-1)
            confidences, selected_indices = torch.max(probabilities, dim=-1)

        end_time = time.perf_counter()
        total_latency_ms = (end_time - start_time) * 1000
        per_query_latency_ms = total_latency_ms / len(queries)

        # Convert to results
        results = []
        for i, (idx, conf) in enumerate(
            zip(selected_indices.tolist(), confidences.tolist())
        ):
            scores = {
                model: float(prob)
                for model, prob in zip(self.model_list, probabilities[i].tolist())
            }

            results.append(
                GraphRoutingResult(
                    selected_model=self.model_list[idx],
                    confidence=conf,
                    latency_ms=per_query_latency_ms,
                    model_scores=scores,
                )
            )

        return results

    def get_memory_usage_gb(self) -> float:
        """Get total GPU memory usage estimate in GB.

        Returns:
            Estimated memory usage in gigabytes.
        """
        self._ensure_loaded()
        assert self._embedder is not None
        assert self._router_model is not None

        embedder_gb = self._embedder.get_memory_footprint_gb()

        # Estimate router model size
        param_bytes = sum(
            p.numel() * p.element_size()
            for p in self._router_model.parameters()
        )
        router_gb = param_bytes / (1024 ** 3)

        return embedder_gb + router_gb

    def warmup(self, num_queries: int = 10) -> float:
        """Run warmup inference passes.

        Args:
            num_queries: Number of warmup queries to run.

        Returns:
            Average latency in milliseconds after warmup.
        """
        self._ensure_loaded()

        dummy_queries = [f"warmup query {i}" for i in range(num_queries)]
        results = self.route_batch(dummy_queries)

        return sum(r.latency_ms for r in results) / len(results)

    def unload(self) -> None:
        """Unload models from memory."""
        if self._embedder is not None:
            self._embedder.unload()
            self._embedder = None
        if self._router_model is not None:
            del self._router_model
            self._router_model = None
        if self._llm_embeddings is not None:
            del self._llm_embeddings
            self._llm_embeddings = None
        self._loaded = False
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


def create_graph_router(
    artifacts_dir: Path | str | None = None,
    device: str | None = None,
    warmup: bool = True,
) -> GraphRouter:
    """Factory function to create and optionally warm up a graph router.

    Args:
        artifacts_dir: Directory containing model artifacts.
        device: Compute device.
        warmup: Whether to run warmup inference passes.

    Returns:
        Initialized GraphRouter instance.
    """
    router = GraphRouter(artifacts_dir=artifacts_dir, device=device)

    if warmup:
        router.warmup()

    return router
