"""
Tests for CircadianEngine clean shutdown + async lifecycle (Issue #94)

Covers:
  - T01  Clean shutdown from AWAKE state (sessions_slept=1, tasks_cancelled)
  - T02  Clean shutdown from SLEEPING state (no emergency sleep triggered)
  - T03  Context manager __aenter__ / __aexit__ auto-shutdown
  - T04  Crash recovery — state persisted on shutdown, restored on startup
  - T05  Shutdown during dream cycle (active tasks cancelled)
  - T06  Double-shutdown is idempotent (no second sleep triggered)
  - T07  State persisted to DB on shutdown (get_config round-trip)
  - T08  resume_saved_state() ignores stale sleeping/dreaming states
  - T09  SIGTERM handler schedules shutdown task
  - T10  ShutdownReport fields populated correctly
  - T11  Shutdown with multiple active sessions (all emergency-slept)
  - T12  persist_state() survives storage failure (returns False, no raise)
  - T13  Context manager propagates exceptions after clean shutdown
  - T14  Engine in DROWSY state is emergency-slept on shutdown
"""

from __future__ import annotations

import asyncio
import json
import signal
from dataclasses import dataclass, field
from datetime import datetime, timezone, timedelta
from typing import Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch, call

import pytest

from cortex.circadian.engine import (
    CircadianEngine,
    CircadianState,
    ShutdownReport,
    _CIRCADIAN_STATE_KEY,
)
from cortex.types import SessionRecord, SemanticClaim


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _make_session(
    session_id: str = "sess-001",
    owner_id: str = "default",
    fatigue: float = 0.0,
    status: str = "awake",
    started_at: Optional[str] = None,
    ended_at: Optional[str] = None,
    compaction_count: int = 0,
) -> SessionRecord:
    return SessionRecord(
        id=f"rec-{session_id}",
        owner_id=owner_id,
        session_id=session_id,
        started_at=started_at or _now_iso(),
        ended_at=ended_at,
        status=status,
        fatigue_score=fatigue,
        compaction_count=compaction_count,
    )


# ---------------------------------------------------------------------------
# Extended MockStorageAdapter with config support
# ---------------------------------------------------------------------------

class MockStorageAdapter:
    """In-memory storage adapter for CircadianEngine shutdown tests."""

    def __init__(self):
        self._sessions: Dict[str, SessionRecord] = {}
        self._claims: Dict[str, List[SemanticClaim]] = {}
        self._orphans: List[SessionRecord] = []
        self._config: Dict[str, str] = {}

    # --- session ops ---

    async def create_session(self, session_id: str, owner_id: str = "default") -> SessionRecord:
        rec = _make_session(session_id=session_id, owner_id=owner_id)
        self._sessions[session_id] = rec
        return rec

    async def update_session(self, session_id: str, **kwargs) -> SessionRecord:
        rec = self._sessions.get(session_id)
        if rec is None:
            raise KeyError(f"Session not found: {session_id}")
        for k, v in kwargs.items():
            if hasattr(rec, k):
                setattr(rec, k, v)
        return rec

    async def get_active_session(self, owner_id: str = "default") -> Optional[SessionRecord]:
        for rec in reversed(list(self._sessions.values())):
            if rec.owner_id == owner_id and rec.status == "awake":
                return rec
        return None

    async def get_active_session_by_id(self, session_id: str) -> Optional[SessionRecord]:
        return self._sessions.get(session_id)

    async def get_orphaned_sessions(self, idle_seconds: int) -> List[SessionRecord]:
        return list(self._orphans)

    async def get_interrupted_sessions(
        self, owner_id: Optional[str] = None
    ) -> List[SessionRecord]:
        results = [
            s for s in self._sessions.values()
            if s.status in ("sleeping", "dreaming") and (not hasattr(s, "ended_at") or s.ended_at is None)
        ]
        if owner_id is not None:
            results = [s for s in results if s.owner_id == owner_id]
        return results

    async def get_session_claims(self, session_id: str) -> List[SemanticClaim]:
        return self._claims.get(session_id, [])

    # --- config ops (cortex_config table) ---

    async def get_config(self, key: str) -> Optional[str]:
        return self._config.get(key)

    async def set_config(self, key: str, value: str) -> None:
        self._config[key] = value

    # --- helpers for test setup ---

    def add_session(self, rec: SessionRecord) -> None:
        self._sessions[rec.session_id] = rec

    def add_claim(self, session_id: str, n: int = 1) -> None:
        claims = self._claims.setdefault(session_id, [])
        for i in range(n):
            claims.append(SemanticClaim(
                id=f"claim-{session_id}-{i}",
                owner_id="default",
                subject="test",
                predicate="has",
                object_value="value",
                source_session=session_id,
            ))


