from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Callable, Mapping

from ember._internal.exceptions import ProviderAPIError
from ..providers.capabilities import build_model_caps, supports
from ..providers.features import Feature, feature_names, required_features
from ember.models.providers.utils import ResponsesRequest
from keiro.ensemble.eb1 import CandidatePlan, Eb1Plan, JudgePlan


class Mode(Enum):
    """Execution policy for the EB1 ensemble."""

    SINGLE = "single"
    ORCHESTRATE = "orchestrate"
    JUDGE_TOOLS = "judge_tools"


@dataclass(frozen=True, slots=True)
class ExecutablePlan:
    """Filtered EB1 plan paired with execution mode metadata."""

    candidates: tuple[CandidatePlan, ...]
    judge: JudgePlan
    mode: Mode
    reason: str
    features: Feature

    def to_plan(self) -> Eb1Plan:
        return Eb1Plan(candidates=self.candidates, judge=self.judge)


def _has_credentials(provider: str, checker: Callable[[str], bool]) -> bool:
    if provider in {"takumi", "unknown"}:
        return True
    return checker(provider)


def _select_judge(
    plan: Eb1Plan,
    candidates: tuple[CandidatePlan, ...],
    features: Feature,
    has_credentials: Callable[[str], bool],
    mode_name: str,
) -> JudgePlan:
    """Return a judge compatible with the filtered candidate set."""

    judge_caps = build_model_caps(plan.judge.model)
    judge_mode = f"{mode_name}.judge"
    if _has_credentials(judge_caps.provider, has_credentials) and supports(
        judge_caps, features, mode=judge_mode
    ):
        return plan.judge

    first_candidate = candidates[0]
    preferred_provider = build_model_caps(first_candidate.model).provider
    replacement = next(
        (
            candidate
            for candidate in candidates
            if build_model_caps(candidate.model).provider == preferred_provider
        ),
        first_candidate,
    )
    return JudgePlan(
        model=replacement.model,
        params=plan.judge.params,
        prompt_template=plan.judge.prompt_template,
        system_template=plan.judge.system_template,
    )


def _features_from_plan(plan: Eb1Plan) -> Feature:
    """Return feature requirements implied by default plan parameters."""

    implied = Feature.NONE
    for component in (*plan.candidates, plan.judge):
        params: Mapping[str, object] = component.params
        if "tools" in params:
            implied |= Feature.TOOLS
        if "tool_choice" in params:
            implied |= Feature.TOOL_CHOICE
        if "response_format" in params:
            implied |= Feature.RESPONSE_FORMAT
    return implied


def compile_plan(
    plan: Eb1Plan,
    request: ResponsesRequest,
    *,
    has_credentials: Callable[[str], bool],
    mode_name: str = "responses",
) -> ExecutablePlan:
    """Compile an EB1 plan into an executable subset with a selected mode."""

    features = required_features(request) | _features_from_plan(plan)

    surviving: list[CandidatePlan] = []
    missing_providers: set[str] = set()
    unsupported_models: list[str] = []

    for candidate in plan.candidates:
        caps = build_model_caps(candidate.model)
        if not _has_credentials(caps.provider, has_credentials):
            missing_providers.add(caps.provider)
            continue
        if not supports(caps, features, mode=mode_name):
            unsupported_models.append(candidate.model)
            continue
        surviving.append(candidate)

    if not surviving:
        context: dict[str, object] = {}
        labels = feature_names(features)
        if labels:
            context["features"] = labels
        if missing_providers:
            context["missing_providers"] = sorted(missing_providers)
        if unsupported_models:
            context["unsupported_models"] = unsupported_models
        raise ProviderAPIError(
            "No EB1 candidate models support the requested configuration.",
            context=context or None,
        )

    candidates_tuple = tuple(surviving)
    judge = _select_judge(
        plan,
        candidates_tuple,
        features,
        has_credentials,
        mode_name,
    )

    if len(candidates_tuple) == 1:
        mode = Mode.SINGLE
        reason = "single surviving candidate"
    elif features & (Feature.TOOLS | Feature.TOOL_CHOICE):
        mode = Mode.JUDGE_TOOLS
        reason = "tools requested; using judge-only tools policy"
    else:
        mode = Mode.ORCHESTRATE
        reason = "multiple candidates; no tools"

    return ExecutablePlan(
        candidates=candidates_tuple,
        judge=judge,
        mode=mode,
        reason=reason,
        features=features,
    )
