from __future__ import annotations
import os
import tomllib
from pathlib import Path
from typing import Mapping, TypeAlias
from pydantic import ValidationError
from returns.result import Failure, Result, Success

from anima.core.schemas import (
    AnimaGlobalConfig,
    CompactionConfig,
    InternalError,
    InvalidInputError,
)


_ToolOutputResolved: TypeAlias = dict[str, object]


DEFAULT_CONFIG_PATH = Path("~/.anima/config.toml").expanduser()
DEFAULT_MODEL_CONTEXT_LIMIT_TOKENS = 128_000
DEFAULT_TOOL_OUTPUT_PER_CALL_CAP_BYTES = 32_000
DEFAULT_TOOL_OUTPUT_PER_TURN_CAP_BYTES = 64_000
_RUNTIME_CONTEXT_LIMIT_ATTRS: tuple[str, ...] = (
    "context_window_tokens",
    "context_window",
    "max_input_tokens",
    "max_tokens",
)


_ConfigLoadResult: TypeAlias = Result[
    AnimaGlobalConfig, InvalidInputError | InternalError
]
_CONFIG_CACHE: dict[Path, tuple[int | None, _ConfigLoadResult]] = {}


# @invar:allow shell_result: scalar normalization helper for precedence resolver
def _as_non_empty_str(value: object) -> str | None:
    if isinstance(value, str):
        normalized = value.strip()
        if normalized:
            return normalized
    return None


# @invar:allow shell_result: normalization helper for profile-map lookup
def _as_mapping(value: object) -> Mapping[str, object] | None:
    if isinstance(value, Mapping):
        return value
    model_dump = getattr(value, "model_dump", None)
    if callable(model_dump):
        dumped = model_dump()
        if isinstance(dumped, Mapping):
            return dumped
    return None


# @invar:allow shell_result: deterministic model identifier parsing helper
def _split_model_identifier(value: object) -> tuple[str | None, str | None]:
    normalized = _as_non_empty_str(value)
    if normalized is None:
        return (None, None)
    if ":" not in normalized:
        return (None, normalized)
    provider_raw, model_raw = normalized.split(":", 1)
    provider = _as_non_empty_str(provider_raw)
    model = _as_non_empty_str(model_raw)
    return (provider, model)


# @invar:allow shell_result: deterministic map lookup by normalized key
# @shell_complexity: key-normalization fallback keeps profile lookup deterministic
def _resolve_profile_map_entry(
    profiles: object,
    *,
    candidates: tuple[str | None, ...],
) -> tuple[str | None, Mapping[str, object] | None]:
    mapping = _as_mapping(profiles)
    if mapping is None or not mapping:
        return (None, None)

    normalized_index: dict[str, str] = {}
    for key in mapping:
        normalized_key = _as_non_empty_str(key)
        if normalized_key is None:
            continue
        lowered = normalized_key.lower()
        normalized_index.setdefault(lowered, normalized_key)

    for candidate in candidates:
        normalized_candidate = _as_non_empty_str(candidate)
        if normalized_candidate is None:
            continue
        direct = mapping.get(normalized_candidate)
        if direct is not None:
            resolved = _as_mapping(direct)
            if resolved is not None:
                return (normalized_candidate, resolved)

        lowered = normalized_candidate.lower()
        resolved_key = normalized_index.get(lowered)
        if resolved_key is None:
            continue
        resolved_value = mapping.get(resolved_key)
        resolved = _as_mapping(resolved_value)
        if resolved is not None:
            return (resolved_key, resolved)

    return (None, None)


# @invar:allow shell_result: env parser returns partial override mapping
def _tool_output_env_overrides(env: Mapping[str, str]) -> dict[str, object]:
    """Collect ANIMA_TOOL_OUTPUT_* env overrides in config shape."""
    overrides: dict[str, object] = {}

    profile = _as_non_empty_str(env.get("ANIMA_TOOL_OUTPUT_PROFILE"))
    if profile is not None:
        overrides["profile"] = profile

    per_call = _as_positive_int(env.get("ANIMA_TOOL_OUTPUT_PER_CALL_CAP_BYTES"))
    if per_call is not None:
        overrides["per_call_cap_bytes"] = per_call

    per_turn = _as_positive_int(env.get("ANIMA_TOOL_OUTPUT_PER_TURN_CAP_BYTES"))
    if per_turn is not None:
        overrides["per_turn_cap_bytes"] = per_turn

    return overrides


