"""TieredRouter strategy wrapper for evaluation.

This module wraps the production TieredRouter as a Strategy for the evaluation
system, enabling systematic comparison against baselines and other routers.
"""

import asyncio
import time
from pathlib import Path
from typing import Any

from embergraphrouter.evals.strategies.base import (
    RoutingDecision,
    Strategy,
    StrategyResult,
    StrategyType,
)


def _extract_usage(response: Any) -> tuple[int, int]:
    """Extract token counts from response object.

    Handles both ChatResponse.usage (UsageStats) and Response.usage (dict).
    """
    if not hasattr(response, "usage") or response.usage is None:
        return 0, 0

    usage = response.usage

    # ChatResponse.usage is UsageStats with prompt_tokens/completion_tokens
    if hasattr(usage, "prompt_tokens"):
        return usage.prompt_tokens, usage.completion_tokens

    # Response.usage is a dict with prompt_tokens/completion_tokens
    if isinstance(usage, dict):
        return usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0)

    return 0, 0


class TieredRouterStrategy(Strategy):
    """Strategy wrapper for TieredRouter.

    Wraps the production TieredRouter to enable evaluation against baselines.
    Extracts full routing signals for analysis.

    Example:
        strategy = TieredRouterStrategy.from_registry("eb1-gnn-preview-latest")
        result = await strategy.infer("What is 2 + 2?")
        print(f"Tier: {result.tier}, Model: {result.model_used}")
    """

    name = "tiered_router"
    strategy_type = StrategyType.ROUTER
    description = "EB1-GNN tiered router with difficulty-aware model selection"
    requires_embedding = True

    def __init__(
        self,
        registry_name: str = "eb1-gnn-preview-latest",
        registry_root: str | Path | None = None,
        device: str = "cpu",
        execute_inference: bool = True,
    ):
        """Initialize TieredRouterStrategy.

        Args:
            registry_name: Name of model registry entry.
            registry_root: Root directory for registry. None uses default.
            device: Device for GNN inference ("cpu" or "cuda").
            execute_inference: If True, execute full inference.
                If False, only do routing classification (useful for testing).
        """
        self._registry_name = registry_name
        self._registry_root = registry_root
        self._device = device
        self._execute_inference = execute_inference
        self._router = None  # Lazy initialization
        self.name = f"tiered_router:{registry_name}"

    def _get_router(self):
        """Lazy-load TieredRouter."""
        if self._router is None:
            from embergraphrouter.inference.tiered_router import TieredRouter

            self._router = TieredRouter.from_registry(
                self._registry_name,
                registry_root=self._registry_root,
                device=self._device,
            )
        return self._router

    def warm_up(self) -> None:
        """Pre-load the TieredRouter and GNN model."""
        self._get_router()

    async def infer(
        self,
        query: str,
        system_prompt: str | None = None,
        **kwargs: Any,
    ) -> StrategyResult:
        """Run inference through TieredRouter.

        Args:
            query: Input query text.
            system_prompt: Optional system prompt.
            **kwargs: Additional options passed to model.

        Returns:
            StrategyResult with prediction and routing details.
        """
        router = self._get_router()

        start_time = time.perf_counter()
        try:
            if self._execute_inference:
                # Full inference
                response = await asyncio.to_thread(
                    router.route_and_execute,
                    query,
                    system=system_prompt,
                    shared_options=kwargs or None,
                )

                # Extract prediction and metadata from response
                # route_and_execute returns Response wrapper or ChatResponse
                # Response wrapper has .text, ChatResponse has .data
                if hasattr(response, "text"):
                    prediction = response.text  # Response wrapper
                else:
                    prediction = response.data  # ChatResponse dataclass

                model_used = getattr(response, "model_id", None) or "unknown"

                # Extract routing metadata from response
                # Response wrapper has .chat_response.raw_output, ChatResponse has .raw_output
                if hasattr(response, "chat_response"):
                    raw_output = getattr(response.chat_response, "raw_output", None)
                else:
                    raw_output = getattr(response, "raw_output", None)
                routing_meta = (
                    raw_output.get("tiered_routing", {})
                    if isinstance(raw_output, dict)
                    else {}
                )

                tokens_in, tokens_out = _extract_usage(response)

            else:
                # Classification only (for testing)
                routing_result = await asyncio.to_thread(router.classify, query)
                prediction = ""
                model_used = routing_result.selected_model

                routing_meta = {
                    "tier": routing_result.tier.name,
                    "selected_model": routing_result.selected_model,
                    "rationale": routing_result.rationale,
                    "signals": {
                        "mean_logit": routing_result.signals.mean_logit,
                        "max_prob": routing_result.signals.max_prob,
                        "logit_margin": routing_result.signals.logit_margin,
                        "logit_std": routing_result.signals.logit_std,
                        "entropy": routing_result.signals.entropy,
                        "top1_model": routing_result.signals.top1_model,
                        "top2_model": routing_result.signals.top2_model,
                    },
                    "latency_ms": {
                        "embedding": routing_result.embedding_latency_ms,
                        "classification": routing_result.classification_latency_ms,
                        "total_routing": routing_result.total_routing_latency_ms,
                    },
                }
                tokens_in = 0
                tokens_out = 0

            elapsed_ms = (time.perf_counter() - start_time) * 1000

            # Build routing decision from extracted metadata
            signals = routing_meta.get("signals", {})
            routing = RoutingDecision(
                tier=routing_meta.get("tier"),
                selected_model=routing_meta.get("selected_model"),
                confidence=signals.get("max_prob"),
                signals={
                    "mean_logit": signals.get("mean_logit", 0.0),
                    "max_prob": signals.get("max_prob", 0.0),
                    "logit_margin": signals.get("logit_margin", 0.0),
                    "logit_std": signals.get("logit_std", 0.0),
                    "entropy": signals.get("entropy", 0.0),
                },
                rationale=routing_meta.get("rationale"),
            )

            # Estimate cost
            cost_usd = self.estimate_cost(tokens_in, tokens_out, model_used)

            return StrategyResult(
                prediction=prediction,
                model_used=model_used,
                routing=routing,
                latency_ms=elapsed_ms,
                cost_usd=cost_usd,
                tokens_in=tokens_in,
                tokens_out=tokens_out,
                metadata={
                    "registry_name": self._registry_name,
                    "routing_latency_ms": routing_meta.get("latency_ms", {}),
                },
            )

        except Exception as e:
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return StrategyResult(
                prediction="",
                model_used="error",
                latency_ms=elapsed_ms,
                error=str(e),
            )

    @classmethod
    def from_registry(
        cls,
        registry_name: str,
        registry_root: str | Path | None = None,
        device: str = "cpu",
        execute_inference: bool = True,
    ) -> "TieredRouterStrategy":
        """Create strategy from registry entry.

        Args:
            registry_name: Name of model registry entry.
            registry_root: Root directory for registry.
            device: Device for GNN inference.
            execute_inference: Whether to execute full inference.

        Returns:
            Configured TieredRouterStrategy.
        """
        return cls(
            registry_name=registry_name,
            registry_root=registry_root,
            device=device,
            execute_inference=execute_inference,
        )

    @property
    def available_models(self) -> list[str]:
        """Return models available through this router."""
        router = self._get_router()
        return router.llm_names


