"""Tests for PriorityPruning strategy."""

from context_manager.models import ContextBlock, Priority
from context_manager.strategies.priority import PriorityPruning


def _make_block(content: str, priority: Priority, tokens: int) -> ContextBlock:
    """Helper to create a block with a preset token count."""
    b = ContextBlock(content=content, priority=priority)
    b.token_count = tokens
    return b


class TestPriorityPruning:
    def test_all_fit(self):
        blocks = [
            _make_block("a", Priority.LOW, 10),
            _make_block("b", Priority.HIGH, 10),
        ]
        result = PriorityPruning().prune(blocks, token_budget=100)
        assert len(result) == 2

    def test_drops_low_first(self):
        blocks = [
            _make_block("critical", Priority.CRITICAL, 50),
            _make_block("low", Priority.LOW, 50),
            _make_block("high", Priority.HIGH, 50),
        ]
        result = PriorityPruning().prune(blocks, token_budget=100)
        assert len(result) == 2
        priorities = {b.priority for b in result}
        assert Priority.LOW not in priorities

    def test_preserves_insertion_order(self):
        blocks = [
            _make_block("first-high", Priority.HIGH, 10),
            _make_block("low", Priority.LOW, 10),
            _make_block("second-high", Priority.HIGH, 10),
        ]
        result = PriorityPruning().prune(blocks, token_budget=30)
        contents = [b.content for b in result]
        assert contents == ["first-high", "low", "second-high"]

    def test_drops_lowest_when_tight(self):
        blocks = [
            _make_block("sys", Priority.CRITICAL, 40),
            _make_block("doc1", Priority.MEDIUM, 30),
            _make_block("doc2", Priority.LOW, 30),
            _make_block("history", Priority.HIGH, 30),
        ]
        result = PriorityPruning().prune(blocks, token_budget=80)
        kept_contents = {b.content for b in result}
        assert "sys" in kept_contents
        assert "history" in kept_contents
        # doc2 (LOW) should be dropped first, doc1 (MEDIUM) if needed
        assert "doc2" not in kept_contents

    def test_empty_blocks(self):
        assert PriorityPruning().prune([], token_budget=100) == []

    def test_zero_budget(self):
        blocks = [_make_block("a", Priority.CRITICAL, 10)]
        assert PriorityPruning().prune(blocks, token_budget=0) == []
