"""
Tests for cortex/hippo/scheduler.py — adaptive compaction scheduling.

Issue: #1191
"""

from __future__ import annotations

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
import pytest_asyncio

from cortex.hippo.config import HippoConfig
from cortex.hippo.engine import HippoEngine
from cortex.hippo.scheduler import (
    CompactionDecision,
    CompactionMode,
    CompactionScheduler,
)


# ── Helpers ──────────────────────────────────────────────────────────


class FakeClock:
    """Controllable clock for deterministic tests."""

    def __init__(self, start: float = 1000.0):
        self._now = start

    def __call__(self) -> float:
        return self._now

    def advance(self, seconds: float) -> None:
        self._now += seconds


def _make_scheduler(clock: FakeClock | None = None, **kwargs) -> CompactionScheduler:
    defaults = dict(
        idle_threshold_seconds=30.0,
        urgency_bypass=0.95,
        light_threshold=0.5,
        full_threshold=0.8,
    )
    defaults.update(kwargs)
    c = clock or FakeClock()
    return CompactionScheduler(**defaults, _clock=c)


# ── Unit tests: CompactionScheduler ──────────────────────────────────


def test_skip_when_not_urgent():
    """Low pressure + recent messages → SKIP."""
    clock = FakeClock()
    sched = _make_scheduler(clock)
    conv_id = 1

    # Record a message just now
    sched.record_message(conv_id)
    clock.advance(1.0)  # 1 second idle

    # Low pressure: 20% of budget
    decision = sched.evaluate(conv_id, current_tokens=2000, token_budget=10000)

    assert decision.mode == CompactionMode.SKIP
    assert decision.urgency < 0.5
    assert decision.token_pressure == pytest.approx(0.2)


def test_light_when_moderate():
    """Moderate pressure + some idle → LIGHT."""
    clock = FakeClock()
    sched = _make_scheduler(clock)
    conv_id = 1

    sched.record_message(conv_id)
    clock.advance(20.0)  # 20s idle → idle_factor = 20/30 ≈ 0.67

    # 75% pressure → urgency = 0.75 * 0.67 ≈ 0.50
    decision = sched.evaluate(conv_id, current_tokens=7500, token_budget=10000)

    assert decision.mode == CompactionMode.LIGHT
    assert 0.5 <= decision.urgency < 0.8


def test_full_when_idle_and_pressured():
    """High pressure + idle >30s → FULL."""
    clock = FakeClock()
    sched = _make_scheduler(clock)
    conv_id = 1

    sched.record_message(conv_id)
    clock.advance(60.0)  # 60s idle → idle_factor = 1.0

    # 85% pressure → urgency = 0.85 * 1.0 = 0.85 ≥ 0.8
    decision = sched.evaluate(conv_id, current_tokens=8500, token_budget=10000)

    assert decision.mode == CompactionMode.FULL
    assert decision.urgency >= 0.8


def test_force_safety_valve():
    """Pressure >=95% → FORCE regardless of idle."""
    clock = FakeClock()
    sched = _make_scheduler(clock)
    conv_id = 1

    sched.record_message(conv_id)
    clock.advance(5.0)  # very recent

    # 96% pressure → safety valve
    decision = sched.evaluate(conv_id, current_tokens=9600, token_budget=10000)

    assert decision.mode == CompactionMode.FORCE
    assert decision.urgency == 1.0
    assert "Safety valve" in decision.reason


def test_force_overrides_busy():
    """Even at 0s idle, 95%+ pressure → FORCE."""
    clock = FakeClock()
    sched = _make_scheduler(clock)
    conv_id = 1

    sched.record_message(conv_id)
    # Don't advance clock → 0s idle

    decision = sched.evaluate(conv_id, current_tokens=9500, token_budget=10000)

    assert decision.mode == CompactionMode.FORCE


def test_idle_factor_ramps():
    """Idle factor goes from 0 to 1 over threshold seconds."""
    clock = FakeClock()
    sched = _make_scheduler(clock, idle_threshold_seconds=30.0)
    conv_id = 1

    sched.record_message(conv_id)

    # At 0s idle → urgency ≈ 0 (idle_factor ≈ 0)
    decision_0 = sched.evaluate(conv_id, current_tokens=8000, token_budget=10000)

    # At 15s → idle_factor = 0.5 → urgency = 0.8 * 0.5 = 0.4
    clock.advance(15.0)
    decision_15 = sched.evaluate(conv_id, current_tokens=8000, token_budget=10000)

    # At 30s → idle_factor = 1.0 → urgency = 0.8 * 1.0 = 0.8
    clock.advance(15.0)
    decision_30 = sched.evaluate(conv_id, current_tokens=8000, token_budget=10000)

    # At 60s → idle_factor still 1.0 (clamped)
    clock.advance(30.0)
    decision_60 = sched.evaluate(conv_id, current_tokens=8000, token_budget=10000)

    assert decision_0.urgency < decision_15.urgency < decision_30.urgency
    assert decision_30.urgency == pytest.approx(decision_60.urgency)
    assert decision_30.urgency == pytest.approx(0.8, abs=0.01)