# ---------------------------------------------------------------------------
# Engine factory
# ---------------------------------------------------------------------------

def _make_engine(storage: Optional[MockStorageAdapter] = None) -> tuple[CircadianEngine, MockStorageAdapter]:
    """Return a CircadianEngine wired to a MockStorageAdapter."""
    if storage is None:
        storage = MockStorageAdapter()

    config = MagicMock()
    config.session_fatigue_nap_threshold = 0.7
    config.session_fatigue_sleep_threshold = 0.9
    config.dream_cycle_idle_timeout_sec = 3600
    config.dream_cycle_hour_utc = 8
    config.dream_cycle_enabled = False  # Don't auto-start dream scheduler in tests
    config.owner_id = "default"

    llm = MagicMock()
    llm.complete = AsyncMock(return_value=MagicMock(text="[test nap summary]"))

    cornerstone_guardian = MagicMock()
    cornerstone_guardian.load_all = AsyncMock(return_value=[])

    engine = CircadianEngine(
        storage=storage,
        config=config,
        llm=llm,
        cornerstone_guardian=cornerstone_guardian,
    )
    return engine, storage


# ---------------------------------------------------------------------------
# T01 — Clean shutdown from AWAKE state
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_shutdown_from_awake_state():
    """T01: shutdown() from AWAKE triggers emergency sleep and sets state to SHUTDOWN."""
    engine, storage = _make_engine()

    session_id = "sess-shutdown-awake"
    storage.add_session(_make_session(session_id=session_id, status="awake"))
    engine._states[session_id] = CircadianState.AWAKE

    report = await engine.shutdown(triggered_by="test")

    assert isinstance(report, ShutdownReport)
    assert report.sessions_slept == 1, "Active session should have been emergency-slept"
    assert engine._states[session_id] == CircadianState.SHUTDOWN
    assert engine._is_shutdown is True
    assert report.triggered_by == "test"


# ---------------------------------------------------------------------------
# T02 — Clean shutdown from SLEEPING state
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_shutdown_from_sleeping_state():
    """T02: shutdown() skips emergency sleep for sessions already sleeping."""
    engine, storage = _make_engine()

    session_id = "sess-already-asleep"
    storage.add_session(_make_session(session_id=session_id, status="sleeping"))
    engine._states[session_id] = CircadianState.SLEEPING

    report = await engine.shutdown(triggered_by="test")

    assert report.sessions_slept == 0, "SLEEPING session should not trigger emergency sleep"
    assert engine._is_shutdown is True


# ---------------------------------------------------------------------------
# T03 — Context manager auto-shutdown
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_context_manager_auto_shutdown():
    """T03: async with CircadianEngine(...) auto-calls shutdown() on exit."""
    engine, storage = _make_engine()

    session_id = "sess-ctx-mgr"
    storage.add_session(_make_session(session_id=session_id, status="awake"))
    engine._states[session_id] = CircadianState.AWAKE

    async with engine:
        # Engine is live inside the block
        assert engine._is_shutdown is False

    # After the block, engine should be shut down
    assert engine._is_shutdown is True
    assert engine._states[session_id] == CircadianState.SHUTDOWN