class SingleModelStrategy(Strategy):
    """Strategy that always uses a single fixed model.

    Useful as a baseline for comparison with routing strategies.
    """

    name = "single_model"
    strategy_type = StrategyType.STANDALONE
    description = "Fixed single model baseline"
    requires_embedding = False

    def __init__(self, model: str):
        """Initialize with fixed model.

        Args:
            model: Model identifier to use for all queries.
        """
        self._model = model
        self.name = f"single_model:{model}"
        self.available_models = [model]

    async def infer(
        self,
        query: str,
        system_prompt: str | None = None,
        **kwargs: Any,
    ) -> StrategyResult:
        """Run inference with fixed model.

        Args:
            query: Input query text.
            system_prompt: Optional system prompt.
            **kwargs: Additional options passed to model.

        Returns:
            StrategyResult with prediction.
        """
        from ember.api.models import models

        start_time = time.perf_counter()
        try:
            options = dict(kwargs)
            if system_prompt:
                options["system"] = system_prompt

            response = await asyncio.to_thread(
                models.response,
                self._model,
                query,
                **options,
            )

            # models.response() returns Response with .text
            prediction = response.text
            tokens_in, tokens_out = _extract_usage(response)

            elapsed_ms = (time.perf_counter() - start_time) * 1000
            cost_usd = self.estimate_cost(tokens_in, tokens_out, self._model)

            return StrategyResult(
                prediction=prediction,
                model_used=self._model,
                latency_ms=elapsed_ms,
                cost_usd=cost_usd,
                tokens_in=tokens_in,
                tokens_out=tokens_out,
            )

        except Exception as e:
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return StrategyResult(
                prediction="",
                model_used=self._model,
                latency_ms=elapsed_ms,
                error=str(e),
            )


class RandomModelStrategy(Strategy):
    """Strategy that randomly selects from available models.

    Useful as a baseline to show routing adds value.
    """

    name = "random_model"
    strategy_type = StrategyType.RANDOM
    description = "Random model selection baseline"
    requires_embedding = False

    def __init__(self, models: list[str], seed: int | None = None):
        """Initialize with model pool.

        Args:
            models: List of model identifiers to randomly choose from.
            seed: Random seed for reproducibility.
        """
        import random

        self._models = models
        self._rng = random.Random(seed)
        self.available_models = models
        self.name = f"random_model:{len(models)}_models"

    async def infer(
        self,
        query: str,
        system_prompt: str | None = None,
        **kwargs: Any,
    ) -> StrategyResult:
        """Run inference with randomly selected model.

        Args:
            query: Input query text.
            system_prompt: Optional system prompt.
            **kwargs: Additional options passed to model.

        Returns:
            StrategyResult with prediction.
        """
        from ember.api.models import models as models_api

        selected_model = self._rng.choice(self._models)

        start_time = time.perf_counter()
        try:
            options = dict(kwargs)
            if system_prompt:
                options["system"] = system_prompt

            response = await asyncio.to_thread(
                models_api.response,
                selected_model,
                query,
                **options,
            )

            # models_api.response() returns Response with .text
            prediction = response.text
            tokens_in, tokens_out = _extract_usage(response)

            elapsed_ms = (time.perf_counter() - start_time) * 1000
            cost_usd = self.estimate_cost(tokens_in, tokens_out, selected_model)

            return StrategyResult(
                prediction=prediction,
                model_used=selected_model,
                routing=RoutingDecision(
                    tier="random",
                    selected_model=selected_model,
                ),
                latency_ms=elapsed_ms,
                cost_usd=cost_usd,
                tokens_in=tokens_in,
                tokens_out=tokens_out,
            )

        except Exception as e:
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return StrategyResult(
                prediction="",
                model_used=selected_model,
                latency_ms=elapsed_ms,
                error=str(e),
            )


__all__ = [
    "TieredRouterStrategy",
    "SingleModelStrategy",
    "RandomModelStrategy",
]
