"""SWE-bench optimized routing for eb1-swe-preview.

This module provides a rule-based router that selects models based on task
complexity classification. The routing decisions are informed by empirical
SWE-bench benchmark data (January 2026).

Key findings that inform routing:
- Gemini 3 Flash: 78% on SWE-bench Verified (beats Haiku's 73.3%)
- Claude Sonnet 4.5: 77.2% (strong mid-tier)
- GPT-5.2: 75-80% (good reasoning)
- Claude Opus 4.5: 80.9% (SOTA)

Three variants are provided:
- eb1-swe-preview: Balanced cost/performance (default)
- eb1-swe-fast-preview: Optimized for latency/cost (Flash-heavy)
- eb1-swe-pro-preview: Optimized for accuracy (ensembles for complex tasks)

See docs/research/swe_bench_routing_research.md for full analysis.
"""

from __future__ import annotations

import logging
import time
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import Any, Literal, TypeAlias

from ember.models.schemas import ChatResponse, UsageStats
from .model_pool import (
    CLAUDE_OPUS,
    GEMINI_FLASH,
    GPT_5_2,
    GPT_5_2_CODEX,
    GPT_5_2_PRO,
)

_LOG = logging.getLogger("ember.takumi.eb1_swe")

# Type alias for variant selection
SweVariant: TypeAlias = Literal["balanced", "fast", "pro"]

JSONPrimitive: TypeAlias = str | int | float | bool | None
JSON: TypeAlias = JSONPrimitive | Mapping[str, "JSON"] | Sequence["JSON"]


# Priority-ordered fallback tuples: first healthy provider wins.
COMPLEXITY_MODEL_MAP: Mapping[str, tuple[str, ...]] = {
    "trivial": (GEMINI_FLASH,),
    "standard": (GPT_5_2, GEMINI_FLASH),
    "complex": (CLAUDE_OPUS,),
    "architectural": (CLAUDE_OPUS,),
}

# Fast variant: Flash-heavy, trades ~3-5% accuracy for latency/cost.
FAST_COMPLEXITY_MODEL_MAP: Mapping[str, tuple[str, ...]] = {
    "trivial": (GEMINI_FLASH,),
    "standard": (GEMINI_FLASH,),
    "complex": (GPT_5_2, CLAUDE_OPUS),
    "architectural": (GPT_5_2, CLAUDE_OPUS),
}

# Pro variant: Opus floor, ensembles for complex/architectural.
PRO_COMPLEXITY_MODEL_MAP: Mapping[str, tuple[str, ...]] = {
    "trivial": (CLAUDE_OPUS,),
    "standard": (CLAUDE_OPUS,),
    "complex": ("ensemble",),
    "architectural": ("ensemble",),
}

PRO_ENSEMBLE_CANDIDATES: tuple[str, ...] = (
    CLAUDE_OPUS,
    GPT_5_2,
    GPT_5_2_CODEX,
)
PRO_ENSEMBLE_JUDGE = CLAUDE_OPUS

# Per-model parameters. Provider layer translates reasoning effort to native API params.
MODEL_PARAMS: Mapping[str, Mapping[str, JSON]] = {
    CLAUDE_OPUS: {
        "reasoning": {"effort": "high"},
        "max_tokens": 8192,
    },
    "claude-sonnet-4-5-20250514": {
        "reasoning": {"effort": "high"},
        "max_tokens": 8192,
    },
    GPT_5_2: {
        "reasoning": {"effort": "high"},
    },
    GPT_5_2_CODEX: {
        "reasoning": {"effort": "xhigh"},
    },
    GEMINI_FLASH: {},
}

# Temperature settings by complexity
COMPLEXITY_TEMPERATURE: Mapping[str, float] = {
    "trivial": 0.3,
    "standard": 0.2,
    "complex": 0.1,
    "architectural": 0.1,
}

# Fast variant uses slightly higher temperatures (speed over precision)
FAST_COMPLEXITY_TEMPERATURE: Mapping[str, float] = {
    "trivial": 0.4,
    "standard": 0.3,
    "complex": 0.2,
    "architectural": 0.2,
}

# Pro variant uses lower temperatures for consistency
PRO_COMPLEXITY_TEMPERATURE: Mapping[str, float] = {
    "trivial": 0.2,
    "standard": 0.1,
    "complex": 0.1,
    "architectural": 0.1,
}

