"""cortex/circadian/consolidation/sleep_pipeline.py — SleepConsolidator class."""
from __future__ import annotations

import asyncio, json, logging, time
from datetime import datetime, timezone
from typing import TYPE_CHECKING, List, Optional

from cortex.circadian.consolidation._helpers import MAX_CLAIMS_PER_BATCH
from cortex.circadian.consolidation.dream_pipeline import DreamPipelineMixin
from cortex.circadian.sleep import (
    SleepReport, _SESSION_SUMMARY_SYSTEM, _SESSION_SUMMARY_USER,
)
from cortex.types import SEMANTIC_MEMORY_TYPES

if TYPE_CHECKING:
    from cortex.storage.adapter import StorageAdapter
    from cortex.config import CortexConfig
    from cortex.types import LLMProvider, SemanticClaim
    from cortex.identity.cornerstone import CornerstoneGuardian
    from cortex.identity.drift_calculator import DriftCalculator, DriftReport
    from cortex.deprecation.pipeline import DeprecationPipeline
    from cortex.core.entity_resolution import EntityResolver
    from cortex.core.feedback import FeedbackSystem

logger = logging.getLogger(__name__)

class SleepConsolidator(DreamPipelineMixin):
    """
    M6b: Full sleep and dream consolidation logic.

    Designed to be called from CircadianEngine via on_sleep() hook,
    or directly for testing / manual triggers.

    Args:
        storage:               StorageAdapter for persistence.
        config:                CortexConfig with lifecycle parameters.
        llm:                   LLMProvider for summaries and analysis.
        cornerstone_guardian:  CornerstoneGuardian for cornerstone operations.
        drift_calculator:      DriftCalculator for drift detection/restoration.
        deprecation_pipeline:  DeprecationPipeline for memory lifecycle.
        entity_resolver:       EntityResolver for suggest_merges() (optional).
        feedback_system:       FeedbackSystem for apply_pending() (optional).
        peer_manager:          PeerManager for peer refresh during dreams (optional).
    """

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

        # Config helpers
        self._owner_id: str = getattr(config, "owner_id", "default")
        self._idle_timeout_sec: int = getattr(config, "dream_cycle_idle_timeout_sec", 3600)
        self._decay_window_hours: int = getattr(config, "dream_decay_window_hours", 24)
        self._promoted_confidence: float = getattr(config, "promoted_confidence_default", 0.8)
        # Number of claims to process before yielding to the event loop (#488)
        self._consolidation_batch_size: int = getattr(
            config, "sleep_consolidation_batch_size", 50
        )
        self._contradiction_batch_size: int = getattr(
            config, "contradiction_batch_size", MAX_CLAIMS_PER_BATCH
        )

    # ------------------------------------------------------------------
    # Public API — sleep()
    # ------------------------------------------------------------------

    async def sleep(self, session_id: str) -> SleepReport:
        """
        End-of-session consolidation (target: <3 minutes).

        Pipeline:
          1. SESSION MEMORY DUMP
          2. DEPRECATION REVIEW
          3. CORNERSTONE DRIFT CHECK
          4. QUICK CONSOLIDATION
          5. SLEEP CONFIRMATION
        """
        logger.info("[M6b] sleep() — session=%s", session_id)
        t_start = time.monotonic()

        report = SleepReport(session_id=session_id)

        # Step 1: SESSION MEMORY DUMP
        logger.debug("[M6b] step 1/5 — session memory dump for %s", session_id)
        session_claims = await self._step_session_dump(session_id, report)

        # Step 2: DEPRECATION REVIEW
        logger.debug("[M6b] step 2/5 — deprecation review for %s", session_id)
        await self._step_deprecation_review(session_id, report)

        # Step 3: CORNERSTONE DRIFT CHECK
        logger.debug("[M6b] step 3/5 — cornerstone drift check for %s", session_id)
        await self._step_drift_check(report)

        # Step 4: QUICK CONSOLIDATION
        logger.debug("[M6b] step 4/5 — quick consolidation for %s", session_id)
        await self._step_quick_consolidation(session_id, session_claims, report)

        # Step 5: SLEEP CONFIRMATION
        logger.debug("[M6b] step 5/5 — sleep confirmation for %s", session_id)
        await self._step_sleep_confirmation(session_id, report)

        report.duration_ms = int((time.monotonic() - t_start) * 1000)
        report.completed_at = datetime.now(timezone.utc).isoformat()

        # Emit events for promoted and deprecated memories
        if self.event_emitter is not None:
            try:
                from cortex.events.emitter import CortexEvent, EventType
                if report.claims_promoted > 0:
                    await self.event_emitter.emit(CortexEvent(
                        event_type=EventType.MEMORY_PROMOTED,
                        owner_id=self._owner_id,
                        payload={
                            "session_id": session_id,
                            "claims_promoted": report.claims_promoted,
                        },
                        session_id=session_id,
                    ))
                if report.claims_pruned > 0:
                    await self.event_emitter.emit(CortexEvent(
                        event_type=EventType.MEMORY_DEPRECATED,
                        owner_id=self._owner_id,
                        payload={
                            "session_id": session_id,
                            "claims_pruned": report.claims_pruned,
                        },
                        session_id=session_id,
                    ))
            except Exception as exc:
                logger.debug("[M6b] event emit failed (sleep events): %s", exc)

        logger.info(
            "[M6b] sleep() complete — session=%s promoted=%d pruned=%d drift_checks=%d duration_ms=%d",
            session_id, report.claims_promoted, report.claims_pruned,
            len(report.drift_checks), report.duration_ms,
        )
        return report

    # ------------------------------------------------------------------
    # Sleep pipeline steps (private)
    # ------------------------------------------------------------------

    async def _step_session_dump(
        self,
        session_id: str,
        report: SleepReport,
    ) -> List["SemanticClaim"]:
        """
        Step 1: SESSION MEMORY DUMP
        - Retrieve all claims sourced from this session
        - Generate a session summary via LLM
        - Tag all session claims as 'cooling' if still 'active'
        """
        claims = await self.storage.get_session_claims(session_id)
        report.claims_created = len(claims)

        if not claims:
            logger.debug("[M6b] no session claims found for %s", session_id)
            report.summary = "(no session claims)"
            return claims

        # Generate session summary via LLM
        try:
            claims_json = json.dumps(
                [
                    {
                        "id": c.id,
                        "subject": c.subject,
                        "predicate": c.predicate,
                        "object": c.object_value,
                        "confidence": c.confidence,
                    }
                    for c in claims[:50]  # cap at 50 to avoid token explosion
                ],
                ensure_ascii=False,
                indent=2,
            )
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(self.llm, "session_summary")
            _resp = await self.llm.complete(
                system=_SESSION_SUMMARY_SYSTEM,
                user=_SESSION_SUMMARY_USER.format(claims_json=claims_json),
                response_format="text",
                **({"model": _routed_model} if _routed_model else {}),
            )
            summary = _resp.text if hasattr(_resp, "text") else str(_resp)
            report.summary = summary.strip()
        except Exception as exc:
            logger.warning("[M6b] LLM session summary failed: %s", exc)
            report.summary = f"(summary unavailable: {exc})"

        logger.debug(
            "[M6b] session dump complete: %d claims, summary=%r",
            len(claims), report.summary[:80],
        )
        return claims

    async def _step_deprecation_review(
        self,
        session_id: str,
        report: SleepReport,
    ) -> None:
        """
        Step 2: DEPRECATION REVIEW
        - Run DeprecationPipeline.review_session() on all session claims
        - Record promoted / pruned counts in report
        """
        try:
            dep_report = await self.deprecation_pipeline.review_session(session_id)
            report.claims_promoted += dep_report.promoted
            report.claims_pruned += dep_report.deprecated
            logger.debug(
                "[M6b] deprecation review: promoted=%d kept=%d deprecated=%d noise=%d",
                dep_report.promoted, dep_report.kept_session,
                dep_report.deprecated, dep_report.auto_noise,
            )
        except Exception as exc:
            logger.error("[M6b] deprecation review failed for session %s: %s", session_id, exc)

    async def _step_drift_check(self, report: SleepReport) -> None:
        """
        Step 3: CORNERSTONE DRIFT CHECK
        - Run DriftCalculator.check_all() on all cornerstones
        - Auto-restore cornerstones with action='restore' or 'alert'
        """
        try:
            drift_reports = await self.drift_calculator.check_all(owner_id=self._owner_id)
            report.drift_checks = list(drift_reports)

            for dr in drift_reports:
                if dr.action in ("restore", "alert"):
                    try:
                        await self.drift_calculator.restore(dr.cornerstone_id)
                        logger.info(
                            "[M6b] restored cornerstone %s (drift=%.3f action=%s)",
                            dr.cornerstone_id, dr.drift_score, dr.action,
                        )
                    except Exception as exc:
                        logger.error(
                            "[M6b] failed to restore cornerstone %s: %s",
                            dr.cornerstone_id, exc,
                        )
                elif dr.action == "warn":
                    logger.warning(
                        "[M6b] cornerstone drift warning: %s label=%r drift=%.3f",
                        dr.cornerstone_id, dr.label, dr.drift_score,
                    )

            logger.debug("[M6b] drift check complete: %d cornerstones checked", len(drift_reports))
        except Exception as exc:
            logger.error("[M6b] drift check failed: %s", exc)

    async def _step_quick_consolidation(
        self,
        session_id: str,
        session_claims: List["SemanticClaim"],
        report: SleepReport,
    ) -> None:
        """
        Step 4: QUICK CONSOLIDATION
        - Fold promoted claims into the semantic layer (memory_type → 'semantic')
        - Update entity relationships for promoted claims
        - Create memory edges between related promoted claims
        """
        if not session_claims:
            return

        promoted_claims = [
            c for c in session_claims
            if c.memory_type in SEMANTIC_MEMORY_TYPES and c.status == "active"
        ]

        for idx, claim in enumerate(promoted_claims):
            # Yield to the event loop every batch_size claims to avoid blocking (#488)
            if idx > 0 and idx % self._consolidation_batch_size == 0:
                logger.debug(
                    "[M6b] quick consolidation: yielding event loop at claim %d/%d",
                    idx, len(promoted_claims),
                )
                await asyncio.sleep(0)

            try:
                # Ensure confidence is at the promoted level minimum
                new_confidence = max(claim.confidence, self._promoted_confidence)
                if new_confidence != claim.confidence:
                    await self.storage.update_claim(
                        claim.id,
                        confidence=new_confidence,
                        updated_at=datetime.now(timezone.utc).isoformat(),
                    )

                # Changelog: record session memory promotion
                try:
                    history = await self.storage.get_changelog(claim.id)
                    next_ver = len(history) + 1
                    await self.storage.log_changelog(
                        claim_id=claim.id,
                        version=next_ver,
                        content_snapshot=claim.object_value,
                        action="promoted",
                        reason="sleep cycle quick consolidation — session → semantic",
                        changed_by="sleep_cycle",
                        diff_summary=None,
                    )
                except Exception as exc:
                    logger.warning("[M6b] changelog log failed for claim %s: %s", claim.id, exc)

            except Exception as exc:
                logger.warning("[M6b] consolidation update failed for claim %s: %s", claim.id, exc)

        logger.debug(
            "[M6b] quick consolidation: %d/%d claims promoted into semantic layer",
            len(promoted_claims), len(session_claims),
        )

    async def _step_sleep_confirmation(
        self,
        session_id: str,
        report: SleepReport,
    ) -> None:
        """
        Step 5: SLEEP CONFIRMATION
        - Update session_record with sleep stats and summary
        - Set status='sleeping'
        """
        try:
            sleep_report_json = json.dumps(
                {
                    "claims_created": report.claims_created,
                    "claims_promoted": report.claims_promoted,
                    "claims_pruned": report.claims_pruned,
                    "drift_checks": len(report.drift_checks),
                    "summary": report.summary,
                    "completed_at": report.completed_at,
                },
                ensure_ascii=False,
            )
            await self.storage.update_session(
                session_id,
                status="sleeping",
                memories_created=report.claims_created,
                memories_promoted=report.claims_promoted,
                memories_pruned=report.claims_pruned,
                summary=report.summary,
                sleep_report=sleep_report_json,
                ended_at=datetime.now(timezone.utc).isoformat(),
            )
            logger.debug("[M6b] session %s status → sleeping", session_id)
        except Exception as exc:
            logger.error("[M6b] sleep confirmation failed for session %s: %s", session_id, exc)

    # ------------------------------------------------------------------
    # Dream step helper — per-step error isolation
    # ------------------------------------------------------------------