# ---------------------------------------------------------------------------
# T04 — Crash recovery: state persisted and restored on restart
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_crash_recovery_state_persisted_and_restored():
    """T04: state persisted to DB on shutdown; resumed on next startup."""
    storage = MockStorageAdapter()
    engine, _ = _make_engine(storage)

    session_id = "sess-crash"
    storage.add_session(_make_session(session_id=session_id, status="awake"))
    engine._states[session_id] = CircadianState.AWAKE

    # Shutdown persists state
    report = await engine.shutdown()
    assert report.state_persisted is True

    # Verify JSON written to DB
    raw = await storage.get_config(_CIRCADIAN_STATE_KEY)
    assert raw is not None
    payload = json.loads(raw)
    assert "states" in payload
    # After shutdown all states are SHUTDOWN
    assert payload["states"][session_id] == CircadianState.SHUTDOWN.value

    # New engine on same storage — resume_saved_state() should find AWAKE states
    # Set up a scenario where state was AWAKE (simulate pre-shutdown persist)
    awake_payload = {
        "states": {session_id: CircadianState.AWAKE.value},
        "persisted_at": _now_iso(),
    }
    await storage.set_config(_CIRCADIAN_STATE_KEY, json.dumps(awake_payload))

    engine2, _ = _make_engine(storage)
    restored = await engine2.resume_saved_state()

    assert session_id in restored
    assert restored[session_id] == CircadianState.AWAKE
    assert engine2._states[session_id] == CircadianState.AWAKE


# ---------------------------------------------------------------------------
# T05 — Shutdown during dream cycle cancels background tasks
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_shutdown_during_dream_cycle():
    """T05: shutdown() cancels the dream scheduler task if it's running."""
    engine, storage = _make_engine()

    # Manually start a fake dream scheduler task that sleeps forever
    async def _fake_dream_loop():
        await asyncio.sleep(9999)

    engine._dream_scheduler_task = asyncio.create_task(_fake_dream_loop())

    assert not engine._dream_scheduler_task.done()

    report = await engine.shutdown()

    assert report.tasks_cancelled >= 1
    assert engine._dream_scheduler_task.done()
    assert engine._dream_scheduler_task.cancelled()


# ---------------------------------------------------------------------------
# T06 — Double-shutdown is idempotent
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_double_shutdown_is_idempotent():
    """T06: calling shutdown() twice doesn't raise and doesn't double-sleep sessions."""
    engine, storage = _make_engine()

    session_id = "sess-double"
    storage.add_session(_make_session(session_id=session_id, status="awake"))
    engine._states[session_id] = CircadianState.AWAKE

    report1 = await engine.shutdown()
    report2 = await engine.shutdown()

    assert report1.sessions_slept == 1
    assert report2.sessions_slept == 0  # Second call is a no-op
    assert engine._is_shutdown is True


# ---------------------------------------------------------------------------
# T07 — State persisted to DB on shutdown (get_config round-trip)
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_state_persisted_to_db():
    """T07: shutdown() writes state to cortex_config via set_config."""
    storage = MockStorageAdapter()
    engine, _ = _make_engine(storage)

    engine._states["s1"] = CircadianState.AWAKE
    storage.add_session(_make_session(session_id="s1", status="awake"))

    await engine.shutdown()

    raw = await storage.get_config(_CIRCADIAN_STATE_KEY)
    assert raw is not None

    payload = json.loads(raw)
    assert "states" in payload
    assert "persisted_at" in payload
    assert "s1" in payload["states"]


# ---------------------------------------------------------------------------
# T08 — resume_saved_state() ignores stale sleeping/dreaming states
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_resume_saved_state_ignores_sleeping():
    """T08: sleeping/dreaming sessions from saved state are NOT restored (handled by orphan recovery)."""
    storage = MockStorageAdapter()
    engine, _ = _make_engine(storage)

    payload = {
        "states": {
            "sess-awake": CircadianState.AWAKE.value,
            "sess-sleeping": CircadianState.SLEEPING.value,
            "sess-dreaming": CircadianState.DREAMING.value,
        },
        "persisted_at": _now_iso(),
    }
    await storage.set_config(_CIRCADIAN_STATE_KEY, json.dumps(payload))

    restored = await engine.resume_saved_state()

    assert "sess-awake" in restored
    assert "sess-sleeping" not in restored
    assert "sess-dreaming" not in restored


