"""Base classes for inference strategies.

This module defines the abstract interfaces for routing/inference strategies
in the EB1 evaluation system. Strategies wrap different ways of answering
queries: single models, routers, ensembles, etc.
"""

import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any


class StrategyType(Enum):
    """Types of inference strategies."""

    STANDALONE = "standalone"  # Single model, no routing
    ROUTER = "router"  # Model router (TieredRouter, EB1Router)
    ENSEMBLE = "ensemble"  # Multi-model ensemble
    CONDUCTOR = "conductor"  # LLM+GNN conductor system
    ORACLE = "oracle"  # Always picks best model (requires ground truth)
    RANDOM = "random"  # Random model selection (baseline)
    CACHED = "cached"  # Uses cached responses (counterfactual evaluation)


@dataclass
class RoutingDecision:
    """Details of a routing decision (for router strategies).

    Attributes:
        tier: Which tier the query was classified into.
        selected_model: The model chosen for this query.
        confidence: Confidence in the routing decision.
        signals: Raw signals used for routing (mean_logit, margin, etc.)
        rationale: Human-readable explanation of the decision.
        alternative_models: Other models considered (ranked).
    """

    tier: str | None = None
    selected_model: str | None = None
    confidence: float | None = None
    signals: dict[str, float] = field(default_factory=dict)
    rationale: str | None = None
    alternative_models: list[tuple[str, float]] | None = None


@dataclass
class StrategyResult:
    """Result from a single inference.

    Attributes:
        prediction: The model's text output.
        model_used: Which LLM actually generated the response.
        routing: Routing decision details (if applicable).
        latency_ms: End-to-end latency in milliseconds.
        cost_usd: Estimated cost in USD.
        tokens_in: Number of input tokens.
        tokens_out: Number of output tokens.
        metadata: Strategy-specific additional info.
        error: Error message if inference failed.
    """

    prediction: str
    model_used: str
    routing: RoutingDecision | None = None
    latency_ms: float = 0.0
    cost_usd: float = 0.0
    tokens_in: int = 0
    tokens_out: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)
    error: str | None = None

    @property
    def success(self) -> bool:
        """Whether inference completed successfully."""
        return self.error is None

    @property
    def tier(self) -> str | None:
        """Shortcut to routing tier."""
        return self.routing.tier if self.routing else None


class Strategy(ABC):
    """Abstract base class for inference strategies.

    Subclasses must implement:
    - infer(): Run inference on a single query

    Optionally override:
    - infer_batch(): Batch inference with parallelism
    - warm_up(): Pre-load models/caches
    """

    # Class attributes to be set by subclasses
    name: str
    strategy_type: StrategyType
    description: str = ""

    # Whether this strategy requires query embeddings
    requires_embedding: bool = False

    # Models this strategy can use
    available_models: list[str] = []

    @abstractmethod
    async def infer(
        self,
        query: str,
        system_prompt: str | None = None,
        **kwargs: Any,
    ) -> StrategyResult:
        """Run inference on a single query.

        Args:
            query: The input query/prompt.
            system_prompt: Optional system prompt.
            **kwargs: Additional strategy-specific options.

        Returns:
            StrategyResult with prediction and metadata.
        """

    async def infer_batch(
        self,
        queries: list[str],
        system_prompt: str | None = None,
        max_concurrent: int = 10,
        **kwargs: Any,
    ) -> list[StrategyResult]:
        """Run inference on multiple queries.

        Default implementation runs sequentially. Subclasses can override
        for parallel execution.

        Args:
            queries: List of input queries.
            system_prompt: Optional system prompt (applied to all).
            max_concurrent: Maximum concurrent requests.
            **kwargs: Additional strategy-specific options.

        Returns:
            List of StrategyResult, one per query.
        """
        results = []
        for query in queries:
            result = await self.infer(query, system_prompt, **kwargs)
            results.append(result)
        return results

    def warm_up(self) -> None:  # noqa: B027
        """Pre-load models, caches, or other resources.

        Called once before evaluation begins. Default is no-op.
        """
        pass

    def get_cost_per_1k_tokens(self, model: str) -> tuple[float, float]:
        """Get cost per 1K tokens (input, output) for a model.

        Returns (0.0, 0.0) if cost info not available.
        """
        # Default pricing (can be overridden by subclasses)
        pricing = {
            "gemini-2.5-flash-lite": (0.0001, 0.0004),
            "claude-haiku-4.5": (0.001, 0.005),
            "gpt-5-nano": (0.00005, 0.0004),
            "gpt-5.1": (0.00125, 0.01),
            "gpt-5.1-codex": (0.00125, 0.01),
            "claude-opus-4.5": (0.005, 0.025),
            "gemini-3-flash-preview": (0.0005, 0.003),
            "gemini-3-pro-preview": (0.002, 0.012),
        }
        return pricing.get(model.lower(), (0.0, 0.0))

    def estimate_cost(
        self,
        tokens_in: int,
        tokens_out: int,
        model: str,
    ) -> float:
        """Estimate cost in USD for a request."""
        in_cost, out_cost = self.get_cost_per_1k_tokens(model)
        return (tokens_in * in_cost + tokens_out * out_cost) / 1000

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(name='{self.name}')"


class TimedStrategy(Strategy):
    """Mixin that automatically measures latency."""

    async def infer(
        self,
        query: str,
        system_prompt: str | None = None,
        **kwargs: Any,
    ) -> StrategyResult:
        """Wrap inference with timing."""
        start = time.perf_counter()
        result = await self._infer_impl(query, system_prompt, **kwargs)
        elapsed_ms = (time.perf_counter() - start) * 1000

        # Update latency if not already set
        if result.latency_ms == 0.0:
            result.latency_ms = elapsed_ms

        return result

    @abstractmethod
    async def _infer_impl(
        self,
        query: str,
        system_prompt: str | None = None,
        **kwargs: Any,
    ) -> StrategyResult:
        """Actual inference implementation (without timing)."""
