"""
tests/test_token_budget.py — Tests for TokenBudgetManager (Issue #18)

Covers:
  - Token counting accuracy
  - Budget creation and state
  - Layer allocation with priority order
  - Cornerstone reservation guarantee
  - Budget exhaustion and truncation
  - Soft cap enforcement
  - Configurable budgets
  - BudgetAllocation properties
  - MemoryLayer enum helpers
  - Edge cases (empty input, zero budget, single item overflow)
"""

from __future__ import annotations

import pytest

from cortex.core.token_budget import TokenBudget, TokenBudgetManager
from cortex.types import (
    BudgetAllocation,
    CompressedItem,
    CompressionLevel,
    InjectionSlot,
    MemoryLayer,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def manager() -> TokenBudgetManager:
    """Default manager: 4000 total, 800 cornerstone."""
    return TokenBudgetManager()


@pytest.fixture
def small_manager() -> TokenBudgetManager:
    """Small manager for budget-exhaustion tests."""
    return TokenBudgetManager(total_budget=200, cornerstone_budget=50)


# ---------------------------------------------------------------------------
# MemoryLayer enum
# ---------------------------------------------------------------------------

class TestMemoryLayer:
    def test_values(self):
        assert MemoryLayer.CORNERSTONE.value == "cornerstone"
        assert MemoryLayer.SESSION.value == "session"
        assert MemoryLayer.SEMANTIC.value == "semantic"
        assert MemoryLayer.PROCEDURAL.value == "procedural"

    def test_priority_order(self):
        ordered = MemoryLayer.by_priority()
        assert ordered[0] == MemoryLayer.CORNERSTONE
        assert ordered[-1] == MemoryLayer.PROCEDURAL

    def test_truncation_order_is_lowest_first(self):
        trunc = MemoryLayer.truncation_order()
        assert trunc[0] == MemoryLayer.PROCEDURAL
        assert trunc[-1] == MemoryLayer.CORNERSTONE

    def test_priority_values(self):
        assert MemoryLayer.CORNERSTONE.priority < MemoryLayer.SESSION.priority
        assert MemoryLayer.SESSION.priority < MemoryLayer.SEMANTIC.priority
        assert MemoryLayer.SEMANTIC.priority < MemoryLayer.PROCEDURAL.priority


# ---------------------------------------------------------------------------
# BudgetAllocation dataclass
# ---------------------------------------------------------------------------

class TestBudgetAllocation:
    def test_used_calculation(self):
        alloc = BudgetAllocation(
            total_budget=4000,
            cornerstone_reserved=800,
            allocations={"session": 300, "semantic": 200},
        )
        assert alloc.used == 1300  # 800 + 300 + 200

    def test_remaining(self):
        alloc = BudgetAllocation(
            total_budget=4000,
            cornerstone_reserved=800,
            allocations={"session": 200},
        )
        assert alloc.remaining == 3000  # 4000 - 800 - 200

    def test_utilization(self):
        alloc = BudgetAllocation(
            total_budget=4000,
            cornerstone_reserved=800,
            allocations={"session": 200},
        )
        assert alloc.utilization == pytest.approx(1000 / 4000)

    def test_utilization_zero_budget(self):
        alloc = BudgetAllocation(
            total_budget=0,
            cornerstone_reserved=0,
            allocations={},
        )
        assert alloc.utilization == 0.0

    def test_is_over_budget(self):
        alloc = BudgetAllocation(
            total_budget=100,
            cornerstone_reserved=80,
            allocations={"session": 30},  # 80+30 = 110 > 100
        )
        assert alloc.is_over_budget() is True

    def test_not_over_budget(self):
        alloc = BudgetAllocation(
            total_budget=4000,
            cornerstone_reserved=800,
            allocations={"session": 100},
        )
        assert alloc.is_over_budget() is False


# ---------------------------------------------------------------------------
# TokenBudget (internal state)
# ---------------------------------------------------------------------------

class TestTokenBudget:
    def test_remaining_starts_at_total_minus_reserved(self):
        budget = TokenBudget(total=4000, reserved_cornerstones=800)
        assert budget.remaining() == 3200

    def test_allocate_succeeds_within_budget(self):
        budget = TokenBudget(total=4000, reserved_cornerstones=800)
        assert budget.allocate("session", 500) is True
        assert budget.allocated["session"] == 500

    def test_allocate_fails_over_budget(self):
        budget = TokenBudget(total=1000, reserved_cornerstones=800)
        assert budget.allocate("session", 300) is False  # only 200 remaining
        assert "session" not in budget.allocated

    def test_allocate_accumulates(self):
        budget = TokenBudget(total=4000, reserved_cornerstones=0)
        budget.allocate("session", 200)
        budget.allocate("session", 300)
        assert budget.allocated["session"] == 500

    def test_used(self):
        budget = TokenBudget(total=4000, reserved_cornerstones=800)
        budget.allocate("session", 200)
        assert budget.used() == 1000

    def test_remaining_decreases_after_allocation(self):
        budget = TokenBudget(total=4000, reserved_cornerstones=800)
        budget.allocate("session", 500)
        assert budget.remaining() == 2700


# ---------------------------------------------------------------------------
# TokenBudgetManager construction
# ---------------------------------------------------------------------------

class TestManagerConstruction:
    def test_defaults(self, manager):
        assert manager.total_budget == 4000
        assert manager.cornerstone_budget == 800

    def test_custom_budgets(self):
        m = TokenBudgetManager(total_budget=2000, cornerstone_budget=400)
        assert m.total_budget == 2000
        assert m.cornerstone_budget == 400

    def test_cornerstone_exceeds_total_raises(self):
        with pytest.raises(ValueError, match="cornerstone_budget"):
            TokenBudgetManager(total_budget=500, cornerstone_budget=600)

    def test_config_override(self):
        class FakeConfig:
            total_memory_token_budget = 3000
            cornerstone_token_budget = 600

        m = TokenBudgetManager(config=FakeConfig())
        assert m.total_budget == 3000
        assert m.cornerstone_budget == 600

    def test_explicit_args_take_precedence_over_config(self):
        class FakeConfig:
            total_memory_token_budget = 3000
            cornerstone_token_budget = 600

        m = TokenBudgetManager(config=FakeConfig(), total_budget=5000)
        assert m.total_budget == 5000

    def test_repr(self, manager):
        r = repr(manager)
        assert "TokenBudgetManager" in r
        assert "4000" in r


# ---------------------------------------------------------------------------
# Token counting
# ---------------------------------------------------------------------------

class TestTokenCounting:
    def test_empty_string_is_zero(self, manager):
        assert manager.count_tokens("") == 0

    def test_simple_word(self, manager):
        # "hello" is a single token in cl100k_base
        count = manager.count_tokens("hello")
        assert count >= 1

    def test_longer_text_has_more_tokens(self, manager):
        short = manager.count_tokens("hi")
        long = manager.count_tokens("hi " * 50)
        assert long > short

    def test_batch_counting(self, manager):
        texts = ["hello", "world", "foo bar"]
        counts = manager.count_tokens_batch(texts)
        assert len(counts) == 3
        assert all(c >= 1 for c in counts)
        # Individual counts should match
        for text, count in zip(texts, counts):
            assert manager.count_tokens(text) == count

    def test_whitespace_only(self, manager):
        count = manager.count_tokens("   ")
        assert count >= 1  # whitespace is still tokenized

    def test_known_token_count(self, manager):
        # "The quick brown fox" — known to be 4 tokens in cl100k_base
        count = manager.count_tokens("The quick brown fox")
        assert count == 4


# ---------------------------------------------------------------------------
# Budget creation
# ---------------------------------------------------------------------------

class TestCreateBudget:
    def test_creates_correct_total(self, manager):
        b = manager.create_budget()
        assert b.total == 4000

    def test_creates_correct_cornerstone_reservation(self, manager):
        b = manager.create_budget()
        assert b.reserved_cornerstones == 800

    def test_override_total(self, manager):
        b = manager.create_budget(total=2000)
        assert b.total == 2000

    def test_starts_empty(self, manager):
        b = manager.create_budget()
        assert b.allocated == {}

    def test_remaining_equals_total_minus_cornerstones(self, manager):
        b = manager.create_budget()
        assert b.remaining() == 3200


# ---------------------------------------------------------------------------
# Layer soft caps
# ---------------------------------------------------------------------------

class TestLayerSoftCaps:
    def test_session_cap(self, manager):
        # 40% of non-cornerstone = 40% of (4000 - 800) = 40% of 3200 = 1280
        cap = manager.layer_soft_cap(MemoryLayer.SESSION)
        assert cap == 1280

    def test_semantic_cap(self, manager):
        # 35% of 3200 = 1120
        cap = manager.layer_soft_cap(MemoryLayer.SEMANTIC)
        assert cap == 1120

    def test_procedural_cap(self, manager):
        # 25% of 3200 = 800
        cap = manager.layer_soft_cap(MemoryLayer.PROCEDURAL)
        assert cap == 800

    def test_custom_fractions(self):
        m = TokenBudgetManager(
            total_budget=4000,
            cornerstone_budget=0,
            layer_fractions={MemoryLayer.SESSION.value: 0.5},
        )
        assert m.layer_soft_cap(MemoryLayer.SESSION) == 2000


# ---------------------------------------------------------------------------
# Allocate layers — happy path
# ---------------------------------------------------------------------------

class TestAllocateLayers:
    def test_empty_layers_returns_zero_allocation(self, manager):
        result = manager.allocate_layers({})
        assert result.allocations == {}
        assert result.truncated == {}
        assert result.remaining >= 0

    def test_cornerstones_always_allocated(self, manager):
        layers = {
            MemoryLayer.CORNERSTONE: ["identity memory one", "identity memory two"],
        }
        result = manager.allocate_layers(layers)
        assert MemoryLayer.CORNERSTONE.value in result.allocations
        assert result.allocations[MemoryLayer.CORNERSTONE.value] > 0

    def test_session_allocated_after_cornerstones(self, manager):
        layers = {
            MemoryLayer.CORNERSTONE: ["cornerstone text"],
            MemoryLayer.SESSION: ["recent session content here"],
        }
        result = manager.allocate_layers(layers)
        assert MemoryLayer.SESSION.value in result.allocations

    def test_all_layers_allocated_when_budget_sufficient(self, manager):
        layers = {
            MemoryLayer.CORNERSTONE: ["cornerstone"],
            MemoryLayer.SESSION: ["session memory"],
            MemoryLayer.SEMANTIC: ["semantic memory"],
            MemoryLayer.PROCEDURAL: ["procedural step"],
        }
        result = manager.allocate_layers(layers)
        # All layers should have some allocation
        assert MemoryLayer.CORNERSTONE.value in result.allocations
        assert MemoryLayer.SESSION.value in result.allocations
        assert MemoryLayer.SEMANTIC.value in result.allocations
        assert MemoryLayer.PROCEDURAL.value in result.allocations

    def test_total_budget_not_exceeded(self, manager):
        """Total tokens allocated must never exceed the total budget."""
        layers = {
            MemoryLayer.CORNERSTONE: ["cornerstone " * 30],
            MemoryLayer.SESSION: ["session content " * 100],
            MemoryLayer.SEMANTIC: ["semantic fact " * 100],
            MemoryLayer.PROCEDURAL: ["procedure step " * 100],
        }
        result = manager.allocate_layers(layers)
        assert result.used <= result.total_budget

    def test_budget_allocation_has_correct_total(self, manager):
        layers = {MemoryLayer.SESSION: ["test"]}
        result = manager.allocate_layers(layers)
        assert result.total_budget == 4000
        assert result.cornerstone_reserved == 800


# ---------------------------------------------------------------------------
# Cornerstone reservation guarantee
# ---------------------------------------------------------------------------

class TestCornerstoneReservation:
    def test_cornerstones_fit_within_reserved_budget(self, manager):
        # Build text that fits exactly in cornerstone budget
        text = "word " * 100  # ~100 tokens
        layers = {
            MemoryLayer.CORNERSTONE: [text] * 8,  # ~800 tokens
            MemoryLayer.SESSION: ["session " * 200],  # should NOT eat cornerstone space
        }
        result = manager.allocate_layers(layers)
        # Session tokens should not appear in cornerstone allocation
        cornerstone_used = result.allocations.get(MemoryLayer.CORNERSTONE.value, 0)
        assert cornerstone_used <= 800

    def test_cornerstone_budget_independent_of_other_layers(self, small_manager):
        """Even when other layers overflow, cornerstones get reserved space."""
        # Fill session/semantic with massive content
        big_content = ["token " * 200]  # ~200 tokens each
        layers = {
            MemoryLayer.CORNERSTONE: ["identity"],
            MemoryLayer.SESSION: big_content,
            MemoryLayer.SEMANTIC: big_content,
            MemoryLayer.PROCEDURAL: big_content,
        }
        result = small_manager.allocate_layers(layers)
        # Cornerstones should still be allocated despite budget pressure
        assert MemoryLayer.CORNERSTONE.value in result.allocations

    def test_cornerstone_overflow_recorded_in_truncated(self, small_manager):
        """When cornerstones exceed their own budget, overflow is recorded."""
        # Build content larger than cornerstone_budget (50 tokens)
        huge_cornerstone = "word " * 100  # ~100 tokens, > 50 budget
        layers = {
            MemoryLayer.CORNERSTONE: [huge_cornerstone],
        }
        result = small_manager.allocate_layers(layers)
        # Some should be truncated
        assert MemoryLayer.CORNERSTONE.value in result.truncated


# ---------------------------------------------------------------------------
# Budget exhaustion and truncation
# ---------------------------------------------------------------------------

class TestBudgetExhaustion:
    def test_lowest_priority_truncated_first(self, small_manager):
        """When budget is tight, procedural is truncated before session."""
        medium_content = "word " * 20  # ~20 tokens each
        layers = {
            MemoryLayer.SESSION: [medium_content] * 5,
            MemoryLayer.PROCEDURAL: [medium_content] * 5,
        }
        result = small_manager.allocate_layers(layers)
        # Either procedural is truncated, or it gets less than session
        session_tokens = result.allocations.get(MemoryLayer.SESSION.value, 0)
        procedural_tokens = result.allocations.get(MemoryLayer.PROCEDURAL.value, 0)
        # Session should get at least as much as procedural (higher priority)
        assert session_tokens >= procedural_tokens

    def test_truncated_tokens_recorded(self, small_manager):
        """Overflow tokens should appear in result.truncated."""
        big_content = "word " * 100  # ~100 tokens
        layers = {
            MemoryLayer.SESSION: [big_content, big_content, big_content],
        }
        result = small_manager.allocate_layers(layers)
        # With 200 total - 50 cornerstone = 150 non-cornerstone, 3×100 = 300 tokens
        # Some should be truncated
        assert MemoryLayer.SESSION.value in result.truncated

    def test_total_never_exceeds_budget_under_pressure(self, small_manager):
        """With massively oversized content, budget must still be respected."""
        enormous = "word " * 5000  # thousands of tokens
        layers = {
            MemoryLayer.CORNERSTONE: [enormous],
            MemoryLayer.SESSION: [enormous],
            MemoryLayer.SEMANTIC: [enormous],
            MemoryLayer.PROCEDURAL: [enormous],
        }
        result = small_manager.allocate_layers(layers)
        assert result.used <= result.total_budget


# ---------------------------------------------------------------------------
# Truncation helpers
# ---------------------------------------------------------------------------

class TestTruncateHelpers:
    def test_truncate_to_budget_within_limit(self, manager):
        text = "short text"
        result = manager.truncate_to_budget(text, max_tokens=50)
        assert result == text

    def test_truncate_to_budget_over_limit(self, manager):
        text = "word " * 100  # ~100 tokens
        result = manager.truncate_to_budget(text, max_tokens=20)
        assert manager.count_tokens(result) <= 20

    def test_truncate_to_budget_zero_limit(self, manager):
        result = manager.truncate_to_budget("hello world", max_tokens=0)
        assert result == ""

    def test_truncate_appends_suffix(self, manager):
        text = "word " * 100
        result = manager.truncate_to_budget(text, max_tokens=10)
        assert result.endswith("…")

    def test_truncate_no_suffix_when_fits(self, manager):
        text = "hi"
        result = manager.truncate_to_budget(text, max_tokens=100)
        assert "…" not in result

    def test_fit_items_to_budget_selects_front(self, manager):
        items = ["item one", "item two", "item three", "item four", "item five"]
        # Force a very small budget to limit selection
        selected, used = manager.fit_items_to_budget(items, max_tokens=3)
        # At least some items selected, front items preferred
        assert len(selected) > 0
        assert selected[0] == items[0]
        assert used <= 3

    def test_fit_items_to_budget_empty_list(self, manager):
        selected, used = manager.fit_items_to_budget([], max_tokens=100)
        assert selected == []
        assert used == 0

    def test_fit_items_to_budget_all_fit(self, manager):
        items = ["a", "b", "c"]
        selected, used = manager.fit_items_to_budget(items, max_tokens=1000)
        assert selected == items
        assert used == sum(manager.count_tokens(i) for i in items)

    def test_fit_items_total_within_budget(self, manager):
        items = ["word " * 20 for _ in range(10)]
        selected, used = manager.fit_items_to_budget(items, max_tokens=100)
        assert used <= 100


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------

class TestEdgeCases:
    def test_zero_total_budget(self):
        m = TokenBudgetManager(total_budget=0, cornerstone_budget=0)
        result = m.allocate_layers({MemoryLayer.SESSION: ["content"]})
        assert result.used == 0

    def test_no_layers_provided(self, manager):
        result = manager.allocate_layers({})
        assert result.used == result.cornerstone_reserved
        assert result.remaining == manager.total_budget - manager.cornerstone_budget

    def test_single_item_exceeds_total(self):
        m = TokenBudgetManager(total_budget=5, cornerstone_budget=0)
        result = m.allocate_layers({
            MemoryLayer.SESSION: ["this sentence has far more than five tokens total"],
        })
        assert result.used <= 5

    def test_unicode_text_counted(self, manager):
        text = "こんにちは世界"  # Japanese
        count = manager.count_tokens(text)
        assert count >= 1

    def test_layer_fractions_respected(self):
        m = TokenBudgetManager(
            total_budget=1000,
            cornerstone_budget=0,
            layer_fractions={
                MemoryLayer.SESSION.value: 0.5,
                MemoryLayer.SEMANTIC.value: 0.5,
                MemoryLayer.PROCEDURAL.value: 0.0,
            },
        )
        # Procedural cap is 0, so nothing should be allocated there
        big = "word " * 200
        result = m.allocate_layers({
            MemoryLayer.PROCEDURAL: [big],
            MemoryLayer.SESSION: ["short"],
        })
        assert MemoryLayer.PROCEDURAL.value not in result.allocations
        assert MemoryLayer.SESSION.value in result.allocations


# ===========================================================================
# Issue #374 — Slot-based allocation
# ===========================================================================

class TestInjectionSlot:
    """Tests for the InjectionSlot enum (Issue #374)."""

    def test_all_slots_exist(self):
        expected = {
            "task_context", "boot_context", "open_loops",
            "semantic_claims", "procedural", "episodic",
        }
        assert {s.value for s in InjectionSlot} == expected

    def test_priority_order(self):
        ordered = InjectionSlot.by_priority()
        assert ordered[0] == InjectionSlot.TASK_CONTEXT
        assert ordered[-1] == InjectionSlot.EPISODIC

    def test_priority_values_ascending(self):
        slots = InjectionSlot.by_priority()
        for a, b in zip(slots, slots[1:]):
            assert a.priority < b.priority

    def test_default_priority_names_matches_order(self):
        names = InjectionSlot.default_priority_names()
        assert names[0] == "task_context"
        assert names[-1] == "episodic"
        assert len(names) == len(list(InjectionSlot))


class TestContextWindowBudgets:
    """Tests for context_window_budgets() (Issue #374)."""

    def test_fractions_sum_to_one(self):
        m = TokenBudgetManager()
        assert abs(m.memory_fraction + m.user_fraction + m.reserve_fraction - 1.0) < 1e-9

    def test_default_fractions(self):
        m = TokenBudgetManager()
        assert m.memory_fraction == pytest.approx(0.25)
        assert m.user_fraction == pytest.approx(0.55)
        assert m.reserve_fraction == pytest.approx(0.20)

    def test_context_window_splits_correctly(self):
        m = TokenBudgetManager()
        splits = m.context_window_budgets(10000)
        assert splits["memory"] == 2500
        assert splits["user"] == 5500
        # reserve gets remainder
        assert splits["reserve"] == 10000 - 2500 - 5500

    def test_context_window_no_negative_reserve(self):
        # memory + user fractions exceed 1.0 due to rounding — reserve should
        # clamp to 0, not go negative.
        m = TokenBudgetManager(memory_fraction=0.60, user_fraction=0.60, reserve_fraction=0.0)
        splits = m.context_window_budgets(100)
        assert splits["reserve"] >= 0

    def test_custom_fractions_accepted(self):
        m = TokenBudgetManager(
            memory_fraction=0.30,
            user_fraction=0.50,
            reserve_fraction=0.20,
        )
        assert m.memory_fraction == pytest.approx(0.30)
        assert m.user_fraction == pytest.approx(0.50)

    def test_zero_context_window(self):
        m = TokenBudgetManager()
        splits = m.context_window_budgets(0)
        assert splits["memory"] == 0
        assert splits["user"] == 0
        assert splits["reserve"] == 0


class TestAllocateSlots:
    """Tests for allocate_slots() (Issue #374)."""

    def test_single_slot_gets_full_budget(self):
        m = TokenBudgetManager()
        result = m.allocate_slots(
            total_budget=500,
            slots_with_items={
                InjectionSlot.TASK_CONTEXT: ["short task text"],
            },
        )
        assert "task_context" in result
        assert result["task_context"] > 0
        assert result["task_context"] <= 500

    def test_priority_order_allocation(self):
        """High-priority slot consumes budget before low-priority slot."""
        m = TokenBudgetManager()
        # Use a tiny budget so only the first slot can be allocated
        big_text = "word " * 80  # ~80 tokens
        result = m.allocate_slots(
            total_budget=100,
            slots_with_items={
                InjectionSlot.TASK_CONTEXT: [big_text],
                InjectionSlot.EPISODIC: ["episodic item"],
            },
        )
        # TASK_CONTEXT is highest priority and should be allocated
        assert "task_context" in result
        # EPISODIC may or may not fit depending on remaining budget
        task_tokens = result.get("task_context", 0)
        episodic_tokens = result.get("episodic", 0)
        assert task_tokens + episodic_tokens <= 100

    def test_empty_slots_omitted_from_result(self):
        m = TokenBudgetManager()
        result = m.allocate_slots(
            total_budget=500,
            slots_with_items={
                InjectionSlot.TASK_CONTEXT: ["task text"],
                InjectionSlot.OPEN_LOOPS: [],  # empty — should be omitted
            },
        )
        assert "open_loops" not in result

    def test_zero_budget_returns_empty(self):
        m = TokenBudgetManager()
        result = m.allocate_slots(
            total_budget=0,
            slots_with_items={
                InjectionSlot.TASK_CONTEXT: ["task text"],
            },
        )
        assert result == {}

    def test_total_allocated_within_budget(self):
        m = TokenBudgetManager()
        result = m.allocate_slots(
            total_budget=200,
            slots_with_items={
                InjectionSlot.TASK_CONTEXT: ["task context text"],
                InjectionSlot.BOOT_CONTEXT: ["boot context text"],
                InjectionSlot.OPEN_LOOPS: ["open loop item"],
                InjectionSlot.SEMANTIC_CLAIMS: ["semantic claim"],
                InjectionSlot.PROCEDURAL: ["procedure step"],
                InjectionSlot.EPISODIC: ["historical event"],
            },
        )
        total_allocated = sum(result.values())
        assert total_allocated <= 200

    def test_overflow_trimming(self):
        """Items exceeding slot budget are dropped (not truncated)."""
        m = TokenBudgetManager()
        # Single item that is too long for the tiny budget
        long_item = "overflow " * 100  # ~100 tokens
        result = m.allocate_slots(
            total_budget=20,
            slots_with_items={
                InjectionSlot.TASK_CONTEXT: [long_item],
                InjectionSlot.EPISODIC: ["short"],
            },
        )
        # Long item should NOT be allocated (exceeds 20-token budget)
        task_tokens = result.get("task_context", 0)
        assert task_tokens == 0 or task_tokens <= 20

    def test_custom_slot_priority_order_respected(self):
        """slot_priorities config controls allocation order."""
        # Reverse priority: EPISODIC first, TASK_CONTEXT last
        m = TokenBudgetManager(
            slot_priorities=["episodic", "task_context"],
        )
        big_text = "word " * 80  # ~80 tokens
        result = m.allocate_slots(
            total_budget=100,
            slots_with_items={
                InjectionSlot.TASK_CONTEXT: [big_text],
                InjectionSlot.EPISODIC: [big_text],
            },
        )
        # With reversed priority, EPISODIC gets the budget first
        assert "episodic" in result

    def test_all_slots_empty_returns_empty(self):
        m = TokenBudgetManager()
        result = m.allocate_slots(
            total_budget=1000,
            slots_with_items={},
        )
        assert result == {}


# ===========================================================================
# Issue #373 — Bundle compression
# ===========================================================================

class TestCompressionLevel:
    """Tests for the CompressionLevel enum (Issue #373)."""

    def test_values_are_ordered(self):
        assert CompressionLevel.POINTER < CompressionLevel.COMPACT
        assert CompressionLevel.COMPACT < CompressionLevel.EVIDENCE

    def test_integer_values(self):
        assert CompressionLevel.POINTER.value == 1
        assert CompressionLevel.COMPACT.value == 2
        assert CompressionLevel.EVIDENCE.value == 3

    def test_all_three_levels_exist(self):
        levels = {lv.name for lv in CompressionLevel}
        assert levels == {"POINTER", "COMPACT", "EVIDENCE"}


class TestCompressedItem:
    """Tests for the CompressedItem dataclass."""

    def test_compression_ratio_full(self):
        item = CompressedItem(
            content="hello world",
            compression=CompressionLevel.EVIDENCE,
            token_count=10,
            original_tokens=10,
        )
        assert item.compression_ratio == pytest.approx(1.0)

    def test_compression_ratio_partial(self):
        item = CompressedItem(
            content="short label",
            compression=CompressionLevel.POINTER,
            token_count=3,
            original_tokens=60,
        )
        assert item.compression_ratio == pytest.approx(3 / 60)

    def test_tokens_saved(self):
        item = CompressedItem(
            content="label",
            compression=CompressionLevel.POINTER,
            token_count=5,
            original_tokens=50,
        )
        assert item.tokens_saved == 45

    def test_zero_original_tokens_ratio(self):
        item = CompressedItem(
            content="",
            compression=CompressionLevel.POINTER,
            token_count=0,
            original_tokens=0,
        )
        assert item.compression_ratio == pytest.approx(1.0)


class TestCompressToFit:
    """Tests for TokenBudgetManager.compress_to_fit() (Issue #373)."""

    def test_items_fit_at_evidence_level(self):
        """When budget is generous and min_compression=EVIDENCE, items use EVIDENCE."""
        m = TokenBudgetManager()
        items = ["This is a short evidence item."]
        # Audit fix: must explicitly allow EVIDENCE (default min is COMPACT)
        result = m.compress_to_fit(items, target_tokens=1000, min_compression=CompressionLevel.EVIDENCE)
        assert len(result) == 1
        assert result[0].compression == CompressionLevel.EVIDENCE

    def test_items_compressed_to_compact_on_tight_budget(self):
        """Tight budget forces items into COMPACT compression."""
        m = TokenBudgetManager()
        # A longer item that won't fit at EVIDENCE within 30 tokens
        long_item = "word " * 50  # ~50 tokens
        result = m.compress_to_fit([long_item], target_tokens=30)
        assert len(result) == 1
        assert result[0].compression in (
            CompressionLevel.COMPACT, CompressionLevel.POINTER
        )
        assert result[0].token_count <= 30

    def test_items_compressed_to_pointer_on_very_tight_budget(self):
        """Very tight budget forces POINTER compression."""
        m = TokenBudgetManager()
        long_item = "word " * 100  # ~100 tokens
        # Audit fix: pointer_max is now 12, so budget=10 ensures POINTER level
        result = m.compress_to_fit([long_item], target_tokens=10)
        assert len(result) == 1
        assert result[0].compression == CompressionLevel.POINTER
        assert result[0].token_count <= 12

    def test_zero_budget_returns_empty(self):
        m = TokenBudgetManager()
        result = m.compress_to_fit(["some text"], target_tokens=0)
        assert result == []

    def test_all_items_compressed_to_pointer_edge_case(self):
        """All items are forced to pointer when budget is very small."""
        m = TokenBudgetManager()
        items = ["item one", "item two", "item three"]
        # Tiny budget — barely enough for a few pointer-level items
        result = m.compress_to_fit(items, target_tokens=45)
        total_tokens = sum(r.token_count for r in result)
        assert total_tokens <= 45
        for r in result:
            assert r.compression in (
                CompressionLevel.POINTER,
                CompressionLevel.COMPACT,
                CompressionLevel.EVIDENCE,
            )

    def test_min_compression_compact_prevents_evidence(self):
        """min_compression=COMPACT prevents EVIDENCE even with large budget."""
        m = TokenBudgetManager()
        # Use an item longer than pointer_max (12) so it doesn't fall to POINTER
        item = "This is a moderately long item that could be evidence level and is definitely longer."
        result = m.compress_to_fit(
            [item],
            target_tokens=1000,
            min_compression=CompressionLevel.COMPACT,
        )
        assert len(result) == 1
        # Audit fix: must be exactly COMPACT, not EVIDENCE (COMPACT blocks EVIDENCE)
        assert result[0].compression == CompressionLevel.COMPACT

    def test_min_compression_pointer_forces_all_to_pointer(self):
        """min_compression=POINTER means everything is pointer-level."""
        m = TokenBudgetManager()
        items = ["item A here", "item B here", "item C here"]
        result = m.compress_to_fit(
            items,
            target_tokens=500,
            min_compression=CompressionLevel.POINTER,
        )
        for r in result:
            assert r.compression == CompressionLevel.POINTER

    def test_original_tokens_recorded(self):
        """CompressedItem.original_tokens reflects the raw item token count."""
        m = TokenBudgetManager()
        item = "hello world"
        original_count = m.count_tokens(item)
        result = m.compress_to_fit([item], target_tokens=1000)
        assert len(result) == 1
        assert result[0].original_tokens == original_count

    def test_total_output_tokens_within_target(self):
        """Sum of output token counts must never exceed target_tokens."""
        m = TokenBudgetManager()
        items = [f"item number {i} with some content" for i in range(10)]
        target = 50
        result = m.compress_to_fit(items, target_tokens=target)
        total = sum(r.token_count for r in result)
        assert total <= target

    def test_items_skipped_when_cannot_fit_at_any_level(self):
        """Items that cannot fit at POINTER level are omitted entirely."""
        m = TokenBudgetManager()
        # Budget so small that even a pointer for ANY item can't fit
        result = m.compress_to_fit(["some text"], target_tokens=1)
        # Either result is empty or the single item fits within 1 token
        if result:
            assert result[0].token_count <= 1

    def test_empty_items_list(self):
        m = TokenBudgetManager()
        result = m.compress_to_fit([], target_tokens=500)
        assert result == []


# ===========================================================================
# Backward compatibility — allocate_layers() must still work unchanged
# ===========================================================================

class TestBackwardCompat:
    """Ensure Issue #374/#373 changes don't break the existing allocate_layers API."""

    def test_allocate_layers_still_returns_budget_allocation(self):
        m = TokenBudgetManager()
        result = m.allocate_layers(
            layers={
                MemoryLayer.CORNERSTONE: ["cornerstone text"],
                MemoryLayer.SESSION: ["session text"],
            }
        )
        assert isinstance(result, BudgetAllocation)

    def test_allocate_layers_priority_unchanged(self):
        """Cornerstone is still guaranteed; session before procedural."""
        m = TokenBudgetManager(total_budget=1000, cornerstone_budget=200)
        big = "word " * 200  # ~200 tokens
        result = m.allocate_layers({
            MemoryLayer.CORNERSTONE: ["cornerstone text"],
            MemoryLayer.SESSION: [big],
            MemoryLayer.PROCEDURAL: [big],
        })
        assert result.cornerstone_reserved == 200
        # Session has higher priority than procedural — should be allocated first
        session_alloc = result.allocations.get(MemoryLayer.SESSION.value, 0)
        procedural_alloc = result.allocations.get(MemoryLayer.PROCEDURAL.value, 0)
        # Both can't exceed total budget
        assert result.used <= 1000

    def test_allocate_layers_with_new_manager_config(self):
        """Manager constructed with new fraction params still handles old API."""
        m = TokenBudgetManager(
            total_budget=2000,
            cornerstone_budget=400,
            memory_fraction=0.30,
            user_fraction=0.50,
            reserve_fraction=0.20,
        )
        result = m.allocate_layers({
            MemoryLayer.SESSION: ["session content here"],
            MemoryLayer.SEMANTIC: ["semantic claim"],
        })
        assert isinstance(result, BudgetAllocation)
        assert result.total_budget == 2000

    def test_existing_properties_unchanged(self):
        m = TokenBudgetManager(total_budget=5000, cornerstone_budget=1000)
        assert m.total_budget == 5000
        assert m.cornerstone_budget == 1000

    def test_new_properties_accessible(self):
        m = TokenBudgetManager(
            memory_fraction=0.28,
            user_fraction=0.52,
            reserve_fraction=0.20,
            slot_priorities=["task_context", "episodic"],
        )
        assert m.memory_fraction == pytest.approx(0.28)
        assert m.user_fraction == pytest.approx(0.52)
        assert m.reserve_fraction == pytest.approx(0.20)
        assert m.slot_priorities == ["task_context", "episodic"]


# ---------------------------------------------------------------------------
# Audit fix: fraction validation
# ---------------------------------------------------------------------------

class TestFractionValidation:
    def test_negative_memory_fraction_raises(self):
        with pytest.raises(ValueError, match="memory_fraction"):
            TokenBudgetManager(memory_fraction=-0.1)

    def test_negative_user_fraction_raises(self):
        with pytest.raises(ValueError, match="user_fraction"):
            TokenBudgetManager(user_fraction=-0.5)

    def test_negative_reserve_fraction_raises(self):
        with pytest.raises(ValueError, match="reserve_fraction"):
            TokenBudgetManager(reserve_fraction=-0.01)

    def test_fraction_above_one_raises(self):
        with pytest.raises(ValueError, match="memory_fraction"):
            TokenBudgetManager(memory_fraction=1.5)

    def test_valid_fractions_accepted(self):
        m = TokenBudgetManager(memory_fraction=0.0, user_fraction=1.0, reserve_fraction=0.0)
        assert m.memory_fraction == 0.0
        assert m.user_fraction == 1.0


# ---------------------------------------------------------------------------
# Audit fix: min_compression=COMPACT blocks EVIDENCE (additional coverage)
# ---------------------------------------------------------------------------

class TestMinCompressionCompactBlocksEvidence:
    def test_generous_budget_compact_min_gives_compact_not_evidence(self):
        """With generous budget and min_compression=COMPACT, result is COMPACT."""
        m = TokenBudgetManager()
        # 21-token item, budget=500 — would be EVIDENCE if allowed
        item = "word " * 20  # 21 tokens
        result = m.compress_to_fit([item], target_tokens=500, min_compression=CompressionLevel.COMPACT)
        assert len(result) == 1
        assert result[0].compression == CompressionLevel.COMPACT

    def test_evidence_min_allows_evidence(self):
        """min_compression=EVIDENCE allows EVIDENCE level."""
        m = TokenBudgetManager()
        item = "word " * 20  # 21 tokens
        result = m.compress_to_fit([item], target_tokens=500, min_compression=CompressionLevel.EVIDENCE)
        assert len(result) == 1
        assert result[0].compression == CompressionLevel.EVIDENCE