# Classification signals. Default to "standard" (80%+ of SWE-bench issues).
# Escalate only for clear multi-file, concurrency, or architectural patterns.

COMPLEX_SIGNALS: tuple[str, ...] = (
    "multiple files",
    "across files",
    "cross-file",
    "race condition",
    "deadlock",
    "thread safe",
    "thread-safe",
    "concurrency",
    "concurrent",
    "async/await",
    "memory leak",
    "memory usage",
    "performance regression",
    "security",
    "vulnerability",
    "injection",
    "xss",
    "csrf",
)

ARCHITECTURAL_SIGNALS: tuple[str, ...] = (
    "refactor",
    "new feature",
    "api change",
    "breaking change",
    "backwards compatibility",
    "deprecate",
    "migration",
)

TRIVIAL_SIGNALS: tuple[str, ...] = (
    "typo",
    "spelling",
    "docstring",
    "comment only",
    "whitespace",
)


@dataclass(frozen=True, slots=True)
class RoutingDecision:
    """Captures the routing decision for telemetry and debugging."""

    complexity: str
    model: str
    temperature: float
    signals_matched: tuple[str, ...]
    prompt_length: int
    rationale: str
    variant: str = "balanced"
    uses_ensemble: bool = False


@dataclass(frozen=True, slots=True)
class SweResult:
    """Result of an eb1-swe-preview invocation."""

    response: ChatResponse
    routing: RoutingDecision
    latency_ms: float

    @property
    def text(self) -> str:
        """Primary text output."""
        return self.response.text

    @property
    def usage(self) -> UsageStats | None:
        """Token usage statistics."""
        return self.response.usage


def classify_complexity(prompt: str) -> tuple[str, tuple[str, ...]]:
    """Classify task complexity from keyword signals.

    Args:
        prompt: The task description or issue text.

    Returns:
        Tuple of (complexity_level, matched_signals).
    """
    prompt_lower = prompt.lower()

    for signal in COMPLEX_SIGNALS:
        if signal in prompt_lower:
            return "complex", (f"complex:{signal}",)

    for signal in ARCHITECTURAL_SIGNALS:
        if signal in prompt_lower:
            return "architectural", (f"arch:{signal}",)

    for signal in TRIVIAL_SIGNALS:
        if signal in prompt_lower:
            return "trivial", (f"trivial:{signal}",)

    return "standard", ("default",)


def _get_model_params(model: str) -> Mapping[str, JSON]:
    """Get model-specific parameters."""
    return MODEL_PARAMS.get(model, {})


def _get_model_map(variant: SweVariant) -> Mapping[str, tuple[str, ...]]:
    """Get the complexity-to-model map for a given variant."""
    if variant == "fast":
        return FAST_COMPLEXITY_MODEL_MAP
    if variant == "pro":
        return PRO_COMPLEXITY_MODEL_MAP
    return COMPLEXITY_MODEL_MAP


def _resolve_first_available(candidates: tuple[str, ...]) -> str:
    """Return the first model whose provider is not unhealthy.

    Falls back to the first candidate (highest quality) when all are
    unhealthy or health cannot be determined.
    """
    if len(candidates) == 1:
        return candidates[0]

    for model_id in candidates:
        if not _is_provider_unhealthy(model_id):
            return model_id

    # All unhealthy -- return the preferred model and let the error surface.
    _LOG.warning(
        "All fallback models unhealthy, using preferred: %s",
        candidates[0],
    )
    return candidates[0]


# Minimum failure count before a provider is considered unhealthy.
_HEALTH_FAILURE_THRESHOLD = 3


def _is_provider_unhealthy(model_id: str) -> bool:
    """Check if the provider for a model has been consistently failing.

    Uses the gateway's MetricsRegistry when available.  Falls back to
    ``False`` (assume healthy) if metrics are not wired.
    """
    try:
        from keiro.gateway.infra.metrics import METRICS
    except ImportError:
        return False

    provider = _model_id_to_provider(model_id)
    if provider is None:
        return False

    stats = METRICS.for_provider(provider)
    total = stats.successes + stats.failures
    return total >= _HEALTH_FAILURE_THRESHOLD and stats.failures == total


