"""cortex/circadian/consolidation/dream_steps.py — Dream step implementations."""
from __future__ import annotations

import json, logging
from datetime import datetime, timezone, timedelta
from typing import TYPE_CHECKING, List, Optional

from cortex.circadian.consolidation._helpers import _parse_dt, _parse_json_list
from cortex.circadian.consolidation.dream_reconsolidation import DreamReconsolidationMixin
from cortex.circadian.sleep import (
    _CORNERSTONE_PROPOSAL_SYSTEM, _CORNERSTONE_PROPOSAL_USER,
)
from cortex.types import SEMANTIC_MEMORY_TYPES

if TYPE_CHECKING:
    from cortex.types import SemanticClaim

logger = logging.getLogger(__name__)


class DreamStepsMixin(DreamReconsolidationMixin):
    """Dream step implementations."""

    # B-OS3: Max times a claimless-awake session can be skipped before being
    # marked as 'orphan' to prevent infinite re-evaluation every dream cycle.
    _CLAIMLESS_SKIP_LIMIT = 3

    async def _dream_recover_orphans(self, owner_id: str) -> int:
        """
        Step 1: ORPHAN RECOVERY
        Find sessions that are still 'awake' past idle_timeout and sleep them.

        B-OS3 fix: Claimless sessions that are skipped repeatedly (>= _CLAIMLESS_SKIP_LIMIT
        consecutive cycles) are marked as 'orphan' to stop infinite re-evaluation.
        The skip count is tracked in the session's sleep_report JSON field.
        """
        orphans_recovered = 0
        try:
            orphaned = await self.storage.get_orphaned_sessions(
                idle_seconds=self._idle_timeout_sec
            )
            for orphan in orphaned:
                if orphan.owner_id != owner_id:
                    continue
                # R-321: Skip claimless sessions — avoid wasting I/O on sessions
                # with no associated claims (e.g. batch-ingest chunk sessions whose
                # source_session FK uses content hashes that don't match session_id).
                try:
                    claim_count = await self.storage.count_session_claims(orphan.session_id)
                    if claim_count == 0:
                        # B-OS3: Track skip count and mark as orphan after threshold
                        skip_count = await self._increment_claimless_skip_count(orphan.session_id)
                        if skip_count >= self._CLAIMLESS_SKIP_LIMIT:
                            logger.info(
                                "[M6b] marking claimless session %s as orphan after %d skips (owner=%s)",
                                orphan.session_id, skip_count, orphan.owner_id,
                            )
                            await self._mark_session_orphan(orphan.session_id)
                        else:
                            logger.debug(
                                "[M6b] skipping claimless orphan session %s (owner=%s, skip_count=%d/%d)",
                                orphan.session_id, orphan.owner_id, skip_count, self._CLAIMLESS_SKIP_LIMIT,
                            )
                        continue
                except Exception as exc:
                    logger.debug(
                        "[M6b] claim count check failed for %s, proceeding: %s",
                        orphan.session_id, exc,
                    )
                try:
                    logger.info(
                        "[M6b] recovering orphaned session %s (owner=%s)",
                        orphan.session_id, orphan.owner_id,
                    )
                    await self.sleep(orphan.session_id)
                    orphans_recovered += 1
                except Exception as exc:
                    logger.error(
                        "[M6b] failed to recover orphan session %s: %s",
                        orphan.session_id, exc,
                    )
        except Exception as exc:
            logger.error("[M6b] orphan recovery failed: %s", exc)
        return orphans_recovered

    async def _increment_claimless_skip_count(self, session_id: str) -> int:
        """Increment and return the claimless skip count for a session.

        B-OS3: Tracks how many dream cycles have skipped this session due to
        having zero claims. Stored in sleep_report JSON as {"claimless_skip_count": N}.
        """
        conn = await self.storage._get_conn()
        async with conn.execute(
            "SELECT sleep_report FROM session_record WHERE session_id = ?",
            (session_id,),
        ) as cur:
            row = await cur.fetchone()

        current_count = 0
        if row and row[0]:
            try:
                report = json.loads(row[0])
                current_count = report.get("claimless_skip_count", 0)
            except (json.JSONDecodeError, TypeError):
                pass

        new_count = current_count + 1
        # Merge into existing sleep_report JSON (preserve other fields)
        report_data = {}
        if row and row[0]:
            try:
                report_data = json.loads(row[0])
            except (json.JSONDecodeError, TypeError):
                pass
        report_data["claimless_skip_count"] = new_count

        await conn.execute(
            "UPDATE session_record SET sleep_report = ? WHERE session_id = ?",
            (json.dumps(report_data), session_id),
        )
        await conn.commit()
        return new_count

    async def _mark_session_orphan(self, session_id: str) -> None:
        """Mark a session as 'orphan', preserving previous status in sleep_report.

        B-OS3: Called when a claimless session exceeds the skip threshold.
        Stores previous_status in sleep_report JSON for potential rollback.
        """
        conn = await self.storage._get_conn()

        # Get current status and sleep_report
        async with conn.execute(
            "SELECT status, sleep_report FROM session_record WHERE session_id = ?",
            (session_id,),
        ) as cur:
            row = await cur.fetchone()

        if not row:
            return

        report_data = {}
        if row[1]:
            try:
                report_data = json.loads(row[1])
            except (json.JSONDecodeError, TypeError):
                pass

        report_data["previous_status"] = row[0]

        await conn.execute(
            "UPDATE session_record SET status = 'orphan', sleep_report = ? WHERE session_id = ?",
            (json.dumps(report_data), session_id),
        )
        await conn.commit()

    async def _dream_deep_consolidation(self, owner_id: str) -> int:
        """
        Step 2: DEEP CONSOLIDATION
        Review all session memories from past 24h, look for patterns/contradictions.
        Returns count of claims reviewed.
        """
        claims_reviewed = 0
        try:
            # Get all active session claims across recent sessions
            session_claims = await self.storage.get_claims_by_status(
                "active", owner_id=owner_id
            )
            # Filter to session-type memories created in past 24h
            cutoff = datetime.now(timezone.utc) - timedelta(hours=self._decay_window_hours)
            recent_session_claims = [
                c for c in session_claims
                if c.memory_type == "session"
                and c.created_at is not None
                and _parse_dt(c.created_at) >= cutoff
            ]
            claims_reviewed = len(recent_session_claims)
            logger.debug(
                "[M6b] deep consolidation: %d recent session claims in past %dh",
                claims_reviewed, self._decay_window_hours,
            )
        except Exception as exc:
            logger.error("[M6b] deep consolidation scan failed: %s", exc)
        return claims_reviewed

    async def _dream_cornerstone_evolution(self, owner_id: str) -> int:
        """
        Step 3: CORNERSTONE EVOLUTION
        Propose new cornerstones based on recurring high-confidence patterns.
        Returns count of proposals queued.
        """
        proposals_created = 0
        try:
            # Gather high-confidence semantic claims from the past 24h
            all_claims = await self.storage.get_claims_by_status("active", owner_id=owner_id)
            cutoff = datetime.now(timezone.utc) - timedelta(hours=self._decay_window_hours)
            high_confidence_claims = [
                c for c in all_claims
                if c.confidence >= 0.75
                and c.memory_type in SEMANTIC_MEMORY_TYPES
                and c.created_at is not None
                and _parse_dt(c.created_at) >= cutoff
            ]

            if not high_confidence_claims:
                logger.debug("[M6b] no high-confidence claims for cornerstone proposals")
                return 0

            # Gather existing cornerstones to avoid re-proposing
            existing = await self.storage.get_cornerstones(owner_id=owner_id)

            claims_json = json.dumps(
                [
                    {
                        "subject": c.subject,
                        "predicate": c.predicate,
                        "object": c.object_value,
                        "confidence": c.confidence,
                    }
                    for c in high_confidence_claims[:30]
                ],
                ensure_ascii=False,
                indent=2,
            )
            cornerstones_json = json.dumps(
                [{"label": cs.label, "content": cs.content[:200]} for cs in existing],
                ensure_ascii=False,
                indent=2,
            )

            from cortex.llm.router import resolve_routed_model
            _cs_model = resolve_routed_model(self.llm, "cornerstone_proposal")
            _resp = await self.llm.complete(
                system=_CORNERSTONE_PROPOSAL_SYSTEM,
                user=_CORNERSTONE_PROPOSAL_USER.format(
                    claims_json=claims_json,
                    cornerstones_json=cornerstones_json,
                ),
                response_format="json",
                **({"model": _cs_model} if _cs_model else {}),
            )
            response = _resp.text if hasattr(_resp, "text") else str(_resp)

            proposals = _parse_json_list(response)
            for prop in proposals[:3]:  # cap at 3
                label = prop.get("label", "").strip()
                content = prop.get("content", "").strip()
                reason = prop.get("reason", "")
                if label and content:
                    try:
                        # Propose evolution (queued for operator review)
                        await self.cornerstone_guardian.propose_evolution(
                            label=label,
                            proposed_content=content,
                            reason=reason,
                            owner_id=owner_id,
                        )
                        proposals_created += 1
                        logger.info(
                            "[M6b] cornerstone proposal queued: %r", label
                        )
                    except Exception as exc:
                        logger.warning(
                            "[M6b] cornerstone proposal failed for %r: %s", label, exc
                        )

        except Exception as exc:
            logger.error("[M6b] cornerstone evolution failed: %s", exc)
        return proposals_created


    async def _dream_memory_decay(self, owner_id: str) -> tuple[int, int]:
        """
        Step 5: MEMORY DECAY
        Run DeprecationPipeline.run_decay() for time-based confidence reduction.
        Returns (claims_decayed, claims_archived).
        """
        claims_decayed = 0
        claims_archived = 0
        try:
            decay_report = await self.deprecation_pipeline.run_decay(owner_id=owner_id)
            claims_decayed = (
                decay_report.transitioned_to_cooling
                + decay_report.transitioned_to_deprecated
                + decay_report.confidence_reduced
            )
            claims_archived = decay_report.transitioned_to_archived
            logger.debug(
                "[M6b] decay: cooling=%d deprecated=%d archived=%d confidence_reduced=%d",
                decay_report.transitioned_to_cooling,
                decay_report.transitioned_to_deprecated,
                decay_report.transitioned_to_archived,
                decay_report.confidence_reduced,
            )
        except Exception as exc:
            logger.error("[M6b] memory decay failed: %s", exc)
        return claims_decayed, claims_archived

    async def _dream_purge_soft_deleted(self, owner_id: str) -> int:
        """
        Step 5b: SOFT-DELETE PURGE
        Permanently remove claims that were soft-deleted more than 90 days ago.
        Returns count of permanently removed claims.
        """
        try:
            purged = await self.storage.purge_deleted_claims(retention_days=90)
            if purged > 0:
                logger.info("[M6b] purged %d expired soft-deleted claims (>90 days)", purged)
            return purged
        except Exception as exc:
            logger.error("[M6b] soft-delete purge failed: %s", exc)
            return 0

    async def _dream_entity_merges(self, owner_id: str) -> int:
        """
        Step 6: ENTITY MERGE SUGGESTIONS
        Run entity_resolver.suggest_merges() and store results.
        Returns count of suggestions generated.
        """
        if self.entity_resolver is None:
            logger.debug("[M6b] entity_resolver not configured, skipping merge suggestions")
            return 0

        try:
            suggestions = await self.entity_resolver.suggest_merges(owner_id=owner_id)
            # Suggestions are returned but not auto-applied (operator review required)
            # They are stored in-memory and available via the resolver's return value.
            # If storage supports a merge_suggestions table, we'd persist here.
            logger.debug(
                "[M6b] entity merge suggestions: %d pairs found", len(suggestions)
            )
            return len(suggestions)
        except Exception as exc:
            logger.error("[M6b] entity merge suggestions failed: %s", exc)
            return 0

    async def _dream_peer_refresh(
        self, owner_id: str
    ) -> tuple[int, list[dict]]:
        """
        Step 7: PEER REFRESH
        For each active peer (claims in last 7 days, max 20), refresh their
        profile representation via PeerManager.build_representation() and
        detect trait changes since the last dream cycle.

        Returns:
            (peers_refreshed, peer_changes) where peer_changes is a list of
            dicts {peer_id, name, changes} for peers whose traits changed.
        """
        if self.peer_manager is None:
            logger.debug("[M6b] peer_manager not configured, skipping peer refresh")
            return 0, []

        peers_refreshed = 0
        peer_changes: list[dict] = []

        try:
            active_peers = await self.peer_manager.get_active_peers(
                owner_id=owner_id, days=7, limit=20,
            )

            for peer_info in active_peers:
                peer_id = peer_info["peer_id"]
                display_name = peer_info["display_name"]
                entity_id = peer_info.get("entity_id")
                old_rep = peer_info.get("representation", {})

                try:
                    # Gather claims referencing this peer's entity
                    claims: list = []
                    if entity_id and hasattr(self.storage, "get_claims_by_entity"):
                        claims = await self.storage.get_claims_by_entity(
                            entity_id,
                            owner_id=owner_id,
                        )

                    if not claims:
                        logger.debug(
                            "[M6b] peer %s (%s): no active claims, skipping",
                            peer_id, display_name,
                        )
                        continue

                    # Build refreshed representation
                    new_rep = await self.peer_manager.build_representation(
                        peer_id, claims,
                    )
                    peers_refreshed += 1

                    # Detect trait changes
                    changes = self._detect_trait_changes(old_rep, new_rep)
                    if changes:
                        peer_changes.append({
                            "peer_id": peer_id,
                            "name": display_name,
                            "changes": changes,
                        })
                        logger.info(
                            "[M6b] peer %s (%s): %d trait changes detected",
                            peer_id, display_name, len(changes),
                        )

                except Exception as exc:
                    logger.warning(
                        "[M6b] peer refresh failed for %s (%s): %s",
                        peer_id, display_name, exc,
                    )

        except Exception as exc:
            logger.error("[M6b] peer refresh failed: %s", exc)

        logger.debug(
            "[M6b] peer refresh complete: %d refreshed, %d with changes",
            peers_refreshed, len(peer_changes),
        )
        return peers_refreshed, peer_changes

    @staticmethod
    def _detect_trait_changes(
        old_rep: dict, new_rep: dict
    ) -> list[str]:
        """Compare old and new representations to detect trait changes.

        Checks each of the four trait dimensions for additions.

        Returns:
            List of human-readable change descriptions.
        """
        changes: list[str] = []
        dimensions = [
            "communication_style",
            "topics_of_interest",
            "preferences",
            "relationship_dynamics",
        ]
        for dim in dimensions:
            old_values = set(
                v.lower() for v in old_rep.get(dim, [])
            )
            new_values = set(
                v.lower() for v in new_rep.get(dim, [])
            )
            added = new_values - old_values
            for val in sorted(added):
                changes.append(f"+{dim}: {val}")
        return changes

    async def _dream_brain_graph_refresh(self, owner_id: str) -> int:
        """
        Step 7: BRAIN GRAPH REFRESH
        Check for stale brain graph nodes (source hashes changed).
        Returns count of graphs checked (not rebuilt — rebuild is M10).
        """
        graphs_checked = 0
        try:
            # We check all brain graph nodes to detect staleness
            # Actual rebuild execution is M10 — we only flag here.
            # For now: count graphs we'd need to check (stub-safe).
            # Full implementation would call storage.get_stale_brain_graphs()
            # when that API is available.
            logger.debug("[M6b] brain graph refresh check (M10 handles actual rebuild)")
            graphs_checked = 0  # no-op until M10 API is available
        except Exception as exc:
            logger.error("[M6b] brain graph refresh failed: %s", exc)
        return graphs_checked

    async def _dream_apply_feedback(self, owner_id: str) -> int:
        """
        Step 8: APPLY PENDING FEEDBACK
        Process all memory_feedback records with status='pending'.
        Returns count of feedback records applied.
        """
        if self.feedback_system is None:
            # Fallback: use storage directly if feedback_system not injected
            try:
                pending = await self.storage.get_pending_feedback(owner_id=owner_id)
                applied = 0
                for record in pending:
                    try:
                        await self.storage.apply_feedback(record.id)
                        applied += 1
                    except Exception as exc:
                        logger.warning(
                            "[M6b] failed to apply feedback %s: %s", record.id, exc
                        )
                logger.debug("[M6b] feedback (direct): %d applied", applied)
                return applied
            except Exception as exc:
                logger.error("[M6b] feedback application failed: %s", exc)
                return 0

        try:
            applied = await self.feedback_system.process_pending(owner_id=owner_id)
            logger.debug("[M6b] feedback (FeedbackSystem): %d applied", applied)
            return applied
        except Exception as exc:
            logger.error("[M6b] FeedbackSystem.process_pending failed: %s", exc)
            return 0

    async def _dream_curiosity_and_journal(
        self, owner_id: str, dream_stats: dict
    ) -> Optional["CuriosityReport"]:
        """
        Step 9: CURIOSITY SEEDS + CROSS-MEMORY CONNECTIONS + DREAM JOURNAL (M6c)
        Runs CuriosityEngine on the most recent session's claims.
        Returns CuriosityReport (or None on failure).
        """
        try:
            from cortex.circadian.curiosity import CuriosityEngine  # local import — no circular dep
            engine = CuriosityEngine(
                storage=self.storage,
                llm=self.llm,
                config=self.config,
            )
            # Fetch recent session claims for this owner
            try:
                active = await self.storage.get_active_session(owner_id=owner_id)
                session_claims: List = []
                if active:
                    session_claims = await self.storage.get_session_claims(
                        active.session_id
                    )
                else:
                    # Fallback: get claims from status='active' without a specific session
                    session_claims = await self.storage.get_claims_by_status(
                        "active", owner_id=owner_id
                    )
            except Exception as exc:
                logger.warning("[M6c] could not fetch session claims: %s", exc)
                session_claims = []

            report = await engine.run(
                owner_id=owner_id,
                session_claims=session_claims,
                dream_stats=dream_stats,
            )
            return report
        except Exception as exc:
            logger.error("[M6b] _dream_curiosity_and_journal failed: %s", exc)
            return None

