"""Model-family reasoning configuration and translation.

Provides a centralized mapping of reasoning effort levels to
provider-specific parameters. Different model families have
different reasoning APIs:

- OpenAI (o1/o3, GPT-5): Uses ``output_config.effort`` natively
- Anthropic (Claude): Uses ``output_config.effort`` natively (SDK >=0.75.0)
- Bedrock (Converse API): Uses ``thinking.budget_tokens`` with synthetic mapping
- Google (Gemini): Not yet supported

This module translates between the unified effort-based API
and provider-specific implementations.
"""

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Final, Literal

from ember._internal.exceptions import ProviderConfigError

# Canonical effort levels after normalization by _normalize_effort().
ReasoningEffort = Literal["low", "medium", "high", "xhigh"]

# Maximum reasoning budget in tokens (128K) for providers that support
# explicit budget_tokens (Anthropic, Bedrock).
_MAX_REASONING_BUDGET: Final[int] = 131072

# Minimum explicit budget_tokens for native-effort providers that also
# accept direct budget overrides (e.g., Anthropic).
_NATIVE_MIN_BUDGET: Final[int] = 1024


@dataclass(frozen=True, slots=True)
class ProviderReasoningPolicy:
    """Model-family-specific reasoning configuration.

    Attributes:
        supports_reasoning: Whether this provider supports reasoning/thinking.
        effort_to_tokens: Mapping of effort levels to token budgets.
            Only populated for providers that need synthetic mapping (e.g.,
            Bedrock). None for providers that support native effort.
        max_budget: Maximum allowed budget_tokens value.
        default_effort: Default effort level when not specified.
        uses_native_effort: When True, the provider accepts effort strings
            directly via its API (e.g., ``output_config: {effort: "high"}``).
            When False, effort is translated to ``budget_tokens``.
    """

    supports_reasoning: bool
    effort_to_tokens: Mapping[ReasoningEffort, int] | None
    max_budget: int
    default_effort: ReasoningEffort
    uses_native_effort: bool = False


# Model family configurations.
#
# Native effort providers (Anthropic, OpenAI) pass effort strings directly
# to the API — no synthetic budget mapping. The API handles token allocation
# internally based on task complexity.
#
# Synthetic mapping providers (Bedrock) still need effort-to-tokens translation
# because their API only accepts explicit budget_tokens.
REASONING_CONFIGS: Final[dict[str, ProviderReasoningPolicy]] = {
    "anthropic": ProviderReasoningPolicy(
        supports_reasoning=True,
        effort_to_tokens=None,  # Native effort — no synthetic mapping
        max_budget=_MAX_REASONING_BUDGET,  # Used for explicit budget_tokens validation
        default_effort="medium",
        uses_native_effort=True,
    ),
    "openai": ProviderReasoningPolicy(
        supports_reasoning=True,
        effort_to_tokens=None,
        max_budget=0,
        default_effort="medium",
        uses_native_effort=True,
    ),
    "google": ProviderReasoningPolicy(
        supports_reasoning=False,
        effort_to_tokens=None,
        max_budget=0,
        default_effort="medium",
    ),
    "bedrock": ProviderReasoningPolicy(
        supports_reasoning=True,
        effort_to_tokens={
            "low": 1024,
            "medium": 4096,
            "high": 16384,
            "xhigh": _MAX_REASONING_BUDGET,
        },
        max_budget=_MAX_REASONING_BUDGET,
        default_effort="medium",
    ),
    "vertex": ProviderReasoningPolicy(
        supports_reasoning=False,
        effort_to_tokens=None,
        max_budget=0,
        default_effort="medium",
    ),
}


def get_reasoning_config(provider: str) -> ProviderReasoningPolicy | None:
    """Get reasoning configuration for a provider.

    Args:
        provider: Provider name (e.g., "openai", "anthropic").

    Returns:
        ReasoningConfig for the provider, or None if unknown.
    """
    return REASONING_CONFIGS.get(provider)