def _model_id_to_provider(model_id: str) -> str | None:
    """Map a model ID to its provider name for health checks."""
    lower = model_id.lower()
    if "gpt" in lower or "o1" in lower or "o3" in lower or "codex" in lower:
        return "openai"
    if "claude" in lower or "opus" in lower or "sonnet" in lower or "haiku" in lower:
        return "anthropic"
    if "gemini" in lower or "flash" in lower:
        return "google"
    return None


def _get_temperature_map(variant: SweVariant) -> Mapping[str, float]:
    """Get the complexity-to-temperature map for a given variant."""
    if variant == "fast":
        return FAST_COMPLEXITY_TEMPERATURE
    if variant == "pro":
        return PRO_COMPLEXITY_TEMPERATURE
    return COMPLEXITY_TEMPERATURE


def _run_swe_single_model(
    prompt: str,
    *,
    system: str | None,
    model: str,
    temperature: float,
    additional_params: Mapping[str, Any] | None,
) -> ChatResponse:
    """Invoke a single model for SWE task completion."""
    from ember.api import models

    params: dict[str, Any] = {
        "temperature": temperature,
        **_get_model_params(model),
    }
    if additional_params:
        params.update(additional_params)
    if system:
        params["system"] = system

    response = models.response(model, prompt, **params)
    return response.chat_response


def _run_swe_ensemble(
    prompt: str,
    *,
    system: str | None,
    temperature: float,
    additional_params: Mapping[str, Any] | None,
) -> ChatResponse:
    """Run ensemble for complex/architectural SWE tasks (pro variant).

    Uses the eb1 ensemble pattern: multiple candidates + judge synthesis.
    """
    from .eb1 import CandidatePlan, Eb1Plan, JudgePlan, run_plan

    # Build candidate plans with SWE-optimized parameters
    candidates = tuple(
        CandidatePlan(model=model, params=dict(_get_model_params(model)))
        for model in PRO_ENSEMBLE_CANDIDATES
    )

    # Build judge plan
    judge = JudgePlan(
        model=PRO_ENSEMBLE_JUDGE,
        params=dict(_get_model_params(PRO_ENSEMBLE_JUDGE)),
    )

    plan = Eb1Plan(candidates=candidates, judge=judge)

    shared_options: dict[str, Any] = {"temperature": temperature}
    if additional_params:
        shared_options.update(additional_params)

    return run_plan(
        plan,
        prompt,
        system=system,
        shared_options=shared_options,
    )


def run_eb1_swe_preview(
    prompt: str,
    *,
    system: str | None = None,
    complexity_override: str | None = None,
    model_override: str | None = None,
    additional_params: Mapping[str, Any] | None = None,
    variant: SweVariant = "balanced",
) -> SweResult:
    """Route SWE-bench tasks to optimal models based on complexity.

    This function classifies task complexity and routes to the appropriate
    model based on empirical SWE-bench benchmark data.

    Args:
        prompt: The task description (issue text, bug report, etc.).
        system: Optional system instructions to prepend.
        complexity_override: Force a specific complexity level
            ("trivial", "standard", "complex", "architectural").
        model_override: Force a specific model (bypasses routing).
        additional_params: Extra parameters to pass to the model.
        variant: Routing variant ("balanced", "fast", "pro").
            - balanced: Default, good cost/performance tradeoff
            - fast: Optimized for latency/cost (Flash-heavy)
            - pro: Optimized for accuracy (ensembles for complex tasks)

    Returns:
        SweResult containing the response, routing decision, and timing.

    Examples:
        >>> result = run_eb1_swe_preview("Fix typo in error message")
        >>> result.routing.complexity
        'trivial'
        >>> result.routing.model
        'gemini-3-flash-preview'

        >>> result = run_eb1_swe_preview("Refactor auth module", variant="pro")
        >>> result.routing.uses_ensemble
        True
    """
    start_time = time.perf_counter()

    # Get variant-specific maps
    model_map = _get_model_map(variant)
    temp_map = _get_temperature_map(variant)

    # Classify complexity
    if complexity_override:
        complexity = complexity_override
        signals = (f"override:{complexity_override}",)
    else:
        complexity, signals = classify_complexity(prompt)

    # Select model from fallback chain
    if model_override:
        model = model_override
        uses_ensemble = False
    else:
        candidates = model_map.get(
            complexity,
            model_map.get("standard", ("claude-sonnet-4-5-20250514",)),
        )
        model = _resolve_first_available(candidates)
        uses_ensemble = model == "ensemble"

    temperature = temp_map.get(complexity, 0.2)

    # Build routing decision for telemetry
    routing = RoutingDecision(
        complexity=complexity,
        model=model if not uses_ensemble else "ensemble",
        temperature=temperature,
        signals_matched=signals,
        prompt_length=len(prompt),
        rationale=f"Variant={variant}, Complexity={complexity}, Signals={signals}",
        variant=variant,
        uses_ensemble=uses_ensemble,
    )

    _LOG.info(
        "eb1_swe_routing",
        extra={
            "variant": variant,
            "complexity": complexity,
            "model": model,
            "uses_ensemble": uses_ensemble,
            "signals": signals,
            "prompt_length": len(prompt),
        },
    )

    # Execute based on routing decision
    if uses_ensemble:
        response = _run_swe_ensemble(
            prompt,
            system=system,
            temperature=temperature,
            additional_params=additional_params,
        )
    else:
        response = _run_swe_single_model(
            prompt,
            system=system,
            model=model,
            temperature=temperature,
            additional_params=additional_params,
        )

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

    return SweResult(
        response=response,
        routing=routing,
        latency_ms=elapsed_ms,
    )


