"""Helpers for Amazon Bedrock prompt caching."""

from __future__ import annotations

import os
from typing import Any


BEDROCK_PROMPT_CACHE_TTL_ENV_VAR = "BEDROCK_PROMPT_CACHE_TTL"
SUPPORTED_BEDROCK_PROMPT_CACHE_TTLS = frozenset({"5m", "1h"})
_DISABLED_TTL_VALUES = frozenset({"", "0", "false", "none", "off", "disabled"})
_ONE_HOUR_CACHE_MODEL_MARKERS = (
    "claude-sonnet-4-5",
    "claude-haiku-4-5",
    "claude-opus-4-5",
)


def model_supports_bedrock_one_hour_prompt_caching(model_id: str | None) -> bool:
    """Return whether the model supports Bedrock's 1-hour prompt cache TTL."""
    if not model_id:
        return False

    model_id = model_id.lower()
    return any(marker in model_id for marker in _ONE_HOUR_CACHE_MODEL_MARKERS)


def default_bedrock_prompt_cache_ttl(model_id: str | None) -> str | None:
    """Use Bedrock's 1-hour cache automatically for the Claude 4.5 family."""
    if model_supports_bedrock_one_hour_prompt_caching(model_id):
        return "1h"
    return None


def normalize_bedrock_prompt_cache_ttl(
    ttl: Any, model_id: str | None = None
) -> str | None:
    """Normalize and validate a requested Bedrock prompt-cache TTL."""
    if ttl is None:
        return None

    ttl = str(ttl).strip().lower()
    if ttl in _DISABLED_TTL_VALUES:
        return None
    if ttl not in SUPPORTED_BEDROCK_PROMPT_CACHE_TTLS:
        raise ValueError(
            "Unsupported Bedrock prompt cache TTL "
            f"{ttl!r}. Expected one of: {sorted(SUPPORTED_BEDROCK_PROMPT_CACHE_TTLS)}"
        )
    if (
        ttl == "1h"
        and model_id
        and not model_supports_bedrock_one_hour_prompt_caching(model_id)
    ):
        raise ValueError(
            "Bedrock 1-hour prompt caching is only supported on Claude Sonnet 4.5, "
            f"Claude Haiku 4.5, and Claude Opus 4.5. Got model {model_id!r}."
        )
    return ttl


def resolve_bedrock_prompt_cache_ttl(
    model_id: str | None = None,
    configured_ttl: Any = None,
) -> str | None:
    """Resolve the effective Bedrock prompt-cache TTL.

    Priority:
    1. Explicit parameter
    2. `BEDROCK_PROMPT_CACHE_TTL`
    3. Automatic `1h` for Claude 4.5 Bedrock models
    """
    if configured_ttl is None:
        env_ttl = os.getenv(BEDROCK_PROMPT_CACHE_TTL_ENV_VAR)
        configured_ttl = (
            env_ttl
            if env_ttl is not None
            else default_bedrock_prompt_cache_ttl(model_id)
        )

    return normalize_bedrock_prompt_cache_ttl(configured_ttl, model_id=model_id)


def build_bedrock_cache_control(ttl: str | None) -> dict[str, str]:
    """Build an Anthropic-style Bedrock cache control block."""
    cache_control = {"type": "ephemeral"}
    if ttl:
        cache_control["ttl"] = ttl
    return cache_control


def build_bedrock_cache_point(ttl: str | None) -> dict[str, str]:
    """Build a Converse API cache point block."""
    cache_point = {"type": "default"}
    if ttl:
        cache_point["ttl"] = ttl
    return cache_point


def apply_bedrock_cache_control_to_recent_messages(
    messages: list[dict[str, Any]],
    ttl: str | None,
    recent_user_turns: int = 2,
) -> None:
    """Attach cache control to the most recent user turns in-place."""
    if not ttl:
        return

    cache_control = build_bedrock_cache_control(ttl)
    remaining = recent_user_turns

    for message in reversed(messages):
        if message.get("role") != "user":
            continue

        content = message.get("content")
        if not isinstance(content, list) or not content:
            continue

        last_block = content[-1]
        if not isinstance(last_block, dict):
            continue

        if remaining > 0:
            last_block["cache_control"] = dict(cache_control)
            remaining -= 1
        else:
            last_block.pop("cache_control", None)
