"""Provider and model capability helpers used by the Takumi gateway."""

from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache
from typing import TYPE_CHECKING, Iterable, Mapping, MutableMapping

from ..provider_registry import get_capability_map
from ..providers.capabilities_defaults import (
    PROVIDER_DEFAULTS,
    resolve_family_capabilities,
)
from ..providers.features import REQUIREMENTS, Cap, Feature
from ..settings import GatewaySettings, get_gateway_settings
from ember.models.capabilities import PROVIDER_ALIASES, build_provider_capabilities
from ember.models.catalog import get_model_info
from ember.models.providers import resolve_model_id
from ember.types import ProviderCapabilities

if TYPE_CHECKING:  # pragma: no cover
    from ..providers.adapters import ProviderAdapter


@dataclass(frozen=True, slots=True)
class ModelCaps:
    """Describes the capabilities Ember expects for a specific model."""

    model: str
    provider: str
    caps: Cap = Cap.NONE


_CAPABILITY_MAP: dict[str, Cap] = {
    "tools": Cap.TOOLS,
    "json": Cap.JSON,
    "structured_output": Cap.STRUCTURED_OUTPUT,
    "reasoning": Cap.REASONING,
    "thinking": Cap.REASONING,
}


def caps_from_strings(values: Iterable[str] | None) -> Cap:
    """Translate catalog capability strings into the internal bitset."""

    if not values:
        return Cap.NONE

    caps = Cap.NONE
    for value in values:
        mapped = _CAPABILITY_MAP.get(value.lower())
        if mapped is not None:
            caps |= mapped
    return caps


RESPONSES_MODE = "responses"
RESPONSES_JUDGE_MODE = f"{RESPONSES_MODE}.judge"
RESPONSES_WORKER_MODE = f"{RESPONSES_MODE}.worker"


PROVIDER_BLOCKS: dict[str, dict[str, Feature]] = {
    # Anthropic Responses adapter supports tools and reasoning (via thinking/effort
    # translation), but not OpenAI-style response_format (structured JSON schema).
    "anthropic": {
        RESPONSES_MODE: Feature.RESPONSE_FORMAT,
        RESPONSES_JUDGE_MODE: Feature.RESPONSE_FORMAT,
    },
}


def build_model_caps(model_id: str) -> ModelCaps:
    """Construct the capability record for ``model_id`` using hybrid resolution."""

    provider, canonical_model = resolve_model_id(model_id)

    # eb1 is a router, not a model. It transparently supports whatever its
    # underlying candidates support, so claim full capabilities here.
    if canonical_model.startswith("eb1"):
        return ModelCaps(
            model=canonical_model,
            provider=provider,
            caps=Cap.TOOLS | Cap.JSON | Cap.STRUCTURED_OUTPUT | Cap.REASONING,
        )

    return _build_model_caps(provider, canonical_model)


@lru_cache(maxsize=None)
def _build_model_caps(provider: str, canonical_model: str) -> ModelCaps:
    """Resolve capabilities for the canonical model/provider pair."""

    resolved_model = canonical_model
    catalog_provider = provider

    try:
        info = get_model_info(
            canonical_model,
            include_dynamic=False,
            discovery_mode="bootstrap",
        )
    except KeyError:
        info = None
    else:
        resolved_model = info.id
        info_provider = info.provider
        if provider == "unknown":
            catalog_provider = info_provider
        catalog_caps = caps_from_strings(info.capabilities)
        if catalog_caps is not Cap.NONE:
            return ModelCaps(
                model=resolved_model,
                provider=catalog_provider,
                caps=catalog_caps,
            )

    family_caps = resolve_family_capabilities(resolved_model, catalog_provider)
    if family_caps is not Cap.NONE:
        return ModelCaps(model=resolved_model, provider=catalog_provider, caps=family_caps)

    return ModelCaps(
        model=resolved_model,
        provider=catalog_provider,
        caps=PROVIDER_DEFAULTS.get(catalog_provider, Cap.NONE),
    )


