"""
tests/test_injection_slots.py — Tests for InjectionSlotManager (Issue #363).

Covers:
  - Each bundle type → correct slot position
  - Rendering produces valid text
  - Slot ordering is deterministic
  - Empty bundles excluded from output
  - Token counts accurate in assignments
"""

from __future__ import annotations

import pytest

from cortex.core.context_compiler import Bundle, BundleType, CompiledContext, ContextCompiler
from cortex.core.token_budget import TokenBudgetManager
from cortex.core.injection_slots import (
    BUNDLE_SLOT_MAP,
    InjectionAssignment,
    InjectionSlotManager,
    SlotPosition,
)
from cortex.types import CompressionLevel


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

@pytest.fixture
def slot_manager() -> InjectionSlotManager:
    return InjectionSlotManager()


def _make_bundle(
    bundle_type: BundleType,
    contents: list[str] | None = None,
    token_count: int = 10,
) -> Bundle:
    """Helper to create a Bundle with items."""
    if contents is None:
        contents = ["test item"]
    items = [
        {"content": c, "compression": "EVIDENCE", "token_count": 5, "original_tokens": 10}
        for c in contents
    ]
    return Bundle(
        bundle_type=bundle_type,
        items=items,
        token_count=token_count,
        compression_level=CompressionLevel.EVIDENCE,
    )


def _make_compiled(*bundles: Bundle) -> CompiledContext:
    """Helper to create a CompiledContext from bundles."""
    total = sum(b.token_count for b in bundles)
    return CompiledContext(
        bundles=list(bundles),
        total_tokens=total,
        budget_remaining=1000 - total,
    )


# ---------------------------------------------------------------------------
# SlotPosition enum
# ---------------------------------------------------------------------------

class TestSlotPosition:
    def test_values(self):
        assert SlotPosition.SYSTEM_PREFIX.value == "system_prefix"
        assert SlotPosition.SYSTEM_SUFFIX.value == "system_suffix"
        assert SlotPosition.USER_PREFIX.value == "user_prefix"
        assert SlotPosition.USER_SUFFIX.value == "user_suffix"

    def test_ordering(self):
        ordered = SlotPosition.by_order()
        assert ordered == [
            SlotPosition.SYSTEM_PREFIX,
            SlotPosition.SYSTEM_SUFFIX,
            SlotPosition.USER_PREFIX,
            SlotPosition.USER_SUFFIX,
        ]

    def test_four_positions(self):
        assert len(SlotPosition) == 4


# ---------------------------------------------------------------------------
# Bundle → Slot mapping
# ---------------------------------------------------------------------------

class TestBundleSlotMapping:
    def test_boot_to_system_prefix(self):
        assert BUNDLE_SLOT_MAP[BundleType.BOOT] == SlotPosition.SYSTEM_PREFIX

    def test_open_loops_to_system_suffix(self):
        assert BUNDLE_SLOT_MAP[BundleType.OPEN_LOOPS] == SlotPosition.SYSTEM_SUFFIX

    def test_task_to_system_suffix(self):
        assert BUNDLE_SLOT_MAP[BundleType.TASK] == SlotPosition.SYSTEM_SUFFIX

    def test_procedures_to_system_suffix(self):
        assert BUNDLE_SLOT_MAP[BundleType.PROCEDURES] == SlotPosition.SYSTEM_SUFFIX

    def test_recent_history_to_user_prefix(self):
        assert BUNDLE_SLOT_MAP[BundleType.RECENT_HISTORY] == SlotPosition.USER_PREFIX

    def test_defaults_to_user_prefix(self):
        assert BUNDLE_SLOT_MAP[BundleType.DEFAULTS] == SlotPosition.USER_PREFIX

    def test_evidence_to_user_suffix(self):
        assert BUNDLE_SLOT_MAP[BundleType.EVIDENCE] == SlotPosition.USER_SUFFIX

    def test_all_bundle_types_mapped(self):
        """Every BundleType must have a slot mapping."""
        for bt in BundleType:
            assert bt in BUNDLE_SLOT_MAP, f"Missing mapping for {bt}"


# ---------------------------------------------------------------------------
# Slot assignment
# ---------------------------------------------------------------------------