# ---------------------------------------------------------------------------
# T09 — SIGTERM handler schedules shutdown task
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_sigterm_handler_schedules_shutdown():
    """T09: registering signal handlers and sending SIGTERM calls shutdown."""
    engine, storage = _make_engine()

    # We can't actually send SIGTERM in tests, so we verify the handler is registered
    # and that calling it directly schedules a shutdown task.
    loop = asyncio.get_event_loop()

    with patch.object(loop, "add_signal_handler") as mock_add:
        engine.register_signal_handlers()
        # Called for SIGTERM and SIGINT
        assert mock_add.call_count == 2
        calls_signals = {c.args[0] for c in mock_add.call_args_list}
        assert signal.SIGTERM in calls_signals
        assert signal.SIGINT in calls_signals

    assert engine._signal_handlers_registered is True


# ---------------------------------------------------------------------------
# T10 — ShutdownReport fields populated correctly
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_shutdown_report_fields():
    """T10: ShutdownReport has all expected fields with correct types."""
    engine, storage = _make_engine()

    session_id = "sess-report"
    storage.add_session(_make_session(session_id=session_id, status="awake"))
    engine._states[session_id] = CircadianState.AWAKE

    # Add a background idle watcher to cancel
    async def _fake_watcher():
        await asyncio.sleep(9999)

    engine._idle_watcher_task = asyncio.create_task(_fake_watcher())

    report = await engine.shutdown(triggered_by="test_report")

    assert isinstance(report.sessions_slept, int)
    assert isinstance(report.tasks_cancelled, int)
    assert isinstance(report.state_persisted, bool)
    assert isinstance(report.triggered_by, str)
    assert isinstance(report.completed_at, str)
    assert report.triggered_by == "test_report"
    assert report.sessions_slept == 1
    assert report.tasks_cancelled >= 1
    # completed_at should be a valid ISO datetime string
    datetime.fromisoformat(report.completed_at)


# ---------------------------------------------------------------------------
# T11 — Multiple active sessions all emergency-slept
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_shutdown_multiple_active_sessions():
    """T11: all AWAKE and DROWSY sessions are emergency-slept on shutdown."""
    engine, storage = _make_engine()

    sessions = ["multi-1", "multi-2", "multi-3"]
    for sid in sessions:
        storage.add_session(_make_session(session_id=sid, status="awake"))
        engine._states[sid] = CircadianState.AWAKE

    # Make one DROWSY
    engine._states["multi-2"] = CircadianState.DROWSY

    report = await engine.shutdown()

    assert report.sessions_slept == 3
    for sid in sessions:
        assert engine._states[sid] == CircadianState.SHUTDOWN


# ---------------------------------------------------------------------------
# T12 — persist_state() survives storage failure
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_persist_state_survives_storage_failure():
    """T12: persist_state() returns False on storage error, does not raise."""
    engine, storage = _make_engine()

    # Break set_config to raise an exception
    async def _broken_set_config(key: str, value: str) -> None:
        raise RuntimeError("DB connection lost")

    storage.set_config = _broken_set_config

    result = await engine.persist_state()
    assert result is False


# ---------------------------------------------------------------------------
# T13 — Context manager propagates exceptions after clean shutdown
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_context_manager_propagates_exception():
    """T13: exceptions inside 'async with' block are re-raised after shutdown."""
    engine, storage = _make_engine()

    with pytest.raises(ValueError, match="intentional"):
        async with engine:
            raise ValueError("intentional test error")

    # Even though an exception occurred, engine should be shut down
    assert engine._is_shutdown is True


# ---------------------------------------------------------------------------
# T14 — DROWSY session is emergency-slept on shutdown
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_shutdown_drowsy_session():
    """T14: a DROWSY session (not just AWAKE) is treated as active and emergency-slept."""
    engine, storage = _make_engine()

    session_id = "sess-drowsy"
    storage.add_session(_make_session(session_id=session_id, status="awake"))
    engine._states[session_id] = CircadianState.DROWSY

    report = await engine.shutdown()

    assert report.sessions_slept == 1
    assert engine._states[session_id] == CircadianState.SHUTDOWN