def supports(model: ModelCaps, features: Feature, *, mode: str = "responses") -> bool:
    """Return True when ``model`` can satisfy the requested feature bitset."""

    if features == Feature.NONE:
        return True

    blocked = PROVIDER_BLOCKS.get(model.provider, {}).get(mode, Feature.NONE)
    if features & blocked:
        return False

    for feature, requirements in REQUIREMENTS.items():
        if features & feature and model.caps & requirements == Cap.NONE:
            return False
    return True


class ProviderCapabilityRow(ProviderCapabilities, total=False):
    """Capability descriptor exposed by the gateway discovery endpoints."""

    supports_responses_api: bool


ProviderCapabilitiesMap = dict[str, ProviderCapabilityRow]


def list_capabilities(
    adapters: Mapping[str, "ProviderAdapter"] | None = None,
) -> ProviderCapabilitiesMap:
    """Return the provider capability matrix used by Takumi endpoints.

    Args:
        adapters: Optional mapping of provider name → adapter instance. When omitted,
            the default adapters are instantiated (legacy mode only).
    """

    settings = get_gateway_settings()
    if settings.features.capability_registry:
        return _capabilities_from_registry(settings)
    return _capabilities_from_canonical(settings, _resolve_adapters(adapters))


def _capabilities_from_registry(settings: GatewaySettings) -> ProviderCapabilitiesMap:
    registry_map = {_k: _normalize_capabilities(_v) for _k, _v in get_capability_map().items()}
    expanded = _apply_aliases(registry_map)
    _apply_google_streaming_override(expanded, settings.providers.google_streaming)
    return expanded


def _capabilities_from_canonical(
    settings: GatewaySettings,
    adapters: Mapping[str, "ProviderAdapter"],
) -> ProviderCapabilitiesMap:
    base = {_k: _normalize_capabilities(_v) for _k, _v in build_provider_capabilities().items()}
    for data in base.values():
        data["supports_responses_api"] = data.get("supports_responses_api", True)

    expanded = _apply_aliases(base)
    for provider, adapter in adapters.items():
        row = expanded.setdefault(provider, _empty_capability_row())
        row["supports_streaming"] = bool(getattr(adapter, "supports_streaming", False))
        row.setdefault("supports_responses_api", True)

    _apply_google_streaming_override(expanded, settings.providers.google_streaming)
    return expanded


def _apply_aliases(capabilities: ProviderCapabilitiesMap) -> ProviderCapabilitiesMap:
    expanded: ProviderCapabilitiesMap = {name: dict(row) for name, row in capabilities.items()}
    for alias, canonical in PROVIDER_ALIASES.items():
        canonical_row = expanded.get(canonical)
        if canonical_row is not None:
            expanded[alias] = dict(canonical_row)
    return expanded


def _apply_google_streaming_override(
    capabilities: MutableMapping[str, ProviderCapabilityRow],
    override: bool | None,
) -> None:
    if override is None:
        return
    target = capabilities.get("google")
    if target is None:
        target = capabilities.setdefault("google", _empty_capability_row())
    target["supports_streaming"] = override
    for alias, canonical in PROVIDER_ALIASES.items():
        if canonical == "google" and alias in capabilities:
            capabilities[alias]["supports_streaming"] = override


def _normalize_capabilities(payload: Mapping[str, object]) -> ProviderCapabilityRow:
    return ProviderCapabilityRow(
        supports_streaming=bool(payload.get("supports_streaming", False)),
        json_strict=bool(payload.get("json_strict", False)),
        tools=bool(payload.get("tools", False)),
        vision=bool(payload.get("vision", False)),
        supports_responses_api=bool(payload.get("supports_responses_api", True)),
    )


def _empty_capability_row() -> ProviderCapabilityRow:
    return ProviderCapabilityRow(
        supports_streaming=False,
        json_strict=False,
        tools=False,
        vision=False,
    )


def _resolve_adapters(
    adapters: Mapping[str, "ProviderAdapter"] | None,
) -> Mapping[str, "ProviderAdapter"]:
    if adapters is not None:
        return adapters
    from ..providers.registry import default_adapters

    return default_adapters()