# @invar:allow shell_result: merge helper returns normalized config payload
def _merge_tool_output_overrides(
    parsed: dict[str, object], env: Mapping[str, str]
) -> dict[str, object]:
    """Apply env tool-output overrides onto parsed TOML payload."""
    overrides = _tool_output_env_overrides(env)
    if not overrides:
        return parsed

    merged = dict(parsed)
    raw_tool_output = merged.get("tool_output")
    tool_output = dict(raw_tool_output) if isinstance(raw_tool_output, dict) else {}
    tool_output.update(overrides)
    merged["tool_output"] = tool_output
    return merged


# @invar:allow shell_result: merges process env with explicit call overrides
def _compose_env(env: Mapping[str, str] | None) -> Mapping[str, str]:
    if env is None:
        return os.environ
    composed = dict(os.environ)
    composed.update(env)
    return composed


def _apply_env_defaults(config: AnimaGlobalConfig) -> None:
    """Apply config [env] without overriding existing process variables."""
    for key, value in config.env.items():
        os.environ.setdefault(key, value)


# @invar:allow shell_result: lightweight path helper for shell wiring
def resolve_global_config_path(env: Mapping[str, str] | None = None) -> Path:
    source_env = env if env is not None else os.environ
    override = source_env.get("ANIMA_CONFIG")
    if override:
        return Path(override).expanduser()
    return DEFAULT_CONFIG_PATH


# @shell_complexity: Spec source help-groups.help-groups-gate-review-fix-shell-complexity-debt;
# config loader keeps explicit branch-specific error mapping (missing file, read, parse/validation)
def load_global_config(
    env: Mapping[str, str] | None = None,
) -> Result[AnimaGlobalConfig, InvalidInputError | InternalError]:
    source_env = _compose_env(env)
    config_path = resolve_global_config_path(source_env)
    # Spec source: vectl step ``deduplicate-load-global-config-calls``.
    # Cache by config path + file mtime to avoid duplicate parsing while
    # still refreshing when config is modified.
    try:
        mtime_ns = config_path.stat().st_mtime_ns
    except FileNotFoundError:
        mtime_ns = None

    cached = _CONFIG_CACHE.get(config_path)
    if cached is not None and cached[0] == mtime_ns:
        cached_result = cached[1]
        if isinstance(cached_result, Success):
            _apply_env_defaults(cached_result.unwrap())
        return cached_result

    if mtime_ns is None:
        config = AnimaGlobalConfig()
        overrides = _tool_output_env_overrides(source_env)
        if overrides:
            resolved = resolve_tool_output_precedence(
                cli=None,
                env=overrides,
                config=config.tool_output.model_dump(),
            )
            tool_output = config.tool_output.model_copy()
            profile = resolved.get("profile")
            if isinstance(profile, str):
                tool_output.profile = profile
            per_call_cap = resolved.get("per_call_cap_bytes")
            if isinstance(per_call_cap, int) and not isinstance(per_call_cap, bool):
                tool_output.per_call_cap_bytes = per_call_cap
            per_turn_cap = resolved.get("per_turn_cap_bytes")
            if isinstance(per_turn_cap, int) and not isinstance(per_turn_cap, bool):
                tool_output.per_turn_cap_bytes = per_turn_cap
            config.tool_output = tool_output

        _apply_env_defaults(config)
        result: _ConfigLoadResult = Success(config)
        _CONFIG_CACHE[config_path] = (None, result)
        return result
    try:
        raw = config_path.read_bytes()
        parsed = tomllib.loads(raw.decode("utf-8"))
        parsed = _merge_tool_output_overrides(
            parsed,
            source_env,
        )
        config = AnimaGlobalConfig(**parsed)
        overrides = _tool_output_env_overrides(source_env)
        if overrides:
            resolved = resolve_tool_output_precedence(
                cli=None,
                env=overrides,
                config=config.tool_output.model_dump(),
            )
            tool_output = config.tool_output.model_copy()
            profile = resolved.get("profile")
            if isinstance(profile, str):
                tool_output.profile = profile
            per_call_cap = resolved.get("per_call_cap_bytes")
            if isinstance(per_call_cap, int) and not isinstance(per_call_cap, bool):
                tool_output.per_call_cap_bytes = per_call_cap
            per_turn_cap = resolved.get("per_turn_cap_bytes")
            if isinstance(per_turn_cap, int) and not isinstance(per_turn_cap, bool):
                tool_output.per_turn_cap_bytes = per_turn_cap
            config.tool_output = tool_output
        # Spec source: vectl steps
        # ``model-simplify.simplify-create-agent`` and
        # ``deduplicate-load-global-config-calls``.
        # Apply config ``[env]`` entries once loaded, without overriding existing env.
        _apply_env_defaults(config)
        result = Success(config)
        _CONFIG_CACHE[config_path] = (mtime_ns, result)
        return result
    except OSError as exc:
        result = Failure(
            InternalError(
                f"Failed reading config file {config_path}: {exc}",
                details={"path": str(config_path), "error": str(exc)},
            )
        )
        _CONFIG_CACHE[config_path] = (mtime_ns, result)
        return result
    except (UnicodeDecodeError, tomllib.TOMLDecodeError, ValidationError) as exc:
        result = Failure(
            InvalidInputError(
                f"Invalid config file {config_path}: {exc}",
                details={"path": str(config_path), "error": str(exc)},
            )
        )
        _CONFIG_CACHE[config_path] = (mtime_ns, result)
        return result