def run_eb1_swe_fast_preview(
    prompt: str,
    *,
    system: str | None = None,
    complexity_override: str | None = None,
    model_override: str | None = None,
    additional_params: Mapping[str, Any] | None = None,
) -> SweResult:
    """Run SWE task with fast variant (optimized for latency/cost).

    This is a convenience wrapper around run_eb1_swe_preview(variant="fast").
    Uses Gemini Flash heavily for speed and cost efficiency.

    Args:
        prompt: The task description.
        system: Optional system instructions.
        complexity_override: Force a specific complexity level.
        model_override: Force a specific model.
        additional_params: Extra parameters to pass to the model.

    Returns:
        SweResult containing the response, routing decision, and timing.
    """
    return run_eb1_swe_preview(
        prompt,
        system=system,
        complexity_override=complexity_override,
        model_override=model_override,
        additional_params=additional_params,
        variant="fast",
    )


def run_eb1_swe_pro_preview(
    prompt: str,
    *,
    system: str | None = None,
    complexity_override: str | None = None,
    model_override: str | None = None,
    additional_params: Mapping[str, Any] | None = None,
) -> SweResult:
    """Run SWE task with pro variant (optimized for accuracy).

    This is a convenience wrapper around run_eb1_swe_preview(variant="pro").
    Uses ensembles for complex/architectural tasks for maximum accuracy.

    Args:
        prompt: The task description.
        system: Optional system instructions.
        complexity_override: Force a specific complexity level.
        model_override: Force a specific model (bypasses ensemble).
        additional_params: Extra parameters to pass to the model.

    Returns:
        SweResult containing the response, routing decision, and timing.
    """
    return run_eb1_swe_preview(
        prompt,
        system=system,
        complexity_override=complexity_override,
        model_override=model_override,
        additional_params=additional_params,
        variant="pro",
    )


def get_routing_decision(
    prompt: str,
    *,
    variant: SweVariant = "balanced",
) -> RoutingDecision:
    """Get the routing decision without invoking the model.

    Useful for debugging and testing classification logic.

    Args:
        prompt: The task description.
        variant: Routing variant ("balanced", "fast", "pro").

    Returns:
        RoutingDecision showing what model would be selected.
    """
    model_map = _get_model_map(variant)
    temp_map = _get_temperature_map(variant)

    complexity, signals = classify_complexity(prompt)
    candidates = model_map.get(
        complexity,
        model_map.get("standard", ("claude-sonnet-4-5-20250514",)),
    )
    model = _resolve_first_available(candidates)
    temperature = temp_map.get(complexity, 0.2)
    uses_ensemble = model == "ensemble"

    return RoutingDecision(
        complexity=complexity,
        model=model,
        temperature=temperature,
        signals_matched=signals,
        prompt_length=len(prompt),
        rationale=f"Variant={variant}, Complexity={complexity}, Signals={signals}",
        variant=variant,
        uses_ensemble=uses_ensemble,
    )
