"""
MODULE 6a: CIRCADIAN ENGINE — Scheduler + Trigger Logic
=========================================================
Manages the session lifecycle state machine: AWAKE → DROWSY → SLEEPING → DREAMING → AWAKE.

Responsibilities (M6a):
  - wake()           — create session record, load cornerstones, handle orphans
  - update_fatigue() — track session fatigue from context ratio / events / time
  - should_nap()     — threshold check against nap_threshold
  - nap()            — mid-session mini-consolidation, reset fatigue to 0.4
  - get_fatigue()    — current fatigue accessor
  - sleep()          — stub (consolidation logic in M6b); triggers state transition
  - dream()          — stub (deep consolidation in M6b); triggers state transition
  - shutdown()       — clean shutdown: cancel tasks, emergency-sleep active sessions,
                       persist state to DB for crash recovery
  - Asyncio background scheduler for auto-sleep after idle_timeout
  - Asyncio background scheduler for dream cycle at configured UTC hour
  - Manual trigger support for forced dream/sleep cycles
  - wake() / sleep() lifecycle hooks
  - Context manager support (async with CircadianEngine(...) as engine: ...)
  - SIGTERM / SIGINT signal handling for graceful OS-level shutdown

Design principles:
  - All public methods are async
  - State transitions logged to session_record
  - Orphaned session detection and recovery at wake time
  - Fatigue score: 0.0 (fresh) → 1.0 (exhausted)
  - Shutdown state persisted to cortex_config for crash recovery
"""

from __future__ import annotations

import asyncio
import json
import logging
import signal
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import TYPE_CHECKING, Any, List, Optional

if TYPE_CHECKING:
    from cortex.storage.adapter import StorageAdapter
    from cortex.config import CortexConfig
    from cortex.llm.provider import LLMProvider
    from cortex.identity.cornerstone import CornerstoneGuardian, DriftReport
    from cortex.identity.drift_calculator import DriftCalculator
    from cortex.deprecation.pipeline import DeprecationPipeline
    from cortex.core.entity_resolution import EntityResolver
    from cortex.core.feedback import FeedbackSystem
    from cortex.core.promotion import PromotionScorer
    from cortex.circadian.sleep import SleepConsolidator

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# State machine
# ---------------------------------------------------------------------------

class CircadianState(str, Enum):
    """Session lifecycle states."""
    AWAKE = "awake"
    DROWSY = "drowsy"
    SLEEPING = "sleeping"
    DREAMING = "dreaming"
    WAKING = "waking"
    SHUTDOWN = "shutdown"


# Key used in cortex_config table for persisted shutdown state
_CIRCADIAN_STATE_KEY = "circadian_state"


# ---------------------------------------------------------------------------
# Report dataclasses
# ---------------------------------------------------------------------------

@dataclass
class WakeReport:
    """Result of a wake() call."""
    session_id: str
    cornerstones_loaded: int
    cornerstone_tokens: int
    context_retrieved: Optional[str]
    orphans_recovered: int
    started_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


@dataclass
class NapReport:
    """Result of a nap() mini-consolidation."""
    session_id: str
    fatigue_before: float
    fatigue_after: float = 0.4
    claims_flushed: int = 0
    summary: str = ""
    completed_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


@dataclass
class SleepReport:
    """Result of a sleep() call (full consolidation in M6b; stub here)."""
    session_id: str
    duration_ms: int = 0
    claims_created: int = 0
    claims_promoted: int = 0
    claims_pruned: int = 0
    drift_checks: List["DriftReport"] = field(default_factory=list)
    summary: str = ""
    completed_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


@dataclass
class DreamReport:
    """Result of a dream() cycle (deep consolidation in M6b; stub here)."""
    owner_id: str
    duration_ms: int = 0
    orphans_recovered: int = 0
    claims_consolidated: int = 0
    contradictions_resolved: int = 0
    claims_decayed: int = 0
    claims_archived: int = 0
    cornerstone_proposals: int = 0
    entity_merge_suggestions: int = 0
    brain_graphs_refreshed: int = 0
    feedback_applied: int = 0
    journal: str = ""
    completed_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


@dataclass
class ShutdownReport:
    """Result of a shutdown() call."""
    sessions_slept: int = 0
    tasks_cancelled: int = 0
    state_persisted: bool = False
    triggered_by: str = "manual"
    completed_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


# ---------------------------------------------------------------------------
# Fatigue constants
# ---------------------------------------------------------------------------

_FATIGUE_TIME_INCREMENT = 0.01     # +0.01 per interval
# Default value for backward-compatibility; actual value comes from CortexConfig.fatigue_interval_sec
_FATIGUE_TIME_INTERVAL_SEC = 300
_FATIGUE_COMPACTION_BUMP = 0.15    # +0.15 per compaction event
_FATIGUE_ERROR_BUMP = 0.05         # +0.05 per error/retry
_FATIGUE_AFTER_NAP = 0.4           # reset fatigue to this after nap
_FATIGUE_MAX = 1.0


# ---------------------------------------------------------------------------
# CircadianEngine
# ---------------------------------------------------------------------------