def load_compaction_config(
    env: Mapping[str, str] | None = None,
) -> Result[CompactionConfig, InvalidInputError | InternalError]:
    """Load compaction configuration from global config.

    Behavior:
        - Calls load_global_config() internally
        - If global config exists: returns config.compaction
        - If no config file: returns default CompactionConfig()

    Args:
        env: Optional environment mapping (defaults to os.environ)

    Returns:
        Success(CompactionConfig) on success
        Failure(InvalidInputError | InternalError) on load failure

    Examples:
        >>> # When config file exists with compaction settings
        >>> result = load_compaction_config()
        >>> isinstance(result, Success)
        True
        >>> config = result.unwrap()
        >>> config.auto  # uses defaults
        True

        >>> # When no config file exists, returns defaults
        >>> from pathlib import Path
        >>> import tempfile
        >>> with tempfile.TemporaryDirectory() as tmp:
        ...     result = load_compaction_config(
        ...         env={"ANIMA_CONFIG": str(Path(tmp) / "missing.toml")}
        ...     )
        ...     isinstance(result, Success)
        True
    """
    result = load_global_config(env)
    if isinstance(result, Success):
        config = result.unwrap()
        return Success(config.compaction)
    else:
        # Forward the error from load_global_config
        return Failure(result.failure())


# @invar:allow shell_result: pure resolver shared by shell orchestration callers
# @shell_orchestration: kept in shell to co-locate config + runtime precedence policy
def resolve_model_context_limit(
    config: AnimaGlobalConfig | None,
    runtime_model_context_limit: int | None,
    *,
    model_id: str | None = None,
    hard_default: int = DEFAULT_MODEL_CONTEXT_LIMIT_TOKENS,
) -> int:
    """Resolve context window size used by compaction threshold checks.

    Precedence (highest to lowest):
        1. runtime_model_context_limit (from PydanticAI model metadata)
        2. Model registry lookup (litellm-sourced, user cache → embedded snapshot)
        3. config.defaults.context_window_tokens
        4. hard_default (128K)

    Args:
        config: Global config (may contain context_window_tokens).
        runtime_model_context_limit: From PydanticAI model object.
        model_id: Model identifier for registry lookup (e.g., "openai:gpt-5.2").
        hard_default: Fallback when all other sources are unset.
    """
    if runtime_model_context_limit is not None and runtime_model_context_limit > 0:
        return runtime_model_context_limit

    # Registry lookup: user cache → embedded snapshot
    if model_id is not None:
        from anima.core.model_registry import (
            get_context_limit,
            load_embedded_registry,
            load_user_cache,
        )

        # Try user cache first, then embedded
        registry = load_user_cache()
        if registry is None:
            registry = load_embedded_registry()
        if registry:
            limit = get_context_limit(model_id, registry)
            if limit != hard_default:  # found a match, not just the fallback
                return limit

    if config is not None:
        configured_limit = getattr(config.defaults, "context_window_tokens", None)
        if isinstance(configured_limit, int) and configured_limit > 0:
            return configured_limit

    return hard_default


# @invar:allow shell_result: shell helper returns scalar for resolver composition
# @shell_orchestration: scalar normalization supports shell-level config precedence
# @shell_complexity: bounded parsing helpers stay in shell where precedence values are normalized
def _as_positive_int(value: object) -> int | None:
    """Return positive integers; treat everything else as unset."""
    if isinstance(value, bool):
        return None
    if isinstance(value, str):
        if value.isdigit():
            parsed = int(value)
            if parsed > 0:
                return parsed
        return None
    if isinstance(value, int) and value > 0:
        return value
    if isinstance(value, str):
        text = value.strip()
        if not text:
            return None
        try:
            parsed = int(text)
        except ValueError:
            return None
        return parsed if parsed > 0 else None
    return None


