from __future__ import annotations

from typing import Any

from app.ai_provider.config.models import ModelConfig
from app.ai_provider.core.registry import ModelRegistry
from app.ai_provider.routing.multimodal import MultiModalDecision
from app.ai_provider.types.errors import AIProviderError

_SELECTOR_CAPABILITY_MAP: dict[MultiModalDecision, str] = {
    MultiModalDecision.VISION: "vision",
    MultiModalDecision.THINKING: "thinking",
    MultiModalDecision.TOOL_USE: "tool_use",
    MultiModalDecision.STRUCTURED_OUTPUT: "structured_output",
    MultiModalDecision.PLAIN_TEXT: "chat",
}


class ModelSelector:
    """为给定上下文选择最佳模型。

    用法::

        selector = ModelSelector(registry)
        model = selector.select({"capabilities": {"vision", "thinking"}})
        model = selector.select({"preferred": "fast-model"})
        model = selector.select_from_decisions({MultiModalDecision.VISION})
    """

    def __init__(self, registry: ModelRegistry) -> None:
        self._registry = registry

    def select(self, context: dict[str, Any]) -> ModelConfig:
        """根据给定的上下文选择模型。

        上下文键:

        * ``capabilities`` (``set[str]``, 可选) — 必需的能力。
        * ``preferred`` (``str``, 可选) — 偏好的模型名或别名。
        * ``avoid`` (``set[str]``, 可选) — 要排除的模型名。

        如果给了偏好模型但不能满足所有必需的能力，此方法会回退到
        能满足要求的模型。
        """
        capabilities: set[str] = set(context.get("capabilities", set()))
        preferred: str | None = context.get("preferred")
        avoid: set[str] = set(context.get("avoid", set()))

        # If preferred is given, resolve it and check capabilities
        if preferred is not None:
            model = self._registry.resolve(preferred)
            if (
                not capabilities or self._model_satisfies_all(model, capabilities)
            ) and model.name not in avoid:
                return model

        # Find models that satisfy all required capabilities
        if capabilities:
            return self._select_by_capabilities(capabilities, avoid)

        # Fall back to default
        return self._registry.resolve()

    def select_from_decisions(
        self,
        decisions: set[MultiModalDecision],
    ) -> ModelConfig:
        """基于 :class:`MultiModalDecision` 值选择模型。

        将决策翻译为能力字符串并委托给 :meth:`select`。
        """
        capabilities: set[str] = set()
        for d in decisions:
            cap = _SELECTOR_CAPABILITY_MAP.get(d)
            if cap and cap != "chat":
                capabilities.add(cap)
        return self.select({"capabilities": capabilities})

    # ------------------------------------------------------------------

    def _select_by_capabilities(
        self,
        capabilities: set[str],
        avoid: set[str],
    ) -> ModelConfig:
        """找到支持所有必需能力的模型。"""
        # Narrow down by each capability (ModelConfig is unhashable, use lists)
        candidates: list[ModelConfig] = list(self._registry.list_all())

        for cap in capabilities:
            cap_models = list(self._registry.list_by_capability(cap))
            cap_names = {m.name for m in cap_models}
            candidates = [m for m in candidates if m.name in cap_names]

        # Remove avoided models
        candidates = [m for m in candidates if m.name not in avoid]

        if not candidates:
            raise AIProviderError(
                message=f"No active model supports all capabilities: {capabilities}",
                code="CAPABILITY_NOT_FOUND",
                status_code=404,
            )

        # Prefer higher-weight models (higher = better)
        return max(candidates, key=lambda m: m.weight)

    @staticmethod
    def _model_satisfies_all(model: ModelConfig, capabilities: set[str]) -> bool:
        caps = model.resolved_capabilities()
        return all(caps.get(cap, False) for cap in capabilities)