class CircadianEngine:
    """
    Session lifecycle manager implementing the AWAKE → DROWSY → SLEEPING → DREAMING → AWAKE
    state machine.

    Args:
        storage:              StorageAdapter for persistence.
        config:               CortexConfig with circadian parameters.
        llm:                  LLMProvider for generating summaries.
        cornerstone_guardian: CornerstoneGuardian for loading/checking cornerstones.
    """

    def __init__(
        self,
        storage: "StorageAdapter",
        config: "CortexConfig",
        llm: "LLMProvider",
        cornerstone_guardian: "CornerstoneGuardian",
        deprecation_pipeline: Optional["DeprecationPipeline"] = None,
        drift_calculator: Optional["DriftCalculator"] = None,
        entity_resolver: Optional["EntityResolver"] = None,
        feedback_system: Optional["FeedbackSystem"] = None,
        event_emitter: Optional[Any] = None,
        peer_manager: Optional[Any] = None,
        promotion_scorer: Optional["PromotionScorer"] = None,
    ) -> None:
        self.storage = storage
        self.config = config
        self.llm = llm
        self.cornerstone_guardian = cornerstone_guardian
        self.deprecation_pipeline = deprecation_pipeline
        self.drift_calculator = drift_calculator
        self.entity_resolver = entity_resolver
        self.feedback_system = feedback_system
        self.event_emitter = event_emitter
        self.peer_manager = peer_manager
        self.promotion_scorer = promotion_scorer

        # Canonical SleepPipeline — always created (S1→S2→S3→S4 consolidation).
        # Replaces the old per-stage duplication that lived in SleepConsolidator.sleep().
        from cortex.sleep.pipeline import SleepPipeline
        self._sleep_pipeline = SleepPipeline(
            storage=storage,
            llm=llm,
            config=config,
            promotion_scorer=promotion_scorer,
        )

        # Create SleepConsolidator if required deps are available.
        # Still used for dream() deep-consolidation (M6b) which needs
        # DeprecationPipeline + DriftCalculator.
        self._consolidator: Optional["SleepConsolidator"] = None
        if deprecation_pipeline is not None and drift_calculator is not None:
            from cortex.circadian.sleep import SleepConsolidator
            self._consolidator = SleepConsolidator(
                storage=storage,
                config=config,
                llm=llm,
                cornerstone_guardian=cornerstone_guardian,
                drift_calculator=drift_calculator,
                deprecation_pipeline=deprecation_pipeline,
                entity_resolver=entity_resolver,
                feedback_system=feedback_system,
                event_emitter=event_emitter,
                peer_manager=peer_manager,
            )

        # Config helpers
        self._nap_threshold: float = getattr(config, "session_fatigue_nap_threshold", 0.7)
        self._sleep_threshold: float = getattr(config, "session_fatigue_sleep_threshold", 0.9)
        self._idle_timeout_sec: int = getattr(config, "dream_cycle_idle_timeout_sec", 3600)
        self._dream_hour_utc: int = getattr(config, "dream_cycle_hour_utc", 8)
        self._dream_enabled: bool = getattr(config, "dream_cycle_enabled", True)
        self._owner_id: str = getattr(config, "owner_id", "default")

        # In-memory state cache: session_id → state
        self._states: dict[str, CircadianState] = {}
        # Lock protecting all mutations to _states
        self._states_lock: asyncio.Lock = asyncio.Lock()

        # Last activity timestamps: session_id → epoch seconds (float)
        self._last_activity: dict[str, float] = {}

        # Fatigue interval from config (seconds between time-increment ticks)
        _fi = getattr(config, "fatigue_interval_sec", None)
        self._fatigue_interval_sec: int = int(_fi) if isinstance(_fi, (int, float)) else _FATIGUE_TIME_INTERVAL_SEC

        # Background task handles
        self._idle_watcher_task: Optional[asyncio.Task] = None
        self._dream_scheduler_task: Optional[asyncio.Task] = None

        # Lifecycle callbacks (hooks)
        self._wake_hooks: list = []
        self._sleep_hooks: list = []

        # Shutdown guard — prevents double-shutdown
        self._is_shutdown: bool = False
        self._shutdown_lock: asyncio.Lock = asyncio.Lock()

        # Concurrency guards — prevent duplicate dream/sleep executions (#477)
        self._dream_lock: asyncio.Lock = asyncio.Lock()
        self._sleep_lock: asyncio.Lock = asyncio.Lock()

        # Signal handlers registered flag
        self._signal_handlers_registered: bool = False

    # ------------------------------------------------------------------
    # Lifecycle hooks registration
    # ------------------------------------------------------------------

    def on_wake(self, fn) -> None:
        """Register a callable invoked after wake(). Signature: fn(session_id: str)."""
        self._wake_hooks.append(fn)

    def on_sleep(self, fn) -> None:
        """Register a callable invoked after sleep(). Signature: fn(session_id: str)."""
        self._sleep_hooks.append(fn)

    async def _run_hooks(self, hooks: list, session_id: str) -> None:
        for hook in hooks:
            try:
                result = hook(session_id)
                if asyncio.iscoroutine(result):
                    await result
            except Exception as exc:
                logger.warning("Hook %s raised: %s", hook, exc)

    # ------------------------------------------------------------------
    # Session lifecycle — public API (all async)
    # ------------------------------------------------------------------

    async def wake(
        self,
        session_id: str,
        owner_id: str = "default",
    ) -> WakeReport:
        """
        Start a new session ("day").

        1. Check for orphaned sessions → run sleep() on them first.
        2. Create session_record with status='awake'.
        3. Load cornerstones.
        4. Retrieve any pending context from last session.
        5. Set fatigue_score = 0.0.
        6. Start background idle watcher.
        7. Fire wake hooks.

        Returns:
            WakeReport with cornerstones + context summary.
        """
        logger.info("[circadian] wake() — session=%s owner=%s", session_id, owner_id)

        # 0. Register signal handlers (idempotent — only registers once per engine)
        self.register_signal_handlers()

        # 0b. Resume any saved state from previous shutdown (crash recovery)
        await self.resume_saved_state()

        # 1. Orphan recovery: find sessions that were never slept
        orphans_recovered = await self._recover_orphans()

        # 1b. Interrupted consolidation recovery: close sessions stuck mid-sleep/mid-dream
        orphans_recovered += await self._recover_interrupted_sessions(owner_id=owner_id)

        # 2. Create session record
        session = await self.storage.create_session(session_id, owner_id=owner_id)
        logger.debug("[circadian] session_record created: %s status=%s", session_id, session.status)

        # 3. Update in-memory state
        async with self._states_lock:
            self._states[session_id] = CircadianState.AWAKE
        self._last_activity[session_id] = datetime.now(timezone.utc).timestamp()

        # 4. Load cornerstones
        cornerstones = await self.cornerstone_guardian.load_all(owner_id=owner_id)
        cornerstone_tokens = sum(getattr(cs, "token_count", 0) for cs in cornerstones)
        logger.debug("[circadian] cornerstones loaded: %d (%d tokens)", len(cornerstones), cornerstone_tokens)

        # 5. Retrieve pending context from previous session (best-effort)
        context_retrieved: Optional[str] = None
        try:
            context_retrieved = await self._retrieve_pending_context(owner_id)
        except Exception as exc:
            logger.warning("[circadian] context retrieval failed: %s", exc)

        # 6. Start background schedulers (idempotent)
        self._ensure_idle_watcher()
        if self._dream_enabled:
            self._ensure_dream_scheduler()

        # 7. Emit SESSION_STARTED event
        if self.event_emitter is not None:
            try:
                from cortex.events.emitter import CortexEvent, EventType
                await self.event_emitter.emit(CortexEvent(
                    event_type=EventType.SESSION_STARTED,
                    owner_id=owner_id,
                    payload={"session_id": session_id, "cornerstones_loaded": len(cornerstones)},
                    session_id=session_id,
                ))
            except Exception as exc:
                logger.debug("[circadian] event emit failed (SESSION_STARTED): %s", exc)

        # 8. Fire wake hooks
        await self._run_hooks(self._wake_hooks, session_id)

        return WakeReport(
            session_id=session_id,
            cornerstones_loaded=len(cornerstones),
            cornerstone_tokens=cornerstone_tokens,
            context_retrieved=context_retrieved,
            orphans_recovered=orphans_recovered,
        )

    async def update_fatigue(
        self,
        session_id: str,
        context_fill_ratio: float,
        compaction_event: bool = False,
        error_event: bool = False,
    ) -> float:
        """
        Update fatigue score based on:
          - context_fill_ratio (0.0–1.0): how full the context window is
          - compaction_event:  +0.15 per compaction
          - error_event:       +0.05 per error/retry
          - Time decay:        +0.01 per 5-minute interval since session start

        Returns:
            Updated fatigue score (clamped to [0.0, 1.0]).
        """
        session = await self._get_session(session_id)
        if session is None:
            logger.warning("[circadian] update_fatigue: unknown session %s", session_id)
            return 0.0

        current_fatigue = session.fatigue_score

        # Context ratio contribution (direct mapping)
        new_fatigue = context_fill_ratio

        # Event bumps
        if compaction_event:
            new_fatigue += _FATIGUE_COMPACTION_BUMP
        if error_event:
            new_fatigue += _FATIGUE_ERROR_BUMP

        # Time decay: +0.01 per 5-minute interval of session duration
        session_start = datetime.fromisoformat(session.started_at)
        now_utc = datetime.now(timezone.utc)
        # Ensure both are timezone-aware for comparison
        if session_start.tzinfo is None:
            session_start = session_start.replace(tzinfo=timezone.utc)
        elapsed_sec = (now_utc - session_start).total_seconds()
        intervals = int(elapsed_sec // self._fatigue_interval_sec)
        new_fatigue += intervals * _FATIGUE_TIME_INCREMENT

        # Clamp
        new_fatigue = min(_FATIGUE_MAX, max(0.0, new_fatigue))

        # Persist compaction count if event happened
        update_kwargs: dict = {"fatigue_score": new_fatigue}
        if compaction_event:
            update_kwargs["compaction_count"] = (session.compaction_count or 0) + 1

        await self.storage.update_session(session_id, **update_kwargs)

        # Update in-memory state to DROWSY if approaching thresholds
        async with self._states_lock:
            if new_fatigue >= self._nap_threshold and self._states.get(session_id) == CircadianState.AWAKE:
                self._states[session_id] = CircadianState.DROWSY
                logger.info("[circadian] session %s → DROWSY (fatigue=%.2f)", session_id, new_fatigue)

        # Touch activity timestamp
        self._last_activity[session_id] = datetime.now(timezone.utc).timestamp()

        logger.debug(
            "[circadian] fatigue updated session=%s ratio=%.2f compact=%s error=%s → %.2f",
            session_id, context_fill_ratio, compaction_event, error_event, new_fatigue,
        )
        return new_fatigue

    async def should_nap(self, session_id: str) -> bool:
        """
        Return True if the current fatigue score exceeds the nap threshold.

        Threshold is configurable via CortexConfig.session_fatigue_nap_threshold
        (default 0.7).
        """
        fatigue = await self.get_fatigue(session_id)
        return fatigue > self._nap_threshold

    async def nap(self, session_id: str) -> NapReport:
        """
        Mid-session mini-consolidation:

        1. Flush working memory (session claims) to durable storage.
        2. Generate brief session summary (via LLM).
        3. Reset fatigue to 0.4.
        4. Set status back to 'awake' (post-nap).

        Returns:
            NapReport with fatigue before/after and claims flushed.
        """
        logger.info("[circadian] nap() — session=%s", session_id)
        session = await self._get_session(session_id)
        if session is None:
            raise ValueError(f"Cannot nap: session {session_id!r} not found")

        fatigue_before = session.fatigue_score

        # 1. Flush working memory (session claims)
        claims = await self.storage.get_session_claims(session_id)
        claims_flushed = len(claims)
        logger.debug("[circadian] nap: flushed %d claims", claims_flushed)

        # 2. Generate summary
        summary = await self._generate_nap_summary(session_id, claims_count=claims_flushed)

        # 3. Reset fatigue to 0.4 and mark status='awake'
        await self.storage.update_session(
            session_id,
            fatigue_score=_FATIGUE_AFTER_NAP,
            status=CircadianState.AWAKE.value,
            summary=summary,
        )
        async with self._states_lock:
            self._states[session_id] = CircadianState.AWAKE
        self._last_activity[session_id] = datetime.now(timezone.utc).timestamp()

        logger.info("[circadian] nap complete session=%s fatigue %.2f→%.2f",
                    session_id, fatigue_before, _FATIGUE_AFTER_NAP)

        return NapReport(
            session_id=session_id,
            fatigue_before=fatigue_before,
            fatigue_after=_FATIGUE_AFTER_NAP,
            claims_flushed=claims_flushed,
            summary=summary,
        )

    async def get_fatigue(self, session_id: str) -> float:
        """
        Return the current fatigue score for the session.

        Returns 0.0 if the session is unknown.
        """
        session = await self._get_session(session_id)
        if session is None:
            return 0.0
        return session.fatigue_score

    async def sleep(
        self,
        session_id: str,
        *,
        triggered_by: str = "idle_timeout",
    ) -> SleepReport:
        """
        End-of-session consolidation trigger.

        M6a responsibility: transition state SLEEPING, log the event, fire sleep hooks.
        Full consolidation (memory dump, deprecation review, etc.) is implemented in M6b.

        Args:
            session_id:    The session to sleep.
            triggered_by:  Reason string logged to session_record ('idle_timeout',
                           'manual', 'dream_cycle').

        Returns:
            SleepReport (stub — consolidation fields populated by M6b).
        """
        logger.info("[circadian] sleep() — session=%s trigger=%s", session_id, triggered_by)
        start_ts = datetime.now(timezone.utc)

        # Concurrency guard — reject duplicate concurrent sleep() calls (#477).
        # In asyncio (cooperative), the locked() check + async with is race-free
        # because no await point exists between the check and the acquire.
        if self._sleep_lock.locked():
            logger.warning(
                "[circadian] sleep() concurrent call rejected for session=%s — already sleeping",
                session_id,
            )
            return SleepReport(session_id=session_id)

        async with self._sleep_lock:
            session = await self._get_session(session_id)
            if session is None:
                logger.warning("[circadian] sleep: session %s not found", session_id)
                return SleepReport(session_id=session_id)

            # Transition: → SLEEPING
            async with self._states_lock:
                self._states[session_id] = CircadianState.SLEEPING
            await self.storage.update_session(
                session_id,
                status=CircadianState.SLEEPING.value,
            )

            # Emit SESSION_ENDED event
            if self.event_emitter is not None:
                try:
                    from cortex.events.emitter import CortexEvent, EventType
                    await self.event_emitter.emit(CortexEvent(
                        event_type=EventType.SESSION_ENDED,
                        owner_id=self._owner_id,
                        payload={"session_id": session_id, "triggered_by": triggered_by},
                        session_id=session_id,
                    ))
                except Exception as exc:
                    logger.debug("[circadian] event emit failed (SESSION_ENDED): %s", exc)

            # Fire sleep hooks
            await self._run_hooks(self._sleep_hooks, session_id)

            # Delegate actual sleep work to the canonical SleepPipeline (S1→S2→S3→S4).
            # Each stage is internally non-fatal; errors are recorded in pipeline_report.errors.
            pipeline_report = None
            try:
                pipeline_report = await self._sleep_pipeline.run(
                    session_id=session_id,
                    owner_id=self._owner_id,
                )
                if pipeline_report.errors:
                    logger.warning(
                        "[circadian] sleep pipeline finished with stage errors session=%s errors=%s",
                        session_id,
                        list(pipeline_report.errors.keys()),
                    )
                else:
                    logger.info(
                        "[circadian] sleep pipeline complete session=%s promoted=%d",
                        session_id,
                        pipeline_report.claims_promoted,
                    )
            except Exception as exc:
                logger.error("[circadian] sleep pipeline failed for %s: %s", session_id, exc)

            # Mark ended_at and compute duration
            end_ts = datetime.now(timezone.utc)
            session_start = datetime.fromisoformat(session.started_at)
            if session_start.tzinfo is None:
                session_start = session_start.replace(tzinfo=timezone.utc)
            duration_ms = int((end_ts - session_start).total_seconds() * 1000)

            await self.storage.update_session(
                session_id,
                ended_at=end_ts.isoformat(),
                status=CircadianState.SLEEPING.value,
            )

            # Emit SLEEP_COMPLETED event
            if self.event_emitter is not None:
                try:
                    from cortex.events.emitter import CortexEvent, EventType
                    await self.event_emitter.emit(CortexEvent(
                        event_type=EventType.SLEEP_COMPLETED,
                        owner_id=self._owner_id,
                        payload={
                            "session_id": session_id,
                            "duration_ms": duration_ms,
                            "triggered_by": triggered_by,
                        },
                        session_id=session_id,
                    ))
                except Exception as exc:
                    logger.debug("[circadian] event emit failed (SLEEP_COMPLETED): %s", exc)

            logger.info("[circadian] session %s is now SLEEPING (duration=%dms)", session_id, duration_ms)

            # Build report, merging SleepPipeline results where available
            report = SleepReport(
                session_id=session_id,
                duration_ms=duration_ms,
            )
            if pipeline_report is not None:
                report.claims_promoted = pipeline_report.claims_promoted
                report.summary = pipeline_report.s1.summary if pipeline_report.s1 else ""

            return report

    async def dream(self, owner_id: str = "default") -> DreamReport:
        """
        Nightly deep consolidation trigger (Electric Sheep).

        M6a responsibility: transition to DREAMING state, recover orphaned sessions,
        log the dream cycle start. Full dream logic (consolidation, decay, etc.) is M6b.

        Returns:
            DreamReport (stub — detailed fields populated by M6b).
        """
        logger.info("[circadian] dream() — owner=%s", owner_id)
        start_ts = datetime.now(timezone.utc)

        # Concurrency guard — reject duplicate concurrent dream() calls (#477).
        # In asyncio (cooperative), the locked() check + async with is race-free.
        if self._dream_lock.locked():
            logger.warning(
                "[circadian] dream() concurrent call rejected for owner=%s — already dreaming",
                owner_id,
            )
            return DreamReport(
                owner_id=owner_id,
                journal="concurrent dream rejected — dream cycle already in progress",
            )

        async with self._dream_lock:
            # Recover orphaned sessions first
            orphans_recovered = await self._recover_orphans(owner_id=owner_id)

            # Also recover sessions interrupted mid-consolidation (sleeping/dreaming with no ended_at)
            orphans_recovered += await self._recover_interrupted_sessions(owner_id=owner_id)

            # Run SleepConsolidator dream cycle if available (M6b full consolidation)
            consolidator_report = None
            if self._consolidator is not None:
                try:
                    consolidator_report = await self._consolidator.dream(owner_id)
                    logger.info(
                        "[circadian] consolidator dream complete owner=%s "
                        "consolidated=%d contradictions=%d",
                        owner_id,
                        consolidator_report.claims_consolidated,
                        consolidator_report.contradictions_resolved,
                    )
                except Exception as exc:
                    logger.error("[circadian] consolidator dream failed for owner=%s: %s", owner_id, exc)

            # WAL checkpoint after consolidation — prevent unbounded WAL growth (#490)
            try:
                result = await self.storage._db.execute("PRAGMA wal_checkpoint(TRUNCATE)")
                logger.info("[circadian] WAL checkpoint (TRUNCATE) complete owner=%s result=%s", owner_id, result)
            except Exception as exc:
                logger.warning("[circadian] WAL checkpoint failed (non-fatal) owner=%s: %s", owner_id, exc)

            end_ts = datetime.now(timezone.utc)
            duration_ms = int((end_ts - start_ts).total_seconds() * 1000)

            # Build report: merge consolidator data if available
            if consolidator_report is not None:
                report = DreamReport(
                    owner_id=owner_id,
                    duration_ms=consolidator_report.duration_ms or duration_ms,
                    orphans_recovered=orphans_recovered + consolidator_report.orphans_recovered,
                    claims_consolidated=consolidator_report.claims_consolidated,
                    contradictions_resolved=consolidator_report.contradictions_resolved,
                    claims_decayed=consolidator_report.claims_decayed,
                    claims_archived=consolidator_report.claims_archived,
                    cornerstone_proposals=consolidator_report.cornerstone_proposals,
                    entity_merge_suggestions=consolidator_report.entity_merge_suggestions,
                    brain_graphs_refreshed=consolidator_report.brain_graphs_refreshed,
                    feedback_applied=consolidator_report.feedback_applied,
                    journal=consolidator_report.journal,
                )
            else:
                journal = (
                    f"[M6a] Dream cycle triggered at {start_ts.isoformat()} UTC. "
                    f"Orphans recovered: {orphans_recovered}. "
                    "Full consolidation pending M6b."
                )
                report = DreamReport(
                    owner_id=owner_id,
                    duration_ms=duration_ms,
                    orphans_recovered=orphans_recovered,
                    journal=journal,
                )

            # Emit DREAM_COMPLETED event
            if self.event_emitter is not None:
                try:
                    from cortex.events.emitter import CortexEvent, EventType
                    await self.event_emitter.emit(CortexEvent(
                        event_type=EventType.DREAM_COMPLETED,
                        owner_id=owner_id,
                        payload={
                            "duration_ms": report.duration_ms,
                            "orphans_recovered": orphans_recovered,
                            "claims_consolidated": report.claims_consolidated,
                        },
                    ))
                except Exception as exc:
                    logger.debug("[circadian] event emit failed (DREAM_COMPLETED): %s", exc)

            logger.info("[circadian] dream cycle complete owner=%s orphans=%d",
                        owner_id, orphans_recovered)

            return report

    # ------------------------------------------------------------------
    # Manual trigger API
    # ------------------------------------------------------------------

    async def force_sleep(self, session_id: str) -> SleepReport:
        """Force sleep on a specific session (manual trigger)."""
        return await self.sleep(session_id, triggered_by="manual")

    async def force_dream(self, owner_id: str = "default") -> DreamReport:
        """Force dream cycle immediately (manual trigger)."""
        logger.info("[circadian] force_dream() triggered manually for owner=%s", owner_id)
        return await self.dream(owner_id=owner_id)

    # ------------------------------------------------------------------
    # State accessors
    # ------------------------------------------------------------------

    def get_state(self, session_id: str) -> CircadianState:
        """Return the current in-memory state for a session."""
        return self._states.get(session_id, CircadianState.AWAKE)

    async def touch_activity(self, session_id: str) -> None:
        """Update the last-activity timestamp for a session (reset idle timer)."""
        self._last_activity[session_id] = datetime.now(timezone.utc).timestamp()

    # ------------------------------------------------------------------
    # Background scheduler management
    # ------------------------------------------------------------------

    def _ensure_idle_watcher(self) -> None:
        """Start the idle-watcher background task if not already running."""
        if self._idle_watcher_task is None or self._idle_watcher_task.done():
            self._idle_watcher_task = asyncio.create_task(
                self._idle_watcher_loop(),
                name="circadian-idle-watcher",
            )
            logger.debug("[circadian] idle watcher started (timeout=%ds)", self._idle_timeout_sec)

    def _ensure_dream_scheduler(self) -> None:
        """Start the dream-cycle scheduler background task if not already running."""
        if self._dream_scheduler_task is None or self._dream_scheduler_task.done():
            self._dream_scheduler_task = asyncio.create_task(
                self._dream_scheduler_loop(),
                name="circadian-dream-scheduler",
            )
            logger.debug("[circadian] dream scheduler started (hour_utc=%d)", self._dream_hour_utc)

    async def _idle_watcher_loop(self) -> None:
        """
        Background task: poll every 60 seconds and sleep any session that has
        been idle for longer than idle_timeout_sec.
        """
        poll_interval = min(60, self._idle_timeout_sec // 2)
        while True:
            try:
                await asyncio.sleep(poll_interval)
                now = datetime.now(timezone.utc).timestamp()
                # Snapshot to avoid modifying dict during iteration
                active = {
                    sid: ts for sid, ts in self._last_activity.items()
                    if self._states.get(sid) in (CircadianState.AWAKE, CircadianState.DROWSY)
                }
                for session_id, last_ts in active.items():
                    idle_sec = now - last_ts
                    if idle_sec >= self._idle_timeout_sec:
                        logger.info(
                            "[circadian] session %s idle for %.0fs → triggering sleep()",
                            session_id, idle_sec,
                        )
                        try:
                            await self.sleep(session_id, triggered_by="idle_timeout")
                        except Exception as exc:
                            logger.error("[circadian] auto-sleep failed for %s: %s", session_id, exc)
            except asyncio.CancelledError:
                logger.debug("[circadian] idle watcher cancelled")
                break
            except Exception as exc:
                logger.error("[circadian] idle watcher error: %s", exc)
                await asyncio.sleep(5)

    async def _dream_scheduler_loop(self) -> None:
        """
        Background task: fire dream() once per day at the configured UTC hour.
        Waits until next occurrence of dream_cycle_hour_utc, then fires and loops.
        """
        while True:
            try:
                seconds_until_dream = self._seconds_until_dream_hour()
                logger.debug("[circadian] dream scheduler: next dream in %.0fs", seconds_until_dream)
                await asyncio.sleep(seconds_until_dream)
                try:
                    _timeout = getattr(self.config, "dream_cycle_timeout_sec", 600) if self.config else 600
                    await asyncio.wait_for(
                        self.dream(owner_id=self._owner_id),
                        timeout=_timeout,
                    )
                except asyncio.TimeoutError:
                    logger.error(
                        "[circadian] dream cycle timed out after %ds — marking incomplete",
                        _timeout,
                    )
                except Exception as exc:
                    logger.error("[circadian] dream cycle failed: %s", exc)
            except asyncio.CancelledError:
                logger.debug("[circadian] dream scheduler cancelled")
                break
            except Exception as exc:
                logger.error("[circadian] dream scheduler error: %s", exc)
                await asyncio.sleep(60)

    def _seconds_until_dream_hour(self) -> float:
        """
        Compute how many seconds until the next occurrence of dream_cycle_hour_utc.
        Always returns a positive number (minimum 1 second).
        """
        now = datetime.now(timezone.utc)
        target_today = now.replace(
            hour=self._dream_hour_utc,
            minute=0,
            second=0,
            microsecond=0,
        )
        if now >= target_today:
            # Already past today's dream hour — schedule for tomorrow
            from datetime import timedelta
            target = target_today + timedelta(days=1)
        else:
            target = target_today
        delta = (target - now).total_seconds()
        return max(1.0, delta)

    # ------------------------------------------------------------------
    # Context manager support
    # ------------------------------------------------------------------

    async def __aenter__(self) -> "CircadianEngine":
        """Async context manager entry — returns self."""
        return self

    async def __aexit__(
        self,
        exc_type: Optional[type],
        exc_val: Optional[BaseException],
        exc_tb: Any,
    ) -> None:
        """Async context manager exit — triggers clean shutdown."""
        await self.shutdown(triggered_by="context_manager_exit")

    # ------------------------------------------------------------------
    # Signal handler registration
    # ------------------------------------------------------------------

    def register_signal_handlers(self) -> None:
        """
        Register SIGTERM / SIGINT handlers that call shutdown() gracefully.

        Safe to call multiple times — only registers once.
        Only operates when there is a running event loop (skipped in test runners
        that patch signals away).
        """
        if self._signal_handlers_registered:
            return

        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            logger.debug("[circadian] no running loop — signal handlers not registered")
            return

        def _handle_signal(sig: signal.Signals) -> None:
            logger.info("[circadian] received signal %s — initiating graceful shutdown", sig.name)
            # Schedule shutdown as a Task so it runs in the event loop
            loop.create_task(
                self.shutdown(triggered_by=f"signal:{sig.name}"),
                name=f"circadian-shutdown-{sig.name}",
            )

        for sig in (signal.SIGTERM, signal.SIGINT):
            try:
                loop.add_signal_handler(sig, _handle_signal, sig)
                logger.debug("[circadian] registered handler for %s", sig.name)
            except (OSError, NotImplementedError):
                # Windows or environments where signal handlers aren't supported
                logger.debug("[circadian] cannot register signal handler for %s", sig.name)

        self._signal_handlers_registered = True

    # ------------------------------------------------------------------
    # State persistence — crash recovery
    # ------------------------------------------------------------------

    async def persist_state(self) -> bool:
        """
        Write current engine state to DB (cortex_config table, key=_CIRCADIAN_STATE_KEY).

        Stored as JSON: {session_id: state_value, ...} plus a timestamp.

        Returns True if persistence succeeded, False on error.
        """
        try:
            payload: dict[str, Any] = {
                "states": {sid: st.value for sid, st in self._states.items()},
                "persisted_at": datetime.now(timezone.utc).isoformat(),
            }
            await self.storage.set_config(_CIRCADIAN_STATE_KEY, json.dumps(payload))
            logger.debug("[circadian] state persisted to DB: %d sessions", len(self._states))
            return True
        except Exception as exc:
            logger.error("[circadian] failed to persist state to DB: %s", exc)
            return False

    async def resume_saved_state(self) -> dict[str, CircadianState]:
        """
        On startup: check for persisted state in DB and restore it.

        Useful for crash recovery — returns the restored states dict
        (may be empty if no saved state exists or state is stale).

        The loaded state is merged into self._states only for sessions that
        are still present in storage and still active (not yet ended).
        """
        try:
            raw = await self.storage.get_config(_CIRCADIAN_STATE_KEY)
            if not raw:
                logger.debug("[circadian] no persisted state found in DB")
                return {}

            payload = json.loads(raw)
            states_raw: dict = payload.get("states", {})
            persisted_at: str = payload.get("persisted_at", "")
            logger.info(
                "[circadian] found persisted state from %s — %d sessions",
                persisted_at,
                len(states_raw),
            )

            restored: dict[str, CircadianState] = {}
            for sid, state_val in states_raw.items():
                try:
                    state = CircadianState(state_val)
                    # Only restore awake/drowsy sessions — sleeping/dreaming are
                    # handled by _recover_interrupted_sessions() at wake time
                    if state in (CircadianState.AWAKE, CircadianState.DROWSY):
                        async with self._states_lock:
                            self._states[sid] = state
                        restored[sid] = state
                        logger.debug("[circadian] restored session %s → %s", sid, state)
                except ValueError:
                    logger.warning("[circadian] unknown state value in persisted data: %r", state_val)

            return restored
        except json.JSONDecodeError as exc:
            logger.error("[circadian] persisted state JSON is invalid: %s", exc)
            return {}
        except Exception as exc:
            logger.error("[circadian] failed to load persisted state: %s", exc)
            return {}

    # ------------------------------------------------------------------
    # Clean shutdown
    # ------------------------------------------------------------------

    async def shutdown(self, triggered_by: str = "manual") -> ShutdownReport:
        """
        Perform a clean shutdown of the CircadianEngine.

        Steps:
          1. Guard against double-shutdown.
          2. Cancel background tasks (idle watcher, dream scheduler).
          3. For each active (awake/drowsy) session: trigger emergency sleep
             to flush state and avoid data loss.
          4. Mark all remaining sessions as SHUTDOWN in memory.
          5. Persist current state to DB for crash recovery.
          6. Log the shutdown event.

        Args:
            triggered_by: Reason string ('manual', 'context_manager_exit',
                          'signal:SIGTERM', 'signal:SIGINT', etc.)

        Returns:
            ShutdownReport summarising what was done.
        """
        async with self._shutdown_lock:
            if self._is_shutdown:
                logger.debug("[circadian] shutdown() called but already shut down — ignoring")
                return ShutdownReport(triggered_by=triggered_by)

            logger.info("[circadian] shutdown() initiated — trigger=%s", triggered_by)
            report = ShutdownReport(triggered_by=triggered_by)

            # 1. Cancel background tasks
            tasks_cancelled = 0
            for task in (self._idle_watcher_task, self._dream_scheduler_task):
                if task is not None and not task.done():
                    task.cancel()
                    try:
                        await task
                    except asyncio.CancelledError:
                        pass
                    tasks_cancelled += 1
            report.tasks_cancelled = tasks_cancelled
            logger.debug("[circadian] cancelled %d background tasks", tasks_cancelled)

            # 2. Emergency sleep for all active sessions
            active_sessions = [
                sid for sid, state in self._states.items()
                if state in (CircadianState.AWAKE, CircadianState.DROWSY)
            ]
            for session_id in active_sessions:
                logger.warning(
                    "[circadian] emergency sleep for active session %s (shutdown in progress)",
                    session_id,
                )
                try:
                    await self.sleep(session_id, triggered_by=f"emergency_shutdown:{triggered_by}")
                    report.sessions_slept += 1
                except Exception as exc:
                    logger.error(
                        "[circadian] emergency sleep failed for session %s: %s",
                        session_id, exc,
                    )

            # 3. Mark all sessions as SHUTDOWN
            async with self._states_lock:
                for session_id in list(self._states.keys()):
                    self._states[session_id] = CircadianState.SHUTDOWN

            # 4. Persist state to DB
            report.state_persisted = await self.persist_state()

            # 5. Mark engine as shut down
            self._is_shutdown = True

            logger.info(
                "[circadian] shutdown complete — sessions_slept=%d tasks_cancelled=%d "
                "state_persisted=%s trigger=%s",
                report.sessions_slept,
                report.tasks_cancelled,
                report.state_persisted,
                triggered_by,
            )

            return report

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    async def _get_session(self, session_id: str):
        """Fetch the session record from storage (returns None if not found)."""
        try:
            return await self.storage.get_active_session_by_id(session_id)
        except (AttributeError, NotImplementedError):
            # Fallback: try get_active_session if adapter doesn't have by-id lookup
            try:
                session = await self.storage.get_active_session()
                if session and session.session_id == session_id:
                    return session
                return None
            except Exception:
                return None

    async def _recover_orphans(self, owner_id: Optional[str] = None) -> int:
        """
        Find orphaned sessions (awake but idle too long) and sleep them.

        Returns:
            Count of sessions recovered.
        """
        try:
            orphans = await self.storage.get_orphaned_sessions(self._idle_timeout_sec)
        except Exception as exc:
            logger.warning("[circadian] orphan check failed: %s", exc)
            return 0

        if owner_id is not None:
            orphans = [s for s in orphans if s.owner_id == owner_id]

        recovered = 0
        for session in orphans:
            logger.info("[circadian] recovering orphaned session: %s", session.session_id)
            try:
                await self.sleep(session.session_id, triggered_by="orphan_recovery")
                recovered += 1
            except Exception as exc:
                logger.error("[circadian] failed to recover orphan %s: %s", session.session_id, exc)

        return recovered

    async def _recover_interrupted_sessions(self, owner_id: Optional[str] = None) -> int:
        """
        Find sessions stuck in 'sleeping' or 'dreaming' with no ``ended_at``
        and close them by marking status='completed' and recording ended_at.

        These sessions were interrupted mid-consolidation (e.g. process crashed
        during sleep step 3/5 or dream step 5/9).  Recovery is lightweight —
        we do NOT re-run consolidation, we simply seal the session so it no
        longer appears as in-progress.

        Returns:
            Count of sessions recovered.
        """
        try:
            interrupted = await self.storage.get_interrupted_sessions(owner_id=owner_id)
        except (AttributeError, NotImplementedError) as exc:
            logger.warning("[circadian] interrupted session check not supported by adapter: %s", exc)
            return 0
        except Exception as exc:
            logger.warning("[circadian] interrupted session check failed: %s", exc)
            return 0

        recovered = 0
        now = datetime.now(timezone.utc).isoformat()
        for session in interrupted:
            logger.warning(
                "[circadian] recovering interrupted %s session %s "
                "(no ended_at — likely crashed mid-consolidation)",
                session.status,
                session.session_id,
            )
            try:
                await self.storage.update_session(
                    session.session_id,
                    status="completed",
                    ended_at=now,
                )
                recovered += 1
                logger.info(
                    "[circadian] interrupted session %s closed "
                    "(status→completed, ended_at=%s)",
                    session.session_id,
                    now,
                )
            except Exception as exc:
                logger.error(
                    "[circadian] failed to recover interrupted session %s: %s",
                    session.session_id,
                    exc,
                )

        return recovered

    async def _retrieve_pending_context(self, owner_id: str) -> Optional[str]:
        """
        Retrieve summary from the most recent completed session for this owner.
        Returns the summary text, or None if not available.
        """
        try:
            # We look for the most-recently-ended session for context carry-forward
            # Using get_active_session won't work here (that's the CURRENT session)
            # Try a generic query if the adapter supports it
            session = await self.storage.get_active_session(owner_id=owner_id)
            if session and session.summary:
                return session.summary
        except Exception as e:
            logger.warning("cortex_circadian_engine: unhandled exception: %s", e)
        return None

    async def _generate_nap_summary(self, session_id: str, claims_count: int) -> str:
        """Generate a brief summary for the nap checkpoint using the LLM."""
        try:
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(self.llm, "nap_summary")
            response = await self.llm.complete(
                system="You are a concise session summarizer for an AI memory system.",
                user=(
                    f"Write a 1-2 sentence summary of a session checkpoint. "
                    f"Session ID: {session_id}. "
                    f"Working memory flushed: {claims_count} claims. "
                    "This is a mid-session nap consolidation. Keep it concise."
                ),
                max_tokens=80,
                **({"model": _routed_model} if _routed_model else {}),
            )
            return response.text.strip()
        except Exception as exc:
            logger.warning("[circadian] nap summary generation failed: %s", exc)
            return f"[Nap checkpoint — {claims_count} claims flushed]"
