"""cortex/circadian/consolidation/dream_reconsolidation.py — Heavy dream steps."""
from __future__ import annotations

import asyncio, json, logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, List

from cortex.circadian.consolidation._helpers import (
    _parse_dt, _parse_json_list, _parse_json_dict, MAX_CLAIMS_PER_BATCH,
)
from cortex.circadian.sleep import _CONTRADICTION_SYSTEM, _CONTRADICTION_USER
from cortex.types import SEMANTIC_MEMORY_TYPES

if TYPE_CHECKING:
    from cortex.types import SemanticClaim

logger = logging.getLogger(__name__)


class DreamReconsolidationMixin:
    """Extraction queue retry + contradiction detection/resolution."""

    async def _dream_retry_extraction_queue(self, owner_id: str) -> int:
        """
        Step 0: EXTRACTION QUEUE RETRY
        Process pending extraction queue entries from previous failures.
        Re-runs the extraction pipeline on queued conversation payloads.
        Returns count of successfully retried entries.
        """
        retried = 0
        try:
            from cortex.capture.extraction_queue import (
                get_pending_entries,
                mark_processing,
                mark_done,
                mark_failed,
            )
            entries = await get_pending_entries(self.storage, limit=10)
            if not entries:
                return 0

            logger.info("[M6b] extraction queue: %d pending entries to retry", len(entries))

            for entry in entries:
                try:
                    await mark_processing(self.storage, entry.id)
                    import json as _json
                    payload = _json.loads(entry.payload)
                    conversation = payload.get("conversation", [])

                    if not conversation:
                        await mark_done(self.storage, entry.id)
                        continue

                    # Build extraction pipeline — use the consolidator's wired components
                    from cortex.capture.wake_pipeline import run_extraction
                    import asyncio

                    if not hasattr(self, "_extraction_semaphore"):
                        self._extraction_semaphore = asyncio.Semaphore(2)

                    # Get extractor, reconciler etc from the storage/config
                    # These are the same components wired in CortexPlugin
                    from cortex.extraction.extractor import ExtractionPipeline

                    extractor = ExtractionPipeline(llm=self.llm, config=self.config)

                    # Simple reconciler creation for queue retry
                    reconciler = None
                    try:
                        from cortex.reconciliation.engine import ReconciliationEngine
                        reconciler = ReconciliationEngine(
                            storage=self.storage,
                            llm=self.llm,
                        )
                    except Exception as recon_exc:
                        logger.warning("[M6b] queue retry: reconciler init failed: %s", recon_exc)

                    errors: list = []
                    warnings: list = []

                    outcome = await run_extraction(
                        extractor=extractor,
                        storage=self.storage,
                        llm=self.llm,
                        entity_resolver=None,
                        cornerstone_guardian=self.cornerstone_guardian,
                        reconciler=reconciler,
                        event_emitter=None,
                        config=self.config,
                        extraction_semaphore=self._extraction_semaphore,
                        conversation=conversation,
                        sid=payload.get("session_id", entry.session_id or "dream-retry"),
                        db_session_pk=None,
                        effective_owner_id=entry.owner_id,
                        episode_event_id=None,
                        budget_check=lambda n: None,
                        budget_consume=lambda n: None,
                        errors=errors,
                        warnings=warnings,
                        zero_claim_outcomes=0,
                        source_session_id=payload.get("source_session_id"),
                        source_message_id=payload.get("source_message_id"),
                        source_channel=payload.get("source_channel"),
                        source_type=payload.get("source_type"),
                    )

                    if outcome.ok:
                        await mark_done(self.storage, entry.id)
                        retried += 1
                        logger.info(
                            "[M6b] queue retry success: id=%s claims=%d",
                            entry.id, outcome.claims_stored,
                        )
                    else:
                        error_msg = "; ".join(outcome.errors[:3])
                        await mark_failed(self.storage, entry.id, error_msg)
                        logger.warning(
                            "[M6b] queue retry failed again: id=%s error=%s",
                            entry.id, error_msg[:100],
                        )

                except Exception as entry_exc:
                    logger.error(
                        "[M6b] queue retry error for entry %s: %s",
                        entry.id, entry_exc,
                    )
                    try:
                        await mark_failed(self.storage, entry.id, str(entry_exc)[:500])
                    except Exception:
                        pass

        except Exception as exc:
            logger.error("[M6b] extraction queue retry step failed: %s", exc)

        return retried


    async def _dream_reconsolidation(self, owner_id: str) -> int:
        """
        Step 4: RECONSOLIDATION
        Detect contradictions in semantic memory, resolve via LLM evidence.
        Returns count of contradictions resolved.
        """
        contradictions_resolved = 0
        try:
            # Get all active semantic claims, grouped by subject
            all_claims = await self.storage.get_claims_by_status("active", owner_id=owner_id)
            semantic_claims = [c for c in all_claims if c.memory_type in SEMANTIC_MEMORY_TYPES]

            # Group by subject
            by_subject: dict[str, list] = {}
            for claim in semantic_claims:
                by_subject.setdefault(claim.subject, []).append(claim)

            for subject, claims in by_subject.items():
                if len(claims) < 2:
                    continue  # can't have contradictions with one claim
                try:
                    result = await self._detect_and_resolve_contradictions(subject, claims)
                    contradictions_resolved += result
                except Exception as exc:
                    logger.warning(
                        "[M6b] reconsolidation failed for subject %r: %s", subject, exc
                    )

        except Exception as exc:
            logger.error("[M6b] reconsolidation scan failed: %s", exc)
        return contradictions_resolved

    async def _detect_and_resolve_contradictions(
        self,
        subject: str,
        claims: List["SemanticClaim"],
    ) -> int:
        """Check one subject group for contradictions and resolve them.

        Issue #594: Claims are processed in batches of at most
        ``self._contradiction_batch_size`` to prevent context-window blow-out.
        """
        batch_size = self._contradiction_batch_size

        def _sort_key(c):
            return (getattr(c, "created_at", "") or "", getattr(c, "confidence", 0.0))

        sorted_claims = sorted(claims, key=_sort_key, reverse=True)

        # Cap total claims sent to LLM to prevent context-window blow-out (#594)
        max_total = batch_size
        if len(sorted_claims) > max_total:
            logger.warning(
                "Truncating %d claims to %d for contradiction analysis (subject: %s)",
                len(sorted_claims), max_total, subject,
            )
            sorted_claims = sorted_claims[:max_total]

        batches = [sorted_claims[i: i + batch_size] for i in range(0, len(sorted_claims), batch_size)]
        all_contradictions: List[Any] = []

        from cortex.llm.router import resolve_routed_model
        _contra_model = resolve_routed_model(self.llm, "contradiction_analysis")

        for batch in batches:
            claims_json = json.dumps(
                [{"id": c.id, "predicate": c.predicate, "object": c.object_value,
                  "confidence": c.confidence, "created_at": c.created_at}
                 for c in batch],
                ensure_ascii=False, indent=2,
            )
            _resp = await self.llm.complete(
                system=_CONTRADICTION_SYSTEM,
                user=_CONTRADICTION_USER.format(subject=subject, claims_json=claims_json),
                response_format="json",
                **({"model": _contra_model} if _contra_model else {}),
            )
            response = _resp.text if hasattr(_resp, "text") else str(_resp)
            data = _parse_json_dict(response)
            all_contradictions.extend(data.get("contradictions", []))

        contradictions = all_contradictions
        resolved = 0

        for contradiction in contradictions:
            resolution = contradiction.get("resolution", "flag_both")
            claim_a_id = contradiction.get("claim_a_id")
            claim_b_id = contradiction.get("claim_b_id")
            reason = contradiction.get("reason", "")

            if not claim_a_id or not claim_b_id:
                continue

            try:
                now = datetime.now(timezone.utc).isoformat()
                if resolution == "supersede_a":
                    # claim_a is superseded by claim_b
                    await self.storage.update_claim(
                        claim_a_id,
                        status="deprecated",
                        superseded_by=claim_b_id,
                        updated_at=now,
                    )
                    # Changelog: dream reconsolidation updated claim_a
                    try:
                        history = await self.storage.get_changelog(claim_a_id)
                        next_ver = len(history) + 1
                        await self.storage.log_changelog(
                            claim_id=claim_a_id,
                            version=next_ver,
                            content_snapshot=None,
                            action="updated",
                            reason=f"dream reconsolidation: {reason}",
                            changed_by="dream_cycle",
                            diff_summary=f"superseded by {claim_b_id}: {reason}",
                        )
                    except Exception as _exc:
                        logger.warning("[M6b] changelog log failed for %s: %s", claim_a_id, _exc)
                    resolved += 1
                    logger.info(
                        "[M6b] claim %s superseded by %s (%s)", claim_a_id, claim_b_id, reason
                    )
                elif resolution == "supersede_b":
                    # claim_b is superseded by claim_a
                    await self.storage.update_claim(
                        claim_b_id,
                        status="deprecated",
                        superseded_by=claim_a_id,
                        updated_at=now,
                    )
                    # Changelog: dream reconsolidation updated claim_b
                    try:
                        history = await self.storage.get_changelog(claim_b_id)
                        next_ver = len(history) + 1
                        await self.storage.log_changelog(
                            claim_id=claim_b_id,
                            version=next_ver,
                            content_snapshot=None,
                            action="updated",
                            reason=f"dream reconsolidation: {reason}",
                            changed_by="dream_cycle",
                            diff_summary=f"superseded by {claim_a_id}: {reason}",
                        )
                    except Exception as _exc:
                        logger.warning("[M6b] changelog log failed for %s: %s", claim_b_id, _exc)
                    resolved += 1
                    logger.info(
                        "[M6b] claim %s superseded by %s (%s)", claim_b_id, claim_a_id, reason
                    )
                else:
                    # flag_both — reduce confidence on both, mark for review
                    for cid in (claim_a_id, claim_b_id):
                        matching = [c for c in claims if c.id == cid]
                        if matching:
                            old_conf = matching[0].confidence
                            new_conf = max(0.3, old_conf * 0.7)
                            await self.storage.update_claim(
                                cid,
                                confidence=new_conf,
                                updated_at=now,
                            )
                            # Changelog: confidence reduced by dream cycle
                            try:
                                history = await self.storage.get_changelog(cid)
                                next_ver = len(history) + 1
                                await self.storage.log_changelog(
                                    claim_id=cid,
                                    version=next_ver,
                                    content_snapshot=None,
                                    action="updated",
                                    reason=f"dream reconsolidation flag_both: {reason}",
                                    changed_by="dream_cycle",
                                    diff_summary=f"confidence {old_conf:.3f} → {new_conf:.3f}",
                                )
                            except Exception as _exc:
                                logger.warning("[M6b] changelog log failed for %s: %s", cid, _exc)
                    logger.info(
                        "[M6b] contradiction flagged (both reduced): %s vs %s (%s)",
                        claim_a_id, claim_b_id, reason,
                    )
                    # Flagging both still counts as a resolution action
                    resolved += 1

            except Exception as exc:
                logger.warning(
                    "[M6b] could not resolve contradiction %s vs %s: %s",
                    claim_a_id, claim_b_id, exc,
                )

        return resolved