class TestAssignSlots:
    def test_single_bundle_assignment(self, slot_manager: InjectionSlotManager):
        bundle = _make_bundle(BundleType.BOOT, ["identity info"])
        compiled = _make_compiled(bundle)
        assignments = slot_manager.assign_slots(compiled)
        assert len(assignments) == 1
        assert assignments[0].slot_position == SlotPosition.SYSTEM_PREFIX
        assert assignments[0].bundle.bundle_type == BundleType.BOOT

    def test_multiple_bundles_assigned(self, slot_manager: InjectionSlotManager):
        bundles = [
            _make_bundle(BundleType.BOOT, ["boot info"]),
            _make_bundle(BundleType.TASK, ["task info"]),
            _make_bundle(BundleType.EVIDENCE, ["evidence info"]),
        ]
        compiled = _make_compiled(*bundles)
        assignments = slot_manager.assign_slots(compiled)
        assert len(assignments) == 3

    def test_empty_bundles_excluded(self, slot_manager: InjectionSlotManager):
        empty_bundle = Bundle(
            bundle_type=BundleType.EVIDENCE,
            items=[],
            token_count=0,
        )
        compiled = _make_compiled(empty_bundle)
        assignments = slot_manager.assign_slots(compiled)
        assert len(assignments) == 0

    def test_assignment_ordering_deterministic(self, slot_manager: InjectionSlotManager):
        bundles = [
            _make_bundle(BundleType.EVIDENCE, ["evidence"]),
            _make_bundle(BundleType.BOOT, ["boot"]),
            _make_bundle(BundleType.TASK, ["task"]),
        ]
        compiled = _make_compiled(*bundles)
        a1 = slot_manager.assign_slots(compiled)
        a2 = slot_manager.assign_slots(compiled)

        # Same order each time
        assert [a.slot_position for a in a1] == [a.slot_position for a in a2]
        assert [a.bundle.bundle_type for a in a1] == [a.bundle.bundle_type for a in a2]

    def test_slot_order_system_before_user(self, slot_manager: InjectionSlotManager):
        bundles = [
            _make_bundle(BundleType.EVIDENCE, ["evidence"]),
            _make_bundle(BundleType.BOOT, ["boot"]),
            _make_bundle(BundleType.RECENT_HISTORY, ["recent"]),
        ]
        compiled = _make_compiled(*bundles)
        assignments = slot_manager.assign_slots(compiled)

        positions = [a.slot_position for a in assignments]
        assert positions.index(SlotPosition.SYSTEM_PREFIX) < positions.index(SlotPosition.USER_SUFFIX)

    def test_token_counts_in_assignments(self, slot_manager: InjectionSlotManager):
        bundle = _make_bundle(BundleType.BOOT, ["identity info"], token_count=42)
        compiled = _make_compiled(bundle)
        assignments = slot_manager.assign_slots(compiled)
        assert assignments[0].token_count == 42


# ---------------------------------------------------------------------------
# Rendering
# ---------------------------------------------------------------------------

class TestRendering:
    def test_render_produces_valid_text(self, slot_manager: InjectionSlotManager):
        bundle = _make_bundle(BundleType.BOOT, ["I am Eva", "session started"])
        compiled = _make_compiled(bundle)
        assignments = slot_manager.assign_slots(compiled)
        rendered = slot_manager.render(assignments)

        assert SlotPosition.SYSTEM_PREFIX in rendered
        text = rendered[SlotPosition.SYSTEM_PREFIX]
        assert "## Session Context" in text
        assert "- I am Eva" in text
        assert "- session started" in text

    def test_render_multiple_bundles_same_slot(self, slot_manager: InjectionSlotManager):
        # open_loops and task both go to SYSTEM_SUFFIX
        bundles = [
            _make_bundle(BundleType.OPEN_LOOPS, ["loop 1"]),
            _make_bundle(BundleType.TASK, ["task 1"]),
        ]
        compiled = _make_compiled(*bundles)
        assignments = slot_manager.assign_slots(compiled)
        rendered = slot_manager.render(assignments)

        assert SlotPosition.SYSTEM_SUFFIX in rendered
        text = rendered[SlotPosition.SYSTEM_SUFFIX]
        assert "## Open Loops" in text
        assert "## Task Context" in text

    def test_render_empty_assignments(self, slot_manager: InjectionSlotManager):
        rendered = slot_manager.render([])
        assert rendered == {}

    def test_render_preserves_slot_order(self, slot_manager: InjectionSlotManager):
        bundles = [
            _make_bundle(BundleType.EVIDENCE, ["ev"]),
            _make_bundle(BundleType.BOOT, ["boot"]),
        ]
        compiled = _make_compiled(*bundles)
        assignments = slot_manager.assign_slots(compiled)
        rendered = slot_manager.render(assignments)

        keys = list(rendered.keys())
        assert keys == sorted(keys, key=lambda s: s.order)

    def test_render_excludes_empty_content_items(self, slot_manager: InjectionSlotManager):
        bundle = Bundle(
            bundle_type=BundleType.BOOT,
            items=[
                {"content": "real content", "compression": "EVIDENCE", "token_count": 5, "original_tokens": 5},
                {"content": "", "compression": "EVIDENCE", "token_count": 0, "original_tokens": 0},
            ],
            token_count=5,
            compression_level=CompressionLevel.EVIDENCE,
        )
        compiled = _make_compiled(bundle)
        assignments = slot_manager.assign_slots(compiled)
        rendered = slot_manager.render(assignments)

        text = rendered[SlotPosition.SYSTEM_PREFIX]
        assert "- real content" in text
        # Empty content should not produce a bullet
        lines = text.strip().split("\n")
        bullet_lines = [l for l in lines if l.startswith("- ")]
        assert len(bullet_lines) == 1


# ---------------------------------------------------------------------------
# Full pipeline (compile → assign → render)
# ---------------------------------------------------------------------------

class TestFullPipeline:
    def test_end_to_end(self, slot_manager: InjectionSlotManager):
        manager = TokenBudgetManager()
        compiler = ContextCompiler(token_budget=manager)

        items = [
            {"content": "I am Eva", "type": "identity", "source": "handoff",
             "relevance_score": 0.0, "task_mode": "", "stability": "",
             "created_at": "", "session_id": ""},
            {"content": "Ship by Friday", "type": "commitment", "source": "",
             "relevance_score": 0.0, "task_mode": "", "stability": "",
             "created_at": "", "session_id": ""},
        ]

        compiled = compiler.compile(items, 500)
        assignments = slot_manager.assign_slots(compiled)
        rendered = slot_manager.render(assignments)

        # Boot goes to SYSTEM_PREFIX, commitment to SYSTEM_SUFFIX
        assert SlotPosition.SYSTEM_PREFIX in rendered
        assert SlotPosition.SYSTEM_SUFFIX in rendered
        assert compiled.total_tokens <= 500
