"""Tests for ContextEngine."""

from context_manager.budget.char_counter import CharCounter
from context_manager.engine import ContextEngine
from context_manager.models import ContextBlock, Priority
from context_manager.strategies.priority import PriorityPruning


def _engine(token_limit: int = 200) -> ContextEngine:
    """Create an engine with a deterministic char-based counter."""
    return ContextEngine(
        model="test",
        token_limit=token_limit,
        pruning_strategy=PriorityPruning(),
        counter=CharCounter(chars_per_token=1),  # 1 char = 1 token for easy math
    )


class TestContextEngine:
    def test_add_and_compile(self):
        engine = _engine(token_limit=500)
        engine.add(ContextBlock(content="hello", role="system", priority="critical"))
        engine.add(ContextBlock(content="world", role="user", priority="medium"))

        result = engine.compile()
        assert len(result) == 2
        assert result[0]["role"] == "system"
        assert result[1]["content"] == "world"

    def test_pruning_under_budget(self):
        engine = _engine(token_limit=15)
        engine.add(ContextBlock(content="critical!", role="system", priority="critical"))  # 9 tokens
        engine.add(ContextBlock(content="low", role="user", priority="low"))  # 3 tokens
        engine.add(ContextBlock(content="high!", role="user", priority="high"))  # 5 tokens

        result = engine.compile()
        contents = [m["content"] for m in result]
        # 9 + 5 = 14 fits; 9 + 5 + 3 = 17 doesn't
        assert "critical!" in contents
        assert "high!" in contents
        assert "low" not in contents

    def test_total_tokens(self):
        engine = _engine()
        engine.add(ContextBlock(content="abcd", priority="low"))  # 4 tokens
        engine.add(ContextBlock(content="ef", priority="high"))  # 2 tokens
        assert engine.total_tokens == 6

    def test_compiled_tokens_and_remaining(self):
        engine = _engine(token_limit=100)
        engine.add(ContextBlock(content="a" * 30, priority="high"))
        engine.add(ContextBlock(content="b" * 20, priority="low"))
        engine.compile()
        assert engine.compiled_tokens == 50
        assert engine.remaining_tokens == 50

    def test_reset(self):
        engine = _engine()
        engine.add(ContextBlock(content="data"))
        engine.reset()
        assert engine.blocks == []
        assert engine.total_tokens == 0

    def test_method_chaining(self):
        engine = _engine()
        result = engine.add(ContextBlock(content="a"))
        assert result is engine

    def test_summary(self):
        engine = _engine(token_limit=10)
        engine.add(ContextBlock(content="keep", role="system", priority="critical"))  # 4
        engine.add(ContextBlock(content="drop_me!", role="user", priority="low"))  # 8
        engine.compile()

        s = engine.summary()
        assert s["kept_blocks"] == 1
        assert s["dropped_blocks"] == 1
        assert s["compiled_tokens"] == 4
        assert len(s["dropped"]) == 1
        assert s["dropped"][0]["priority"] == "LOW"

    def test_visualize_returns_string(self, capsys):
        engine = _engine(token_limit=100)
        engine.add(ContextBlock(content="a" * 50, role="system", priority="critical"))
        engine.compile()
        output = engine.visualize()
        assert "system" in output
        assert "unused" in output
