"""
Circadian engine — state persistence, crash recovery, and context manager (mixin).
"""

from __future__ import annotations

import asyncio
import json
import logging
import signal
from datetime import datetime, timezone
from typing import Any, Optional

from cortex.circadian.engine.models import (
    CircadianState,
    ShutdownReport,
    _CIRCADIAN_STATE_KEY,
)

logger = logging.getLogger(__name__)


class PersistenceMixin:
    """Mixin providing state save/load, signal handling, and context manager for CircadianEngine."""

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

    async def __aenter__(self) -> "PersistenceMixin":
        """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).
        """
        # Import here to avoid circular — CircadianEngine is the actual class
        from cortex.circadian.engine.lifecycle import CircadianEngine

        if CircadianEngine._signal_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)

        CircadianEngine._signal_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 — session retrieval and recovery
    # ------------------------------------------------------------------

    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