def test_message_cadence_tracking():
    """record_message updates cadence correctly."""
    clock = FakeClock()
    sched = _make_scheduler(clock)
    conv_id = 1

    # No messages → 0 cadence
    assert sched.get_cadence(conv_id) == 0.0

    # 1 message → still 0 (need ≥2 for rate)
    sched.record_message(conv_id)
    assert sched.get_cadence(conv_id) == 0.0

    # 2nd message 30s later → 2 msgs/min
    clock.advance(30.0)
    sched.record_message(conv_id)
    cadence = sched.get_cadence(conv_id)
    assert cadence == pytest.approx(2.0)

    # 3rd message 30s later → still ~2 msgs/min (3 msgs over 60s = 2/min)
    clock.advance(30.0)
    sched.record_message(conv_id)
    cadence = sched.get_cadence(conv_id)
    assert cadence == pytest.approx(2.0)


def test_cadence_window_pruning():
    """Old timestamps pruned from window."""
    clock = FakeClock()
    sched = _make_scheduler(clock)
    conv_id = 1

    # Record a burst
    for _ in range(5):
        sched.record_message(conv_id)
        clock.advance(1.0)

    # Jump past the 120s window
    clock.advance(200.0)

    # Record a new message — old ones should be pruned
    sched.record_message(conv_id)

    # Only 1 message in window now → cadence = 0
    assert sched.get_cadence(conv_id) == 0.0

    # Verify the deque was pruned
    dq = sched._message_times[conv_id]
    assert len(dq) == 1


def test_no_messages_gives_inf_idle():
    """No recorded messages → infinite idle → full if pressure is enough."""
    clock = FakeClock()
    sched = _make_scheduler(clock)
    conv_id = 99

    # No messages recorded, 85% pressure → idle=inf → idle_factor=1.0 → urgency=0.85
    decision = sched.evaluate(conv_id, current_tokens=8500, token_budget=10000)
    assert decision.mode == CompactionMode.FULL
    assert decision.idle_seconds == float("inf")


def test_zero_budget_no_crash():
    """Zero token budget should not crash."""
    clock = FakeClock()
    sched = _make_scheduler(clock)
    decision = sched.evaluate(1, current_tokens=100, token_budget=0)
    assert decision.token_pressure == 0.0
    assert decision.mode == CompactionMode.SKIP


# ── Config integration tests ─────────────────────────────────────────


def test_config_has_adaptive_fields():
    """Config has all adaptive scheduling fields."""
    cfg = HippoConfig()
    assert cfg.adaptive_scheduling_enabled is True
    assert cfg.compaction_idle_threshold_seconds == 30.0
    assert cfg.compaction_urgency_bypass == 0.95
    assert cfg.compaction_light_threshold == 0.5
    assert cfg.compaction_full_threshold == 0.8


@pytest.mark.asyncio
async def test_config_kill_switch():
    """adaptive_scheduling_enabled=False uses old behavior (no scheduler)."""
    config = HippoConfig(
        enabled=True,
        db_path=":memory:",
        adaptive_scheduling_enabled=False,
    )
    engine = HippoEngine(config=config)
    await engine.initialize(db_path=":memory:", owner_id="test")
    try:
        assert engine._scheduler is None
    finally:
        await engine.shutdown()


@pytest.mark.asyncio
async def test_config_enabled_creates_scheduler():
    """adaptive_scheduling_enabled=True creates scheduler."""
    config = HippoConfig(
        enabled=True,
        db_path=":memory:",
        adaptive_scheduling_enabled=True,
    )
    engine = HippoEngine(config=config)
    await engine.initialize(db_path=":memory:", owner_id="test")
    try:
        assert engine._scheduler is not None
        assert isinstance(engine._scheduler, CompactionScheduler)
    finally:
        await engine.shutdown()


# ── Engine integration tests ─────────────────────────────────────────


@pytest.mark.asyncio
async def test_engine_uses_scheduler():
    """Integration: after_turn respects scheduler decision (SKIP stops compaction)."""
    config = HippoConfig(
        enabled=True,
        db_path=":memory:",
        adaptive_scheduling_enabled=True,
        context_window_tokens=100,
        leaf_trigger_margin=0,
    )

    async def fake_summarize(text: str, **kwargs) -> str:
        return "summary"

    engine = HippoEngine(config=config, summarize_fn=fake_summarize)
    conv = await engine.initialize(db_path=":memory:", owner_id="test")

    # Ingest a few messages
    for i in range(3):
        await engine.ingest_message("user", f"message {i}")

    # Patch scheduler to return SKIP
    assert engine._scheduler is not None
    original_evaluate = engine._scheduler.evaluate
    engine._scheduler.evaluate = lambda *a, **kw: CompactionDecision(
        mode=CompactionMode.SKIP,
        urgency=0.1,
        token_pressure=0.1,
        idle_seconds=1.0,
        reason="test: forcing SKIP",
    )

    # Patch compaction engine to track calls
    assert engine._compaction_engine is not None
    compact_called = False
    original_compact = engine._compaction_engine.compact_full_sweep

    async def _track_compact(*a, **kw):
        nonlocal compact_called
        compact_called = True
        return await original_compact(*a, **kw)

    engine._compaction_engine.compact_full_sweep = _track_compact

    # Call after_turn with a budget — scheduler should SKIP
    await engine.after_turn(
        conversation_id=conv.conversation_id,
        messages=[{"role": "user", "content": "hello"}],
        token_budget=10000,
    )

    # Give async tasks a chance to run
    await asyncio.sleep(0.05)

    assert not compact_called, "Compaction should not run when scheduler says SKIP"

    await engine.shutdown()


