"""Layer 4 N-most-recent image pruning for Anthropic message lists.

This module provides client-side pruning of stale images from the Anthropic message
format used by the Haiku micro-pilot in vm_viewport_action (Path B). The anthropic SDK
is an optional dependency — this module is importable and functional without it, but
``is_pruning_available()`` will return False and callers must check before pruning.

Pattern adapted from the Anthropic computer-use-demo
``_maybe_filter_to_n_most_recent_images()`` function.
"""

from __future__ import annotations

import os

try:
    import anthropic as _anthropic  # noqa: F401 — presence check only

    _HAS_ANTHROPIC: bool = True
except ImportError:
    _HAS_ANTHROPIC = False

# Module-level defaults exposed for callers that want to reference the package defaults.
PRUNE_RETAIN_COUNT: int = 3
PRUNE_MIN_REMOVAL_THRESHOLD: int = 5


def is_pruning_available() -> bool:
    """Return True when both the anthropic SDK is installed and an API key is set.

    Returns:
        True if pruning can be applied; False otherwise.
    """
    return _HAS_ANTHROPIC and bool(os.environ.get("ANTHROPIC_API_KEY"))


def prune_stale_images(
    messages: list[dict],
    retain_count: int = PRUNE_RETAIN_COUNT,
    min_removal_threshold: int = PRUNE_MIN_REMOVAL_THRESHOLD,
) -> list[dict]:
    """Remove older image content blocks from a message list, keeping the N most recent.

    Scans ``messages`` for any message whose ``content`` list contains a block with
    ``type == 'image'``.  When the total count of such messages is below
    ``min_removal_threshold``, returns the list unchanged (a copy).  Otherwise the
    oldest image-bearing messages (beyond ``retain_count``) have their image content
    blocks replaced with a text placeholder indicating the prune operation.

    Non-image messages and non-image content blocks within a pruned message are never
    modified.  The input list is never mutated.

    Args:
        messages: A list of Anthropic-format message dicts, each with a ``role`` key
            and a ``content`` key whose value is a list of content-block dicts.
        retain_count: Number of most-recent image-bearing messages to keep intact.
            Defaults to ``PRUNE_RETAIN_COUNT`` (3).
        min_removal_threshold: Minimum number of image-bearing messages required before
            any pruning is applied.  Defaults to ``PRUNE_MIN_REMOVAL_THRESHOLD`` (5).

    Returns:
        A new list of message dicts.  Pruned image blocks are replaced with a text
        block describing how many images were retained.
    """
    # Identify which message indices carry at least one image block.
    image_message_indices: list[int] = [
        idx
        for idx, msg in enumerate(messages)
        if isinstance(msg.get("content"), list)
        and any(isinstance(block, dict) and block.get("type") == "image" for block in msg["content"])
    ]

    total_image_messages = len(image_message_indices)

    # Not enough images to warrant pruning — return a shallow copy unchanged.
    if total_image_messages < min_removal_threshold:
        return list(messages)

    # The indices that should be pruned: everything except the retain_count most recent.
    indices_to_prune: set[int] = (
        set(image_message_indices[:-retain_count]) if retain_count > 0 else set(image_message_indices)
    )

    result: list[dict] = []
    for idx, msg in enumerate(messages):
        if idx not in indices_to_prune:
            result.append(msg)
            continue

        # Replace image blocks in this message with a text placeholder.
        new_content: list[dict] = []
        for block in msg.get("content", []):
            if isinstance(block, dict) and block.get("type") == "image":
                new_content.append({
                    "type": "text",
                    "text": f"[image pruned - {retain_count} of {total_image_messages} retained]",
                })
            else:
                new_content.append(block)

        result.append({**msg, "content": new_content})

    return result
