"""
tests/test_context_compiler_locomo.py — Tests for memory_class classification and
salience-aware ordering in ContextCompiler (Issue #525).

Covers:
  - memory_class="open_loop" → OPEN_LOOPS bundle (Rule 2b)
  - memory_class="procedural" → PROCEDURES bundle (Rule 3b)
  - memory_class="episodic" → RECENT_HISTORY even when older than 24h (Rule 5b)
  - Salience sorting within bundles (high > medium > low)
  - Interaction: memory_class rules don't override higher-priority type rules
"""

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from unittest.mock import MagicMock

import pytest

from cortex.core.context_compiler import (
    BundleType,
    ContextCompiler,
)
from cortex.core.task_mode import TaskMode
from cortex.core.token_budget import TokenBudgetManager


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

@pytest.fixture
def manager() -> TokenBudgetManager:
    return TokenBudgetManager()


@pytest.fixture
def compiler(manager: TokenBudgetManager) -> ContextCompiler:
    return ContextCompiler(token_budget=manager)


def _make_item(
    content: str = "test content",
    type: str = "",
    source: str = "",
    relevance_score: float = 0.0,
    task_mode: str = "",
    stability: str = "",
    created_at: str = "",
    memory_class: str = "",
    salience: str = "",
) -> dict:
    """Helper to build a memory item dict, including memory_class and salience."""
    item: dict = {
        "content": content,
        "type": type,
        "source": source,
        "relevance_score": relevance_score,
        "task_mode": task_mode,
        "stability": stability,
        "created_at": created_at,
    }
    if memory_class:
        item["memory_class"] = memory_class
    if salience:
        item["salience"] = salience
    return item