# @invar:allow shell_result: deterministic scalar map used by shell policy resolution
# @shell_orchestration: centralizes CLI/env/config precedence for tool-output knobs
def resolve_tool_output_precedence(
    *,
    cli: Mapping[str, object] | None,
    env: Mapping[str, object] | None,
    config: Mapping[str, object] | None,
) -> dict[str, object]:
    """Resolve tool-output overrides with precedence CLI > env > config.

    Each field resolves independently.
    """

    def _first_profile() -> str | None:
        for source in (cli, env, config):
            if source is None:
                continue
            profile = _as_non_empty_str(source.get("profile"))
            if profile is not None:
                return profile
        return None

    def _first_positive_int(key: str) -> int | None:
        for source in (cli, env, config):
            if source is None:
                continue
            value = _as_positive_int(source.get(key))
            if value is not None:
                return value
        return None

    return {
        "profile": _first_profile(),
        "per_call_cap_bytes": _first_positive_int("per_call_cap_bytes"),
        "per_turn_cap_bytes": _first_positive_int("per_turn_cap_bytes"),
    }


# @invar:allow shell_result: deterministic scalar map consumed by tool-output assembly
# @shell_orchestration: budget resolver co-located with config/runtime precedence policy
# @invar:allow dead_param: signature intentionally carries context_window_tokens for compatibility
# with existing assembly call sites that pass runtime context metadata
def resolve_tool_output_budget_contract(
    *,
    context_window_tokens: int | None,
    selected_provider_profile: Mapping[str, object] | None,
    selected_model_profile: Mapping[str, object] | None,
    configured_per_call_cap_bytes: int | None = None,
    configured_per_turn_cap_bytes: int | None = None,
    hard_default_per_call_bytes: int = DEFAULT_TOOL_OUTPUT_PER_CALL_CAP_BYTES,
    hard_default_per_turn_bytes: int = DEFAULT_TOOL_OUTPUT_PER_TURN_CAP_BYTES,
) -> dict[str, object]:
    """Normalize effective per-call/per-turn budgets for a tool-output run.

    Conservative rules:
    - Provider hard caps always dominate when present.
    - Missing profiles fall back to explicit hard defaults.
    - Context-window metadata never inflates caps directly.
    """
    del context_window_tokens

    provider_call = None
    provider_turn = None
    provider_prefers_pagination = False
    if selected_provider_profile is not None:
        provider_call = _as_positive_int(
            selected_provider_profile.get("hard_cap_per_call_bytes")
        )
        provider_turn = _as_positive_int(
            selected_provider_profile.get("hard_cap_per_turn_bytes")
        )
        provider_prefers_pagination = bool(
            selected_provider_profile.get("prefers_pagination", False)
        )

    model_call = None
    model_turn = None
    expand_to_context_window = False
    if selected_model_profile is not None:
        model_call = _as_positive_int(selected_model_profile.get("per_call_cap_bytes"))
        model_turn = _as_positive_int(selected_model_profile.get("per_turn_cap_bytes"))
        expand_to_context_window = bool(
            selected_model_profile.get("expand_to_context_window", False)
        )

    configured_call = _as_positive_int(configured_per_call_cap_bytes)
    configured_turn = _as_positive_int(configured_per_turn_cap_bytes)

    per_call_cap_bytes = configured_call or model_call or hard_default_per_call_bytes
    per_turn_cap_bytes = configured_turn or model_turn or hard_default_per_turn_bytes

    if provider_call is not None:
        per_call_cap_bytes = min(per_call_cap_bytes, provider_call)
    if provider_turn is not None:
        per_turn_cap_bytes = min(per_turn_cap_bytes, provider_turn)

    return {
        "per_call_cap_bytes": per_call_cap_bytes,
        "per_turn_cap_bytes": per_turn_cap_bytes,
        "prefers_pagination": provider_prefers_pagination,
        "expand_to_context_window": expand_to_context_window,
    }


