"""
Circadian engine — fatigue tracking and nap logic (mixin).
"""

from __future__ import annotations

import asyncio
import logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING

from cortex.circadian.engine.models import (
    CircadianState,
    NapReport,
    FATIGUE_AFTER_NAP,
    FATIGUE_COMPACTION_BUMP,
    FATIGUE_ERROR_BUMP,
    FATIGUE_MAX,
    FATIGUE_TIME_INCREMENT,
)

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


class FatigueMixin:
    """Mixin providing fatigue tracking and nap methods for CircadianEngine."""

    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 _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]"