@pytest.mark.asyncio
async def test_scheduler_with_budget():
    """Scheduler + budget work together: scheduler allows, budget blocks → no compaction."""
    config = HippoConfig(
        enabled=True,
        db_path=":memory:",
        adaptive_scheduling_enabled=True,
        context_window_tokens=100,
        leaf_trigger_margin=0,
    )

    async def fake_summarize(text: str, **kwargs) -> str:
        return "summary"

    engine = HippoEngine(config=config, summarize_fn=fake_summarize)
    conv = await engine.initialize(db_path=":memory:", owner_id="test")

    # Ingest messages
    for i in range(3):
        await engine.ingest_message("user", f"message {i}")

    # Patch scheduler to return FULL (allows compaction)
    assert engine._scheduler is not None
    engine._scheduler.evaluate = lambda *a, **kw: CompactionDecision(
        mode=CompactionMode.FULL,
        urgency=0.9,
        token_pressure=0.9,
        idle_seconds=60.0,
        reason="test: forcing FULL",
    )

    # Patch budget to block
    from cortex.hippo.budget import BudgetDecision as BDec

    assert engine._budget is not None
    engine._budget.check_budget = lambda sid: BDec(
        allowed=False, reason="budget_exceeded", remaining_tokens=0
    )

    # Track compaction
    compact_called = False
    original_compact = engine._compaction_engine.compact_full_sweep

    async def _track_compact(*a, **kw):
        nonlocal compact_called
        compact_called = True
        return await original_compact(*a, **kw)

    engine._compaction_engine.compact_full_sweep = _track_compact

    await engine.after_turn(
        conversation_id=conv.conversation_id,
        messages=[{"role": "user", "content": "hello"}],
        token_budget=10000,
    )

    await asyncio.sleep(0.05)

    assert not compact_called, "Budget block should prevent compaction even when scheduler allows"

    await engine.shutdown()


@pytest.mark.asyncio
async def test_light_compaction_single_round():
    """LIGHT mode fires compaction (uses compact_full_sweep with force=True)."""
    config = HippoConfig(
        enabled=True,
        db_path=":memory:",
        adaptive_scheduling_enabled=True,
        context_window_tokens=100,
        leaf_trigger_margin=0,
    )

    async def fake_summarize(text: str, **kwargs) -> str:
        return "summary"

    engine = HippoEngine(config=config, summarize_fn=fake_summarize)
    conv = await engine.initialize(db_path=":memory:", owner_id="test")

    for i in range(3):
        await engine.ingest_message("user", f"message {i}")

    # Patch scheduler to return LIGHT
    assert engine._scheduler is not None
    engine._scheduler.evaluate = lambda *a, **kw: CompactionDecision(
        mode=CompactionMode.LIGHT,
        urgency=0.6,
        token_pressure=0.6,
        idle_seconds=20.0,
        reason="test: forcing LIGHT",
    )

    # Track which compaction method was called and with what args
    compact_calls: list[dict] = []
    from cortex.hippo.compaction.types import CompactionResult as CEResult

    async def _track_full_sweep(*a, **kw):
        compact_calls.append({"method": "compact_full_sweep", "args": a, "kwargs": kw})
        return CEResult()

    engine._compaction_engine.compact_full_sweep = _track_full_sweep

    await engine.after_turn(
        conversation_id=conv.conversation_id,
        messages=[{"role": "user", "content": "hello"}],
        token_budget=10000,
    )

    await asyncio.sleep(0.05)

    assert len(compact_calls) == 1, "LIGHT should fire exactly one compaction call"
    assert compact_calls[0]["method"] == "compact_full_sweep"
    assert compact_calls[0]["kwargs"].get("force") is True

    await engine.shutdown()


@pytest.mark.asyncio
async def test_ingest_records_message_in_scheduler():
    """ingest_message() records in the scheduler."""
    config = HippoConfig(
        enabled=True,
        db_path=":memory:",
        adaptive_scheduling_enabled=True,
    )
    engine = HippoEngine(config=config)
    conv = await engine.initialize(db_path=":memory:", owner_id="test")

    assert engine._scheduler is not None
    assert engine._scheduler.get_cadence(conv.conversation_id) == 0.0

    await engine.ingest_message("user", "hello")
    await engine.ingest_message("assistant", "world")

    # 2 messages recorded — cadence should be calculable
    dq = engine._scheduler._message_times.get(conv.conversation_id)
    assert dq is not None
    assert len(dq) == 2

    await engine.shutdown()
