"""Tests for the Layer 4 N-most-recent image pruning module (pruning.py)."""

from __future__ import annotations

import os
from unittest.mock import patch

import pytest

from mcp_vm_blackbox.pruning import (
    _HAS_ANTHROPIC,
    PRUNE_MIN_REMOVAL_THRESHOLD,
    PRUNE_RETAIN_COUNT,
    is_pruning_available,
    prune_stale_images,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _image_msg(data: str = "imgdata") -> dict:
    """Build a minimal Anthropic-format image message."""
    return {"role": "user", "content": [{"type": "image", "source": {"data": data}}]}


def _text_msg(text: str = "hello") -> dict:
    """Build a minimal Anthropic-format text message."""
    return {"role": "user", "content": [{"type": "text", "text": text}]}


def _mixed_msg(text: str = "hello", img_data: str = "imgdata") -> dict:
    """Build a message with both text and image content blocks."""
    return {
        "role": "user",
        "content": [{"type": "text", "text": text}, {"type": "image", "source": {"data": img_data}}],
    }


# ---------------------------------------------------------------------------
# Module constants
# ---------------------------------------------------------------------------


class TestModuleConstants:
    def test_prune_retain_count_default_value(self) -> None:
        assert PRUNE_RETAIN_COUNT == 3

    def test_prune_min_removal_threshold_default_value(self) -> None:
        assert PRUNE_MIN_REMOVAL_THRESHOLD == 5

    def test_has_anthropic_flag_is_bool(self) -> None:
        assert isinstance(_HAS_ANTHROPIC, bool)


# ---------------------------------------------------------------------------
# is_pruning_available
# ---------------------------------------------------------------------------


class TestIsPruningAvailable:
    def test_is_pruning_available_returns_false_when_no_api_key(self) -> None:
        # Arrange — ensure env var absent
        with patch.dict(os.environ, {}, clear=True):
            env = {k: v for k, v in os.environ.items() if k != "ANTHROPIC_API_KEY"}
            with patch.dict(os.environ, env, clear=True):
                # Act & Assert
                result = is_pruning_available()
        # If anthropic is not installed, always False regardless of key
        if not _HAS_ANTHROPIC:
            assert result is False

    def test_is_pruning_available_returns_false_when_has_anthropic_false(self) -> None:
        # Arrange
        with (
            patch("mcp_vm_blackbox.pruning._HAS_ANTHROPIC", False),
            patch.dict(os.environ, {"ANTHROPIC_API_KEY": "sk-fake"}),
        ):
            # Act
            result = is_pruning_available()
        # Assert
        assert result is False

    def test_is_pruning_available_returns_true_when_sdk_and_key_present(self) -> None:
        # Arrange
        with (
            patch("mcp_vm_blackbox.pruning._HAS_ANTHROPIC", True),
            patch.dict(os.environ, {"ANTHROPIC_API_KEY": "sk-fake"}),
        ):
            # Act
            result = is_pruning_available()
        # Assert
        assert result is True

    def test_is_pruning_available_returns_false_when_sdk_present_but_no_key(self) -> None:
        # Arrange — remove key from env
        with patch("mcp_vm_blackbox.pruning._HAS_ANTHROPIC", True), patch.dict(os.environ, {}, clear=True):
            # Act
            result = is_pruning_available()
        # Assert
        assert result is False


# ---------------------------------------------------------------------------
# prune_stale_images — below threshold (no pruning)
# ---------------------------------------------------------------------------


class TestPruneStaleImagesNoOp:
    def test_prune_stale_images_below_threshold_returns_unchanged_copy(self) -> None:
        # Arrange — 4 image messages, threshold is 5
        messages = [_image_msg(str(i)) for i in range(4)]
        # Act
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        # Assert
        assert result == messages
        assert result is not messages  # must be a copy

    def test_prune_stale_images_empty_list_returns_empty_list(self) -> None:
        result = prune_stale_images([], retain_count=3, min_removal_threshold=5)
        assert result == []

    def test_prune_stale_images_no_images_returns_unchanged_copy(self) -> None:
        # Arrange — only text messages
        messages = [_text_msg(str(i)) for i in range(10)]
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        assert result == messages

    def test_prune_stale_images_exactly_at_threshold_is_pruned(self) -> None:
        # Arrange — exactly min_removal_threshold images
        messages = [_image_msg(str(i)) for i in range(5)]
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        kept = sum(1 for m in result if any(b.get("type") == "image" for b in m["content"]))
        assert kept == 3


# ---------------------------------------------------------------------------
# prune_stale_images — pruning applied
# ---------------------------------------------------------------------------


class TestPruneStaleImagesPruning:
    def test_prune_stale_images_retains_n_most_recent_images(self) -> None:
        # Arrange — 10 image messages
        messages = [_image_msg(str(i)) for i in range(10)]
        # Act
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        # Assert
        kept = sum(1 for m in result if any(b.get("type") == "image" for b in m["content"]))
        assert kept == 3

    def test_prune_stale_images_most_recent_images_kept_intact(self) -> None:
        # Arrange — 10 image messages; last 3 should remain image blocks
        messages = [_image_msg(str(i)) for i in range(10)]
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        # The last 3 messages should still have image blocks
        for msg in result[-3:]:
            block_types = [b.get("type") for b in msg["content"]]
            assert "image" in block_types

    def test_prune_stale_images_older_images_replaced_with_text_placeholder(self) -> None:
        # Arrange
        messages = [_image_msg(str(i)) for i in range(10)]
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        # The first 7 messages should have text blocks, not image blocks
        for msg in result[:7]:
            block_types = [b.get("type") for b in msg["content"]]
            assert "image" not in block_types
            assert "text" in block_types

    def test_prune_stale_images_placeholder_text_mentions_retained_count(self) -> None:
        # Arrange
        messages = [_image_msg(str(i)) for i in range(10)]
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        # Check placeholder text content in a pruned message
        pruned_msg = result[0]
        text_block = next(b for b in pruned_msg["content"] if b.get("type") == "text")
        assert "3" in text_block["text"]  # retain_count
        assert "10" in text_block["text"]  # total count

    def test_prune_stale_images_does_not_mutate_input(self) -> None:
        # Arrange
        messages = [_image_msg(str(i)) for i in range(10)]
        original_first_content = [dict(b) for b in messages[0]["content"]]
        # Act
        prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        # Assert — original list and message content unchanged
        assert messages[0]["content"] == original_first_content

    def test_prune_stale_images_result_length_unchanged(self) -> None:
        # Pruning replaces blocks but does not remove messages
        messages = [_image_msg(str(i)) for i in range(10)]
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        assert len(result) == len(messages)

    def test_prune_stale_images_non_image_messages_never_modified(self) -> None:
        # Arrange — interleave text and image messages.
        # Layout: [text-0, img-0, text-1, img-1, text-2, img-2, extra-0..4]
        # Text messages are at indices 0, 2, 4; image messages at 1, 3, 5, 6-10.
        messages = []
        text_indices: list[int] = []
        for i in range(3):
            text_indices.append(len(messages))
            messages.extend((_text_msg(f"text-{i}"), _image_msg(str(i))))
        # 3 text + 3 image = 6 total; add more images to exceed threshold of 5
        messages.extend(_image_msg(f"extra-{i}") for i in range(5))
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        # Pure-text messages (indices 0, 2, 4) must be completely untouched
        for idx in text_indices:
            assert messages[idx] == result[idx]

    def test_prune_stale_images_mixed_message_image_block_replaced(self) -> None:
        # Arrange — mixed messages (text + image blocks in same message)
        # Use 8 images to exceed threshold
        messages = [_mixed_msg(f"t{i}", f"img{i}") for i in range(8)]
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        # The first 5 mixed messages should have their image blocks pruned
        for msg in result[:5]:
            block_types = [b.get("type") for b in msg["content"]]
            assert "image" not in block_types
            # text block should still be present
            assert "text" in block_types

    def test_prune_stale_images_mixed_message_text_block_preserved(self) -> None:
        # Arrange
        messages = [_mixed_msg(f"t{i}", f"img{i}") for i in range(8)]
        result = prune_stale_images(messages, retain_count=3, min_removal_threshold=5)
        # Original text block in a pruned message should still exist
        pruned_msg = result[0]
        text_blocks = [
            b for b in pruned_msg["content"] if b.get("type") == "text" and b.get("text", "").startswith("t")
        ]
        assert len(text_blocks) == 1

    def test_prune_stale_images_retain_count_zero_prunes_all(self) -> None:
        # Arrange
        messages = [_image_msg(str(i)) for i in range(10)]
        result = prune_stale_images(messages, retain_count=0, min_removal_threshold=5)
        kept = sum(1 for m in result if any(b.get("type") == "image" for b in m["content"]))
        assert kept == 0

    def test_prune_stale_images_retain_count_exceeds_total_keeps_all(self) -> None:
        # retain_count > total image count — no pruning occurs
        messages = [_image_msg(str(i)) for i in range(10)]
        result = prune_stale_images(messages, retain_count=20, min_removal_threshold=5)
        kept = sum(1 for m in result if any(b.get("type") == "image" for b in m["content"]))
        assert kept == 10

    def test_prune_stale_images_uses_default_retain_count(self) -> None:
        # Call without explicit retain_count — uses PRUNE_RETAIN_COUNT (3)
        messages = [_image_msg(str(i)) for i in range(10)]
        result = prune_stale_images(messages)
        kept = sum(1 for m in result if any(b.get("type") == "image" for b in m["content"]))
        assert kept == PRUNE_RETAIN_COUNT

    def test_prune_stale_images_uses_default_min_removal_threshold(self) -> None:
        # Call without explicit min_removal_threshold — uses PRUNE_MIN_REMOVAL_THRESHOLD (5)
        # 4 image messages should NOT be pruned (below default threshold of 5)
        messages = [_image_msg(str(i)) for i in range(4)]
        result = prune_stale_images(messages)
        kept = sum(1 for m in result if any(b.get("type") == "image" for b in m["content"]))
        assert kept == 4  # unchanged


# ---------------------------------------------------------------------------
# Importability without anthropic
# ---------------------------------------------------------------------------


class TestImportabilityWithoutAnthropic:
    def test_module_importable_even_when_has_anthropic_false(self) -> None:
        # The module import itself must not raise even if _HAS_ANTHROPIC is False.
        # Since we can already import it here, this is a regression guard.
        from mcp_vm_blackbox import pruning

        assert pruning is not None

    def test_prune_stale_images_callable_without_anthropic(self) -> None:
        with patch("mcp_vm_blackbox.pruning._HAS_ANTHROPIC", False):
            # prune_stale_images does not call the SDK; must work regardless
            messages = [_image_msg(str(i)) for i in range(10)]
            result = prune_stale_images(messages)
            assert len(result) == 10


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