# @invar:allow shell_result: deterministic policy payload consumed by shell assembly
# @shell_orchestration: composes precedence + provider/model profile lookup in one place
# @shell_complexity: combines precedence + profile-selection branches in one wiring helper
# @invar:allow dead_export: public shell resolver consumed by upcoming tool-output assembly seam
def resolve_tool_output_budget_profile(
    *,
    cli: Mapping[str, object] | None,
    env: Mapping[str, object] | None,
    config: Mapping[str, object] | None,
    model_hint: str | None,
    provider_hint: str | None,
    context_window_tokens: int | None,
    hard_default_per_call_bytes: int = DEFAULT_TOOL_OUTPUT_PER_CALL_CAP_BYTES,
    hard_default_per_turn_bytes: int = DEFAULT_TOOL_OUTPUT_PER_TURN_CAP_BYTES,
) -> dict[str, object]:
    """Resolve effective tool-output budget from precedence + profile hints.

    Precedence for scalar overrides remains CLI > env > config. Provider/model
    profile maps are loaded from config and selected conservatively:
    - provider candidates: selected profile name, explicit provider_hint,
      parsed provider from model_hint
    - model candidates: explicit model_hint, model_hint suffix after provider
      prefix
    """
    resolved = resolve_tool_output_precedence(cli=cli, env=env, config=config)

    profile_name = _as_non_empty_str(resolved.get("profile"))
    hinted_provider = _as_non_empty_str(provider_hint)
    parsed_provider, parsed_model = _split_model_identifier(model_hint)

    provider_profile_name, selected_provider_profile = _resolve_profile_map_entry(
        config.get("provider_profiles") if config is not None else None,
        candidates=(profile_name, hinted_provider, parsed_provider),
    )

    model_profile_name, selected_model_profile = _resolve_profile_map_entry(
        config.get("model_profiles") if config is not None else None,
        candidates=(_as_non_empty_str(model_hint), parsed_model),
    )

    budget_contract = resolve_tool_output_budget_contract(
        context_window_tokens=context_window_tokens,
        selected_provider_profile=selected_provider_profile,
        selected_model_profile=selected_model_profile,
        configured_per_call_cap_bytes=_as_positive_int(
            resolved.get("per_call_cap_bytes")
        ),
        configured_per_turn_cap_bytes=_as_positive_int(
            resolved.get("per_turn_cap_bytes")
        ),
        hard_default_per_call_bytes=hard_default_per_call_bytes,
        hard_default_per_turn_bytes=hard_default_per_turn_bytes,
    )

    return {
        "profile": profile_name,
        "provider_profile_name": provider_profile_name,
        "model_profile_name": model_profile_name,
        "selected_provider_profile": dict(selected_provider_profile)
        if selected_provider_profile is not None
        else None,
        "selected_model_profile": dict(selected_model_profile)
        if selected_model_profile is not None
        else None,
        **budget_contract,
    }


# @invar:allow shell_result: deterministic output-policy defaults for execution contract snapshots
def resolve_execution_output_policy_defaults(
    *,
    cli: Mapping[str, object] | None,
    env: Mapping[str, object] | None,
    config: Mapping[str, object] | None,
    model_hint: str | None,
    provider_hint: str | None,
    context_window_tokens: int | None,
) -> dict[str, object]:
    """Resolve contract-level output-policy defaults from shell precedence.

    This helper centralizes shell-side budget/profile resolution used to seed
    ``ExecutionIntentDefaults`` so ``ResolvedExecutionContract.output_policy``
    becomes the persisted single source of truth for runtime consumers.
    """
    resolved_budget = resolve_tool_output_budget_profile(
        cli=cli,
        env=env,
        config=config,
        model_hint=model_hint,
        provider_hint=provider_hint,
        context_window_tokens=context_window_tokens,
    )
    prefers_pagination = bool(resolved_budget.get("prefers_pagination", False))
    truncation_strategy = "paginate_then_truncate" if prefers_pagination else "truncate"
    provider_profile = resolved_budget.get("provider_profile_name")
    if not isinstance(provider_profile, str) or len(provider_profile) == 0:
        provider_profile = "openai-compatible"

    return {
        "default_output_provider_profile": provider_profile,
        "default_per_call_tool_output_chars": resolved_budget.get("per_call_cap_bytes"),
        "default_per_turn_tool_output_chars": resolved_budget.get("per_turn_cap_bytes"),
        "default_truncation_strategy": truncation_strategy,
    }


# @invar:allow shell_result: shell helper returns optional scalar metadata value
# @shell_orchestration: runtime metadata extraction belongs with shell model wiring
# @shell_complexity: checks multiple runtime metadata shapes for provider compatibility
def evaluate_runtime_model_context_limit(runtime_model: object | None) -> int | None:
    """Read context-window limit from runtime model metadata when available.

    Spec source:
    - context-compaction-remediation.spec-clarify-context-limit-source

    The runtime model is the source of truth when it exposes a positive integer
    context-window limit. Missing or non-positive metadata is treated as "not
    set" so callers can fall back to config/default values.
    """
    if runtime_model is None:
        return None

    for attr in _RUNTIME_CONTEXT_LIMIT_ATTRS:
        value = _as_positive_int(getattr(runtime_model, attr, None))
        if value is not None:
            return value

    profile = getattr(runtime_model, "profile", None)
    for attr in _RUNTIME_CONTEXT_LIMIT_ATTRS:
        value = _as_positive_int(getattr(profile, attr, None))
        if value is not None:
            return value

    return None