def translate_reasoning_for_provider(
    reasoning: Mapping[str, Any] | None,
    provider: str,
    model_id: str | None = None,
) -> dict[str, Any] | None:
    """Translate reasoning config for target provider.

    Handles two input formats:
    1. Effort-based: ``{"effort": "high"}`` — translated per provider
    2. Budget-based: ``{"budget_tokens": 8192}`` — explicit thinking budget

    For native-effort providers (Anthropic, OpenAI), effort returns
    ``{"effort": "<level>"}`` and the internal ``xhigh`` level maps to
    ``"max"`` in the API. For synthetic-mapping providers (Bedrock),
    effort returns ``{"type": "enabled", "budget_tokens": N}``.

    Explicit ``budget_tokens`` always returns
    ``{"type": "enabled", "budget_tokens": N}`` regardless of provider.

    Args:
        reasoning: Input reasoning configuration.
        provider: Target provider name.
        model_id: Optional model ID for model-specific adjustments.

    Returns:
        Provider-appropriate reasoning dict, or None if reasoning
        is not requested.

    Raises:
        ProviderConfigError: If the provider is unknown, does not support
            reasoning, or the reasoning payload is invalid.
    """
    if reasoning is None:
        return None

    config = REASONING_CONFIGS.get(provider)
    if config is None:
        raise ProviderConfigError(
            f"Unknown provider '{provider}' for reasoning config"
        )
    if not config.supports_reasoning:
        raise ProviderConfigError(
            f"Provider '{provider}' does not support reasoning"
        )

    # Handle explicit budget_tokens (provider-specific override)
    budget_tokens = reasoning.get("budget_tokens")
    if budget_tokens is not None:
        return _translate_budget(budget_tokens, config, provider)

    # Translate effort level
    effort = _normalize_effort(reasoning.get("effort"), config)

    allowed_efforts = (
        set(config.effort_to_tokens)
        if config.effort_to_tokens
        else {"low", "medium", "high", "xhigh"}
    )
    if effort not in allowed_efforts:
        raise ProviderConfigError(
            f"Reasoning effort '{effort}' is not supported "
            f"for provider '{provider}'"
        )

    # Native effort providers pass effort string directly to the API
    if config.uses_native_effort:
        native_effort = "max" if effort == "xhigh" else effort
        return {"effort": native_effort}

    # Synthetic mapping for non-native providers (Bedrock)
    if config.effort_to_tokens:
        effort_tokens = config.effort_to_tokens.get(effort)
        if effort_tokens is None:
            raise ProviderConfigError(
                f"Reasoning effort '{effort}' is not supported "
                f"for provider '{provider}'"
            )
        return {"type": "enabled", "budget_tokens": effort_tokens}

    # Fallback: pass effort directly
    return {"effort": effort}


def _translate_budget(
    budget_tokens: Any,
    config: ProviderReasoningPolicy,
    provider: str,
) -> dict[str, Any]:
    """Translate explicit budget_tokens for a provider.

    Args:
        budget_tokens: Raw budget_tokens value from reasoning config.
        config: Provider reasoning policy.
        provider: Provider name for error messages.

    Returns:
        Thinking payload with budget_tokens.

    Raises:
        ProviderConfigError: If the provider does not accept budget_tokens
            or the value is out of range.
    """
    # Native-effort providers with max_budget > 0 accept explicit budgets
    # as an override (e.g., Anthropic). Providers with max_budget == 0
    # (e.g., OpenAI) do not accept budget_tokens at all.
    if config.uses_native_effort and config.max_budget == 0:
        raise ProviderConfigError(
            f"Provider '{provider}' does not accept budget_tokens; "
            "use 'effort' instead"
        )
    if not config.uses_native_effort and config.effort_to_tokens is None:
        raise ProviderConfigError(
            f"Provider '{provider}' does not accept budget_tokens; "
            "use 'effort' instead"
        )

    try:
        budget = int(budget_tokens)
    except (TypeError, ValueError) as exc:
        raise ProviderConfigError(
            "Reasoning budget_tokens must be an integer "
            f"for provider '{provider}'"
        ) from exc

    # Determine valid range
    if config.uses_native_effort:
        min_budget = _NATIVE_MIN_BUDGET
    elif config.effort_to_tokens is not None:
        min_budget = min(config.effort_to_tokens.values())
    else:
        min_budget = _NATIVE_MIN_BUDGET

    if budget < min_budget or budget > config.max_budget:
        raise ProviderConfigError(
            f"Reasoning budget_tokens must be between "
            f"{min_budget} and {config.max_budget} for provider '{provider}'"
        )

    return {"type": "enabled", "budget_tokens": budget}


def _normalize_effort(
    raw_effort: Any | None,
    config: ProviderReasoningPolicy,
) -> ReasoningEffort:
    """Normalize an effort value to one of {low, medium, high, xhigh}.

    Args:
        raw_effort: Raw effort value from reasoning config, or None.
        config: Provider reasoning policy (supplies default).

    Returns:
        Normalized effort string.

    Raises:
        ProviderConfigError: If the effort value is unrecognized.
    """
    effort = raw_effort if raw_effort is not None else config.default_effort
    effort = str(effort).lower().strip()

    if effort in ("xhigh", "x-high", "maximum", "max"):
        return "xhigh"
    if effort in ("hi", "high"):
        return "high"
    if effort in ("med", "medium", "moderate"):
        return "medium"
    if effort in ("lo", "low", "minimal"):
        return "low"
    raise ProviderConfigError(
        f"Unrecognized reasoning effort '{effort}'; "
        "expected one of: low, medium, high, xhigh"
    )


def supports_reasoning(provider: str) -> bool:
    """Check if a provider supports reasoning/thinking."""
    config = REASONING_CONFIGS.get(provider)
    return config is not None and config.supports_reasoning


def get_default_effort(provider: str) -> str:
    """Get the default reasoning effort for a provider.

    Raises:
        ProviderConfigError: If the provider is unknown.
    """
    config = REASONING_CONFIGS.get(provider)
    if config is None:
        raise ProviderConfigError(
            f"Unknown provider '{provider}' for reasoning config"
        )
    return config.default_effort


ReasoningConfig = ProviderReasoningPolicy  # Backward-compat alias

__all__ = [
    "ProviderReasoningPolicy",
    "ReasoningConfig",
    "REASONING_CONFIGS",
    "get_default_effort",
    "get_reasoning_config",
    "supports_reasoning",
    "translate_reasoning_for_provider",
]
