"""Oracle strategy that always selects the optimal model.

The Oracle represents the theoretical best routing - it knows which model
will perform best for each query. This enables computing frac_of_oracle,
the key metric showing how close a router gets to optimal performance.
"""

import asyncio
import time
from dataclasses import dataclass, field
from typing import Any, Callable

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


@dataclass
class OracleEntry:
    """Pre-computed performance data for a single query.

    Attributes:
        query_id: Unique identifier for the query.
        query: The query text.
        best_model: Model with highest quality for this query.
        best_score: Quality score achieved by best model.
        model_scores: Mapping of model name to quality score.
        model_responses: Optional cached responses per model.
    """

    query_id: str
    query: str
    best_model: str
    best_score: float
    model_scores: dict[str, float] = field(default_factory=dict)
    model_responses: dict[str, str] = field(default_factory=dict)


class OracleStrategy(Strategy):
    """Strategy that always picks the optimal model for each query.

    The Oracle can operate in two modes:

    1. **Lookup mode** (default): Uses pre-computed oracle data to instantly
       return the best model's response. Requires an oracle bank.

    2. **Compute mode**: Runs all available models on each query and selects
       the one with the highest quality. Expensive but doesn't require pre-computation.

    Example (lookup mode):
        oracle_data = {
            "q1": OracleEntry(query_id="q1", query="...", best_model="gpt-4", ...),
            "q2": OracleEntry(query_id="q2", query="...", best_model="claude", ...),
        }
        oracle = OracleStrategy.from_bank(oracle_data)
        result = await oracle.infer("What is 2+2?", query_id="q1")

    Example (compute mode):
        models = ["gpt-4", "claude-3-opus", "gemini-pro"]
        oracle = OracleStrategy.compute_mode(models, evaluator=my_eval_fn)
        result = await oracle.infer("What is 2+2?")
    """

    name = "oracle"
    strategy_type = StrategyType.ORACLE
    description = "Optimal model selection with perfect foresight"
    requires_embedding = False

    def __init__(
        self,
        oracle_bank: dict[str, OracleEntry] | None = None,
        models: list[str] | None = None,
        evaluator: Callable[[str, str], float] | None = None,
        match_by_content: bool = False,
    ):
        """Initialize OracleStrategy.

        Args:
            oracle_bank: Pre-computed oracle entries keyed by query_id.
            models: List of models for compute mode.
            evaluator: Function(prediction, reference) -> score for compute mode.
            match_by_content: If True, match queries by content hash instead of ID.
        """
        self._oracle_bank = oracle_bank or {}
        self._models = models or []
        self._evaluator = evaluator
        self._match_by_content = match_by_content
        self._query_hash_map: dict[str, str] = {}

        # Build content hash map for matching
        if match_by_content and oracle_bank:
            for query_id, entry in oracle_bank.items():
                content_hash = self._hash_query(entry.query)
                self._query_hash_map[content_hash] = query_id

        self.available_models = list(
            set(
                model
                for entry in self._oracle_bank.values()
                for model in entry.model_scores.keys()
            )
        ) or self._models

    @staticmethod
    def _hash_query(query: str) -> str:
        """Create a hash of query content for matching."""
        import hashlib

        normalized = " ".join(query.lower().split())
        return hashlib.sha256(normalized.encode()).hexdigest()[:16]

    @classmethod
    def from_bank(
        cls,
        oracle_bank: dict[str, OracleEntry],
        match_by_content: bool = False,
    ) -> "OracleStrategy":
        """Create Oracle from pre-computed response bank.

        Args:
            oracle_bank: Mapping of query_id to OracleEntry.
            match_by_content: Match by query content hash instead of ID.

        Returns:
            Configured OracleStrategy in lookup mode.
        """
        return cls(oracle_bank=oracle_bank, match_by_content=match_by_content)

    @classmethod
    def compute_mode(
        cls,
        models: list[str],
        evaluator: Callable[[str, str], float],
    ) -> "OracleStrategy":
        """Create Oracle that computes optimal model per query.

        This mode runs all models on each query and selects the best.
        Expensive but doesn't require pre-computed data.

        Args:
            models: List of model identifiers to try.
            evaluator: Function(prediction, reference) -> quality score.

        Returns:
            Configured OracleStrategy in compute mode.
        """
        return cls(models=models, evaluator=evaluator)

    def _lookup_entry(self, query: str, query_id: str | None) -> OracleEntry | None:
        """Look up oracle entry by ID or content hash."""
        if query_id and query_id in self._oracle_bank:
            return self._oracle_bank[query_id]

        if self._match_by_content:
            content_hash = self._hash_query(query)
            if content_hash in self._query_hash_map:
                entry_id = self._query_hash_map[content_hash]
                return self._oracle_bank.get(entry_id)

        return None

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

        In lookup mode, returns cached best response.
        In compute mode, runs all models and picks best.

        Args:
            query: Input query text.
            system_prompt: Optional system prompt.
            query_id: Query identifier for oracle bank lookup.
            reference: Ground truth for compute mode evaluation.
            **kwargs: Additional options.

        Returns:
            StrategyResult with optimal model's response.
        """
        start_time = time.perf_counter()

        # Try lookup mode first
        entry = self._lookup_entry(query, query_id)
        if entry:
            elapsed_ms = (time.perf_counter() - start_time) * 1000

            # Get cached response if available
            prediction = entry.model_responses.get(entry.best_model, "")

            return StrategyResult(
                prediction=prediction,
                model_used=entry.best_model,
                routing=RoutingDecision(
                    tier="oracle",
                    selected_model=entry.best_model,
                    confidence=1.0,
                    rationale=f"Oracle selection (score={entry.best_score:.3f})",
                    signals={
                        "oracle_score": entry.best_score,
                        "n_models_compared": len(entry.model_scores),
                    },
                ),
                latency_ms=elapsed_ms,
                cost_usd=0.0,  # Lookup is free
                tokens_in=0,
                tokens_out=0,
                metadata={
                    "mode": "lookup",
                    "query_id": entry.query_id,
                    "all_scores": entry.model_scores,
                },
            )

        # Compute mode: run all models
        if self._models and self._evaluator and reference:
            return await self._compute_oracle(
                query, reference, system_prompt, **kwargs
            )

        # No oracle data available
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        return StrategyResult(
            prediction="",
            model_used="none",
            latency_ms=elapsed_ms,
            error="No oracle data available for query",
        )

    async def _compute_oracle(
        self,
        query: str,
        reference: str,
        system_prompt: str | None,
        **kwargs: Any,
    ) -> StrategyResult:
        """Compute oracle by running all models."""
        from ember.api.models import models as models_api

        start_time = time.perf_counter()
        model_results: dict[str, tuple[str, float, float, int, int]] = {}

        # Run all models
        for model in self._models:
            try:
                options = dict(kwargs)
                if system_prompt:
                    options["system"] = system_prompt

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

                # models_api.response() returns Response with .text and .usage dict
                prediction = response.text
                score = self._evaluator(prediction, reference)

                usage = response.usage or {}
                tokens_in = usage.get("prompt_tokens", 0)
                tokens_out = usage.get("completion_tokens", 0)

                cost = self.estimate_cost(tokens_in, tokens_out, model)
                model_results[model] = (prediction, score, cost, tokens_in, tokens_out)

            except Exception:
                model_results[model] = ("", 0.0, 0.0, 0, 0)

        if not model_results:
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return StrategyResult(
                prediction="",
                model_used="error",
                latency_ms=elapsed_ms,
                error="All models failed",
            )

        # Select best model
        best_model = max(model_results.keys(), key=lambda m: model_results[m][1])
        prediction, score, cost, tokens_in, tokens_out = model_results[best_model]

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

        return StrategyResult(
            prediction=prediction,
            model_used=best_model,
            routing=RoutingDecision(
                tier="oracle",
                selected_model=best_model,
                confidence=1.0,
                rationale=f"Computed oracle (score={score:.3f})",
                signals={
                    "oracle_score": score,
                    "n_models_compared": len(self._models),
                },
            ),
            latency_ms=elapsed_ms,
            cost_usd=sum(r[2] for r in model_results.values()),
            tokens_in=sum(r[3] for r in model_results.values()),
            tokens_out=sum(r[4] for r in model_results.values()),
            metadata={
                "mode": "compute",
                "all_scores": {m: r[1] for m, r in model_results.items()},
            },
        )


class CachedResponseStrategy(Strategy):
    """Strategy that returns cached responses from a response bank.

    Enables counterfactual evaluation: test routing decisions without
    re-running expensive LLM calls. Essential for rapid router iteration.

    Example:
        # Build response bank by running all models once
        bank = ResponseBank()
        for query_id, query in queries:
            for model in models:
                response = await run_model(model, query)
                bank.add(query_id, model, response)

        # Evaluate different routers against the same responses
        cached = CachedResponseStrategy(bank, model_selector=my_router)
        result = await cached.infer(query, query_id="q1")
    """

    name = "cached_response"
    strategy_type = StrategyType.CACHED
    description = "Returns cached responses for counterfactual evaluation"
    requires_embedding = False

    def __init__(
        self,
        response_bank: dict[str, dict[str, str]],
        model_selector: Callable[[str], str] | None = None,
        default_model: str | None = None,
    ):
        """Initialize CachedResponseStrategy.

        Args:
            response_bank: Mapping of query_id -> {model: response}.
            model_selector: Function(query) -> model_name for selection.
            default_model: Fallback model if selector not provided.
        """
        self._response_bank = response_bank
        self._model_selector = model_selector
        self._default_model = default_model

        # Collect all available models
        all_models = set()
        for model_responses in response_bank.values():
            all_models.update(model_responses.keys())
        self.available_models = sorted(all_models)

    async def infer(
        self,
        query: str,
        system_prompt: str | None = None,
        query_id: str | None = None,
        model: str | None = None,
        **kwargs: Any,
    ) -> StrategyResult:
        """Return cached response for the query.

        Args:
            query: Input query text.
            system_prompt: Ignored (responses already generated).
            query_id: Required for bank lookup.
            model: Specific model to use, or uses selector/default.
            **kwargs: Additional options (ignored).

        Returns:
            StrategyResult with cached response.
        """
        start_time = time.perf_counter()

        if not query_id or query_id not in self._response_bank:
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return StrategyResult(
                prediction="",
                model_used="none",
                latency_ms=elapsed_ms,
                error=f"Query ID '{query_id}' not in response bank",
            )

        model_responses = self._response_bank[query_id]

        # Select model
        if model and model in model_responses:
            selected_model = model
        elif self._model_selector:
            selected_model = self._model_selector(query)
        elif self._default_model and self._default_model in model_responses:
            selected_model = self._default_model
        elif model_responses:
            selected_model = next(iter(model_responses))
        else:
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return StrategyResult(
                prediction="",
                model_used="none",
                latency_ms=elapsed_ms,
                error="No model responses available",
            )

        if selected_model not in model_responses:
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return StrategyResult(
                prediction="",
                model_used=selected_model,
                latency_ms=elapsed_ms,
                error=f"Model '{selected_model}' not in response bank for query",
            )

        prediction = model_responses[selected_model]
        elapsed_ms = (time.perf_counter() - start_time) * 1000

        return StrategyResult(
            prediction=prediction,
            model_used=selected_model,
            routing=RoutingDecision(
                tier="cached",
                selected_model=selected_model,
            ),
            latency_ms=elapsed_ms,
            cost_usd=0.0,
            tokens_in=0,
            tokens_out=0,
            metadata={
                "mode": "cached",
                "query_id": query_id,
                "available_models": list(model_responses.keys()),
            },
        )


__all__ = [
    "OracleEntry",
    "OracleStrategy",
    "CachedResponseStrategy",
]
