"""Query routing inference pipeline.

This module provides the main inference pipeline that combines text embedding
and model routing into a single interface for query-to-model selection.
"""

from __future__ import annotations

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

import torch

from partner_inference_harness.embedding.embedder import TextEmbedder
from partner_inference_harness.model.router_mlp import RouterMLP, load_model

if TYPE_CHECKING:
    pass


@dataclass
class RoutingResult:
    """Result of a 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 scores.
    """

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


class Router:
    """Query routing inference pipeline.

    Combines text embedding and MLP routing to select an appropriate language
    model for a given query. This is a load-testing harness with placeholder
    weights.

    Attributes:
        device: Compute device (cuda/cpu).
        model_list: List of model identifiers that can be routed to.
    """

    DEFAULT_MODELS = [
        "model_a",
        "model_b",
        "model_c",
        "model_d",
        "model_e",
        "model_f",
        "model_g",
    ]

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

        Args:
            artifacts_dir: Directory containing model artifacts. If None, uses
                default placeholder location.
            device: Compute device. Defaults to CUDA if available.
            model_list: List of model identifiers. Defaults to placeholder list.
        """
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model_list = model_list or self.DEFAULT_MODELS

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

        self._embedder: TextEmbedder | None = None
        self._router_model: RouterMLP | None = None
        self._loaded = False

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

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

        # Load router model
        model_path = self._artifacts_dir / "model.pth"
        if model_path.exists():
            self._router_model = load_model(model_path, device=self.device)
        else:
            # If no weights exist, create with random init for testing
            from partner_inference_harness.model.router_mlp import RouterConfig

            config = RouterConfig(num_models=len(self.model_list))
            self._router_model = RouterMLP(config)
            self._router_model.to(self.device)
            self._router_model.eval()

        # 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"]

        self._loaded = True

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

        Args:
            query: User query text.

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

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

        Args:
            queries: List of user query texts.

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

        start_time = time.perf_counter()

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

        # Route through MLP
        with torch.no_grad():
            logits = self._router_model(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(
                RoutingResult(
                    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()
        router_mb = self._router_model.get_memory_footprint_mb()
        router_gb = router_mb / 1024

        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()

        # Generate dummy queries
        dummy_queries = [f"warmup query {i}" for i in range(num_queries)]

        # Run warmup
        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
        self._loaded = False
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


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

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

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

    if warmup:
        router.warmup()

    return router