def _recent_ts() -> str:
    """ISO timestamp from 1 hour ago (within 24h window)."""
    return (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat()


def _old_ts() -> str:
    """ISO timestamp from 48 hours ago (outside 24h window)."""
    return (datetime.now(timezone.utc) - timedelta(hours=48)).isoformat()


# ---------------------------------------------------------------------------
# Helper: call internal _classify_item
# ---------------------------------------------------------------------------

def _classify(compiler: ContextCompiler, item: dict, task_mode=None) -> BundleType:
    return compiler._classify_item(item, task_mode)


# ---------------------------------------------------------------------------
# Part A: memory_class classification — open_loop (Rule 2b)
# ---------------------------------------------------------------------------

class TestMemoryClassOpenLoop:
    def test_memory_class_open_loop_goes_to_open_loops(self, compiler: ContextCompiler):
        """memory_class='open_loop' without type field → OPEN_LOOPS."""
        item = _make_item(memory_class="open_loop", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.OPEN_LOOPS

    def test_memory_class_open_loop_with_neutral_type(self, compiler: ContextCompiler):
        """memory_class='open_loop' with a non-matching type → still OPEN_LOOPS."""
        item = _make_item(type="fact", memory_class="open_loop", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.OPEN_LOOPS

    def test_type_commitment_still_wins_over_memory_class(self, compiler: ContextCompiler):
        """type='commitment' already triggers Rule 2 before Rule 2b — result same."""
        item = _make_item(type="commitment", memory_class="open_loop", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.OPEN_LOOPS

    def test_memory_class_open_loop_old_timestamp(self, compiler: ContextCompiler):
        """memory_class='open_loop' with old timestamp still goes to OPEN_LOOPS (not EVIDENCE)."""
        item = _make_item(memory_class="open_loop", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.OPEN_LOOPS


# ---------------------------------------------------------------------------
# Part A: memory_class classification — procedural (Rule 3b)
# ---------------------------------------------------------------------------

class TestMemoryClassProcedural:
    def test_memory_class_procedural_goes_to_procedures(self, compiler: ContextCompiler):
        """memory_class='procedural' without type field → PROCEDURES."""
        item = _make_item(memory_class="procedural", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.PROCEDURES

    def test_memory_class_procedural_with_neutral_type(self, compiler: ContextCompiler):
        """memory_class='procedural' with a generic type → PROCEDURES."""
        item = _make_item(type="fact", memory_class="procedural", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.PROCEDURES

    def test_type_procedure_wins_before_rule_3b(self, compiler: ContextCompiler):
        """type='procedure' triggers Rule 3, memory_class='procedural' would trigger 3b — both PROCEDURES."""
        item = _make_item(type="procedure", memory_class="procedural", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.PROCEDURES

    def test_memory_class_procedural_old_timestamp(self, compiler: ContextCompiler):
        """memory_class='procedural' with old timestamp → PROCEDURES (not EVIDENCE)."""
        item = _make_item(memory_class="procedural", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.PROCEDURES

    def test_memory_class_procedural_does_not_get_open_loops(self, compiler: ContextCompiler):
        """memory_class='procedural' must NOT go to OPEN_LOOPS."""
        item = _make_item(memory_class="procedural", created_at=_old_ts())
        assert _classify(compiler, item) != BundleType.OPEN_LOOPS


# ---------------------------------------------------------------------------
# Part B: episodic → RECENT_HISTORY boost (Rule 5b)
# ---------------------------------------------------------------------------

class TestMemoryClassEpisodic:
    def test_episodic_old_item_goes_to_recent_history(self, compiler: ContextCompiler):
        """memory_class='episodic' with 48h-old timestamp → RECENT_HISTORY (bypasses 24h rule)."""
        item = _make_item(memory_class="episodic", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.RECENT_HISTORY

    def test_episodic_very_old_item_still_recent_history(self, compiler: ContextCompiler):
        """memory_class='episodic' from a week ago → RECENT_HISTORY."""
        week_ago = (datetime.now(timezone.utc) - timedelta(days=7)).isoformat()
        item = _make_item(memory_class="episodic", created_at=week_ago)
        assert _classify(compiler, item) == BundleType.RECENT_HISTORY

    def test_episodic_with_recent_ts_stays_recent_history(self, compiler: ContextCompiler):
        """memory_class='episodic' + within 24h → RECENT_HISTORY (Rule 5 fires first, same result)."""
        item = _make_item(memory_class="episodic", created_at=_recent_ts())
        assert _classify(compiler, item) == BundleType.RECENT_HISTORY

    def test_episodic_does_not_override_boot(self, compiler: ContextCompiler):
        """source='handoff' (Rule 1) must beat memory_class='episodic' (Rule 5b)."""
        item = _make_item(source="handoff", memory_class="episodic", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.BOOT

    def test_episodic_does_not_override_open_loops(self, compiler: ContextCompiler):
        """type='commitment' (Rule 2) must beat memory_class='episodic' (Rule 5b)."""
        item = _make_item(type="commitment", memory_class="episodic", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.OPEN_LOOPS

    def test_episodic_does_not_override_procedures(self, compiler: ContextCompiler):
        """type='runbook' (Rule 3) must beat memory_class='episodic' (Rule 5b)."""
        item = _make_item(type="runbook", memory_class="episodic", created_at=_old_ts())
        assert _classify(compiler, item) == BundleType.PROCEDURES

    def test_non_episodic_old_item_not_recent_history(self, compiler: ContextCompiler):
        """memory_class='semantic' (non-episodic) old item → does NOT get recency boost."""
        item = _make_item(memory_class="semantic", created_at=_old_ts())
        assert _classify(compiler, item) != BundleType.RECENT_HISTORY


# ---------------------------------------------------------------------------
# Part C: Salience-aware budget allocation / ordering
# ---------------------------------------------------------------------------

class TestSalienceOrdering:
    """Tests that salience sorting happens correctly within bundles before compression."""

    def _compile_and_get_bundle(
        self,
        compiler: ContextCompiler,
        items: list[dict],
        bundle_type: BundleType,
        budget: int = 2000,
    ):
        """Compile items and return the matching bundle, or None."""
        result = compiler.compile(items, budget_tokens=budget)
        for b in result.bundles:
            if b.bundle_type == bundle_type:
                return b
        return None

    def test_high_salience_item_sorted_first_within_bundle(self, compiler: ContextCompiler):
        """Within EVIDENCE bundle, high-salience item must appear before low-salience."""
        # All evidence items (old, no special type) with different salience
        items = [
            _make_item(content="low salience content", created_at=_old_ts(), salience="low"),
            _make_item(content="medium salience content", created_at=_old_ts(), salience="medium"),
            _make_item(content="high salience content", created_at=_old_ts(), salience="high"),
        ]
        result = compiler.compile(items, budget_tokens=5000)
        evidence_bundle = next(
            (b for b in result.bundles if b.bundle_type == BundleType.EVIDENCE), None
        )
        assert evidence_bundle is not None, "Expected EVIDENCE bundle"
        assert len(evidence_bundle.items) > 0

        # The first item in the bundle should correspond to the high-salience item
        # (compression may change content slightly, but ordering is preserved)
        first_content = evidence_bundle.items[0].get("content", "")
        assert "high salience" in first_content, (
            f"Expected high-salience item first, got: {first_content!r}"
        )

    def test_salience_order_within_open_loops_bundle(self, compiler: ContextCompiler):
        """Within OPEN_LOOPS, high-salience commitment appears before medium/low."""
        items = [
            _make_item(
                content="low priority open loop",
                type="commitment",
                created_at=_old_ts(),
                salience="low",
            ),
            _make_item(
                content="high priority open loop",
                type="commitment",
                created_at=_old_ts(),
                salience="high",
            ),
        ]
        result = compiler.compile(items, budget_tokens=5000)
        open_loops = next(
            (b for b in result.bundles if b.bundle_type == BundleType.OPEN_LOOPS), None
        )
        assert open_loops is not None
        assert len(open_loops.items) >= 1
        first_content = open_loops.items[0].get("content", "")
        assert "high priority" in first_content, (
            f"Expected high-salience item first, got: {first_content!r}"
        )

    def test_items_without_salience_treated_as_medium(self, compiler: ContextCompiler):
        """Items without salience field are treated as medium (between high and low)."""
        items = [
            _make_item(content="no salience item", created_at=_old_ts()),
            _make_item(content="high salience item", created_at=_old_ts(), salience="high"),
            _make_item(content="low salience item", created_at=_old_ts(), salience="low"),
        ]
        result = compiler.compile(items, budget_tokens=5000)
        evidence_bundle = next(
            (b for b in result.bundles if b.bundle_type == BundleType.EVIDENCE), None
        )
        assert evidence_bundle is not None
        assert len(evidence_bundle.items) >= 1
        first_content = evidence_bundle.items[0].get("content", "")
        assert "high salience" in first_content

    def test_all_same_salience_no_error(self, compiler: ContextCompiler):
        """All items with same salience should not raise and should compile normally."""
        items = [
            _make_item(content=f"item {i}", created_at=_old_ts(), salience="medium")
            for i in range(5)
        ]
        result = compiler.compile(items, budget_tokens=5000)
        assert result is not None
        assert result.total_tokens >= 0

    def test_salience_sort_across_memory_class_items(self, compiler: ContextCompiler):
        """memory_class items sorted by salience within their assigned bundle."""
        items = [
            _make_item(
                content="low salience procedural",
                memory_class="procedural",
                created_at=_old_ts(),
                salience="low",
            ),
            _make_item(
                content="high salience procedural",
                memory_class="procedural",
                created_at=_old_ts(),
                salience="high",
            ),
        ]
        result = compiler.compile(items, budget_tokens=5000)
        procedures = next(
            (b for b in result.bundles if b.bundle_type == BundleType.PROCEDURES), None
        )
        assert procedures is not None
        first_content = procedures.items[0].get("content", "")
        assert "high salience" in first_content


# ---------------------------------------------------------------------------
# Integration: combined memory_class + salience
# ---------------------------------------------------------------------------

class TestMemoryClassIntegration:
    def test_multiple_memory_classes_route_to_correct_bundles(self, compiler: ContextCompiler):
        """Items with different memory_class values all route to their correct bundles."""
        items = [
            _make_item(content="open loop item", memory_class="open_loop", created_at=_old_ts()),
            _make_item(content="procedural item", memory_class="procedural", created_at=_old_ts()),
            _make_item(content="episodic item", memory_class="episodic", created_at=_old_ts()),
            _make_item(content="semantic item", memory_class="semantic", created_at=_old_ts()),
        ]
        result = compiler.compile(items, budget_tokens=8000)
        bundle_types = {b.bundle_type for b in result.bundles}

        assert BundleType.OPEN_LOOPS in bundle_types, "open_loop → OPEN_LOOPS missing"
        assert BundleType.PROCEDURES in bundle_types, "procedural → PROCEDURES missing"
        assert BundleType.RECENT_HISTORY in bundle_types, "episodic → RECENT_HISTORY missing"

    def test_memory_class_transient_falls_through_to_evidence(self, compiler: ContextCompiler):
        """memory_class='transient' has no special rule → falls to EVIDENCE."""
        item = _make_item(memory_class="transient", created_at=_old_ts())
        result = compiler._classify_item(item)
        assert result == BundleType.EVIDENCE

    def test_memory_class_semantic_falls_through_to_evidence(self, compiler: ContextCompiler):
        """memory_class='semantic' has no special rule → falls to EVIDENCE."""
        item = _make_item(memory_class="semantic", created_at=_old_ts())
        result = compiler._classify_item(item)
        assert result == BundleType.EVIDENCE

    def test_episodic_with_high_salience_in_recent_history(self, compiler: ContextCompiler):
        """Episodic + high salience should land in RECENT_HISTORY first."""
        items = [
            _make_item(
                content="low salience episodic",
                memory_class="episodic",
                created_at=_old_ts(),
                salience="low",
            ),
            _make_item(
                content="high salience episodic",
                memory_class="episodic",
                created_at=_old_ts(),
                salience="high",
            ),
        ]
        result = compiler.compile(items, budget_tokens=5000)
        rh = next(
            (b for b in result.bundles if b.bundle_type == BundleType.RECENT_HISTORY), None
        )
        assert rh is not None
        first_content = rh.items[0].get("content", "")
        assert "high salience" in first_content
