"""
cortex/identity/drift_calculator.py — Drift Calculator (Module 11 / M11).

Detects when cornerstone memories have drifted from their golden (sealed) copies.
This is the integrity watchdog — it runs periodically (during dream cycles or
on explicit check) and computes a composite drift score for each cornerstone.

Drift score composition (weighted average):
    - Semantic similarity (embedding cosine distance): 40%
      → Catches meaning changes even when words are reshuffled
    - Length ratio (compression detection): 30%
      → Catches summarization/truncation that loses detail
    - Key phrase preservation: 30%
      → Catches loss of specific important terms/phrases

Score interpretation and thresholds (configurable in CortexConfig):
    score < 0.1    → 'none'    (healthy, no action)
    0.1 ≤ s < 0.3  → 'warn'    (log warning, monitor)
    0.3 ≤ s < 0.5  → 'restore' (auto-restore from golden copy)
    s ≥ 0.5        → 'alert'   (restore + alert operator)

Design principles:
    - Drift detected → action is proportional to severity
    - 'restore' auto-restores content to golden_copy (safe because golden is immutable)
    - 'alert' always involves operator notification (never silent)
    - All drift measurements are logged to drift_log table for audit trail
    - Embeddings are optional — if unavailable, semantic component scores 0.0
      and the remaining components are reweighted

Dependencies:
    - cortex.storage.adapter.StorageAdapter (M1) — read/write cornerstones, drift_log
    - cortex.llm.provider.LLMProvider — embed() for semantic similarity (optional)
    - cortex.config.CortexConfig — drift thresholds
    - cortex.identity.cornerstone.CornerstoneMemory — cornerstone data type

Used by:
    - Dream cycle orchestrator — calls check_all() periodically
    - CortexPlugin (M13) — manual drift check via API
"""

from __future__ import annotations

import logging
import math
import re
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING, List, Optional

if TYPE_CHECKING:
    from cortex.storage.adapter import StorageAdapter
    from cortex.config import CortexConfig
    from cortex.llm.provider import LLMProvider
    from cortex.identity.cornerstone import CornerstoneMemory

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------

@dataclass
class DriftReport:
    """
    Result of a single drift check against a cornerstone's golden copy.

    drift_score: 0.0 = identical, 1.0 = completely different
    action:
        'none'    → drift_score < 0.1   (healthy)
        'warn'    → 0.1 ≤ drift_score < 0.3   (log, monitor)
        'restore' → 0.3 ≤ drift_score < 0.5   (auto-restore)
        'alert'   → drift_score ≥ 0.5   (restore + alert operator)
    """
    cornerstone_id: str
    label: str
    drift_score: float          # 0.0=identical, 1.0=completely different
    semantic_sim: float         # cosine similarity (0-1, higher = more similar)
    length_ratio: float         # len(current) / len(golden), clamped to [0,1]
    phrase_preservation: float  # fraction of key phrases still present (0-1)
    action: str                 # 'none' | 'warn' | 'restore' | 'alert'
    checked_at: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _cosine_similarity(a: List[float], b: List[float]) -> float:
    """
    Compute cosine similarity between two vectors.
    Returns value in [-1, 1]; higher = more similar.
    Returns 0.0 on zero-norm vectors.
    """
    if len(a) != len(b):
        raise ValueError(f"Embedding dimension mismatch: {len(a)} vs {len(b)}")
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)


def _length_ratio_score(current: str, golden: str) -> float:
    """
    Compute a length-based drift component.

    Returns a score in [0, 1] where 1.0 = same length (no compression/expansion),
    and 0.0 = completely empty or extremely different length.

    Formula: 1.0 - abs(len_ratio - 1.0), clamped to [0, 1].
    len_ratio = len(current) / len(golden)
    """
    if not golden:
        return 1.0 if not current else 0.0
    len_ratio = len(current) / len(golden)
    return max(0.0, min(1.0, 1.0 - abs(len_ratio - 1.0)))


def _phrase_preservation_score(current: str, phrases: List[str]) -> float:
    """
    Compute what fraction of key phrases from the golden copy are still
    present in the current text. Case-insensitive substring match.

    Returns 1.0 if there are no key phrases (nothing to lose).
    """
    if not phrases:
        return 1.0
    current_lower = current.lower()
    preserved = sum(1 for p in phrases if p.lower() in current_lower)
    return preserved / len(phrases)


def _determine_action(
    drift_score: float,
    warn_threshold: float = 0.1,
    restore_threshold: float = 0.3,
    alert_threshold: float = 0.5,
) -> str:
    """Map a drift score to an action string using configurable thresholds."""
    if drift_score < warn_threshold:
        return "none"
    if drift_score < restore_threshold:
        return "warn"
    if drift_score < alert_threshold:
        return "restore"
    return "alert"


# ---------------------------------------------------------------------------
# DriftCalculator
# ---------------------------------------------------------------------------

class DriftCalculator:
    """
    Measure and respond to drift in cornerstone memories.

    Drift is quantified as a composite score:
      - Semantic similarity via embedding cosine distance:  40%
      - Length ratio (detects compression/expansion):       30%
      - Key phrase preservation:                            30%

    Final drift_score = 1.0 means no similarity detected at all.
    Final drift_score = 0.0 means perfectly preserved.
    """

    def __init__(
        self,
        storage: "StorageAdapter",
        config: "CortexConfig",
        llm: "LLMProvider",
        event_emitter: object = None,
    ) -> None:
        self.storage = storage
        self.config = config
        self.llm = llm
        self.event_emitter = event_emitter

        # Composite scoring weights (must sum to 1.0)
        self.similarity_weight: float = 0.4
        self.length_weight: float = 0.3
        self.phrase_weight: float = 0.3

        # Action thresholds — pulled from config in check_one; defaults here
        self._default_warn_threshold: float = 0.1
        self._default_restore_threshold: float = 0.3
        self._default_alert_threshold: float = 0.5

    # ------------------------------------------------------------------ #
    # Private helpers
    # ------------------------------------------------------------------ #

    def _get_thresholds(self):
        """
        Return (warn, restore, alert) thresholds from config.

        Threshold mapping (note: original config names are swapped relative
        to their semantic meaning — see CortexConfig for properly named aliases):
            warn    = 0.1 (hardcoded, not in config)
            restore = drift_check_threshold  (0.3) = config.drift_restore_score
            alert   = drift_restore_threshold (0.5) = config.drift_alert_score

        Falls back to hardcoded defaults if config attributes are missing.
        """
        warn = getattr(
            self.config, "drift_warn_threshold", self._default_warn_threshold
        )
        # Use properly named aliases when available, fall back to legacy names
        restore = getattr(
            self.config, "drift_restore_score",
            getattr(self.config, "drift_check_threshold", self._default_restore_threshold),
        )
        alert = getattr(
            self.config, "drift_alert_score",
            getattr(self.config, "drift_restore_threshold", self._default_alert_threshold),
        )
        return warn, restore, alert

    async def _get_embedding(self, text: str) -> List[float]:
        """Get a single text embedding from the LLM provider."""
        response = await self.llm.embed(
            texts=[text],
            model=self.config.embedding_model,
            input_type="document",
        )
        return response.embeddings[0]

    # ------------------------------------------------------------------ #
    # Public interface
    # ------------------------------------------------------------------ #

    async def extract_key_phrases(self, text: str) -> List[str]:
        """
        Extract identity-defining phrases from cornerstone text.

        Uses LLM to identify the most distinctive phrases.
        Falls back to simple heuristics if LLM call fails.
        """
        prompt = (
            "Extract 5-10 key phrases from the following text that are most "
            "identity-defining, distinctive, or semantically important. "
            "Return one phrase per line, nothing else.\n\n"
            f"TEXT:\n{text[:2000]}"
        )
        try:
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(
                self.llm, "key_phrase_extraction", self.config.extraction_model,
            )
            response = await self.llm.complete(
                system=(
                    "You are a precise text analyst. Extract key phrases exactly "
                    "as they appear in the source text."
                ),
                user=prompt,
                model=_routed_model,
                temperature=0.0,
                max_tokens=256,
            )
            text_out = response.text if hasattr(response, "text") else str(response)
            phrases = [
                line.strip()
                for line in text_out.strip().splitlines()
                if line.strip()
            ]
            return phrases[:10]
        except Exception as exc:  # pragma: no cover
            logger.warning("Key phrase extraction failed, using fallback: %s", exc)
            # Fallback: split on punctuation, take unique non-trivial segments
            segments = re.split(r"[.!?;]", text)
            phrases = []
            for seg in segments:
                seg = seg.strip()
                if len(seg) > 20:
                    phrases.append(seg[:80])
                if len(phrases) >= 8:
                    break
            return phrases

    async def check_one(self, cs_id: str) -> DriftReport:
        """
        Compare a cornerstone's current content against its golden copy.

        Composite score:
          - Semantic similarity (embedding cosine distance): 40%
          - Length ratio (compression detection): 30%
          - Key phrase preservation: 30%

        Logs result to drift_log and updates cornerstone drift_score in storage.
        """
        # Load the cornerstone
        cornerstones = await self.storage.get_cornerstones(owner_id="", status="active")
        cs: Optional["CornerstoneMemory"] = next(
            (c for c in cornerstones if c.id == cs_id), None
        )
        if cs is None:
            raise ValueError(f"Cornerstone not found: {cs_id}")

        golden = cs.golden_copy
        current = cs.content

        # ── Semantic similarity ──────────────────────────────────────────
        try:
            golden_emb = await self._get_embedding(golden)
            current_emb = await self._get_embedding(current)
            raw_sim = _cosine_similarity(golden_emb, current_emb)
            # Cosine sim is in [-1, 1]; clamp to [0, 1]
            semantic_sim = max(0.0, min(1.0, (raw_sim + 1.0) / 2.0))
            # Normalize to [0, 1] with 1 = perfectly similar
            # But we treat embeddings in the [0, 1] space already
            # Most embedding outputs are in [0, 1] cos sim for similar docs
            semantic_sim = max(0.0, min(1.0, raw_sim))
        except Exception as exc:
            logger.warning(
                "Embedding call failed for cs %s, treating as 0 similarity: %s",
                cs_id,
                exc,
            )
            semantic_sim = 0.0

        # ── Length ratio ─────────────────────────────────────────────────
        length_ratio = _length_ratio_score(current, golden)

        # ── Key phrase preservation ──────────────────────────────────────
        try:
            key_phrases = await self.extract_key_phrases(golden)
            phrase_preservation = _phrase_preservation_score(current, key_phrases)
        except Exception as exc:
            logger.warning(
                "Phrase extraction failed for cs %s: %s", cs_id, exc
            )
            phrase_preservation = 1.0

        # ── Composite drift score ────────────────────────────────────────
        # semantic_sim = 1.0 → no semantic drift
        # length_ratio = 1.0 → no length drift
        # phrase_preservation = 1.0 → all phrases intact
        # drift contribution from each = (1 - metric)
        drift_score = (
            self.similarity_weight * (1.0 - semantic_sim)
            + self.length_weight * (1.0 - length_ratio)
            + self.phrase_weight * (1.0 - phrase_preservation)
        )
        drift_score = max(0.0, min(1.0, drift_score))

        warn_t, restore_t, alert_t = self._get_thresholds()
        action = _determine_action(drift_score, warn_t, restore_t, alert_t)
        checked_at = datetime.now(timezone.utc).isoformat()

        # ── Persist results ──────────────────────────────────────────────
        try:
            await self.storage.update_cornerstone_drift(cs_id, drift_score)
        except Exception as exc:
            logger.warning("Failed to update cornerstone drift score: %s", exc)

        try:
            await self.storage.log_drift(
                cornerstone_id=cs_id,
                drift_score=drift_score,
                semantic_sim=semantic_sim,
                length_ratio=length_ratio,
                phrase_preservation=phrase_preservation,
                action_taken=action,
                current_text=current,
                golden_text=golden,
                checked_at=checked_at,
            )
        except Exception as exc:
            logger.warning("Failed to log drift: %s", exc)

        report = DriftReport(
            cornerstone_id=cs_id,
            label=cs.label,
            drift_score=drift_score,
            semantic_sim=semantic_sim,
            length_ratio=length_ratio,
            phrase_preservation=phrase_preservation,
            action=action,
            checked_at=checked_at,
        )

        if action in ("restore", "alert"):
            logger.warning(
                "Drift action=%r for cornerstone %r (score=%.3f) — restoring.",
                action,
                cs.label,
                drift_score,
            )
            # Emit CORNERSTONE_DRIFT_DETECTED event
            if self.event_emitter is not None:
                try:
                    from cortex.events.emitter import CortexEvent, EventType
                    await self.event_emitter.emit(CortexEvent(
                        event_type=EventType.CORNERSTONE_DRIFT_DETECTED,
                        owner_id=cs.owner_id if hasattr(cs, "owner_id") else "default",
                        payload={
                            "cornerstone_id": cs_id,
                            "label": cs.label,
                            "drift_score": drift_score,
                            "action": action,
                        },
                    ))
                except Exception as emit_exc:
                    logger.debug("Event emit failed (CORNERSTONE_DRIFT_DETECTED): %s", emit_exc)
            try:
                await self.restore(cs_id)
            except Exception as exc:
                logger.error("Failed to restore cornerstone %s: %s", cs_id, exc)

        if action == "alert":
            logger.error(
                "CRITICAL DRIFT ALERT: cornerstone %r (id=%s) score=%.3f exceeds "
                "alert threshold %.2f. Restore attempted.",
                cs.label,
                cs_id,
                drift_score,
                alert_t,
            )

        return report

    async def check_all(self, owner_id: str = "default") -> List[DriftReport]:
        """
        Check all active cornerstones for the given owner.

        Designed to be called during sleep/dream cycles.
        Returns a list of DriftReport — one per cornerstone.
        Reports with action != 'none' are logged at WARNING level.
        """
        cornerstones = await self.storage.get_cornerstones(
            owner_id=owner_id, status="active"
        )

        if not cornerstones:
            logger.debug("No active cornerstones found for owner_id=%r", owner_id)
            return []

        reports: List[DriftReport] = []
        for cs in cornerstones:
            try:
                report = await self.check_one(cs.id)
                reports.append(report)
                if report.action != "none":
                    logger.warning(
                        "Drift detected: label=%r, score=%.3f, action=%r",
                        report.label,
                        report.drift_score,
                        report.action,
                    )
            except Exception as exc:
                logger.error(
                    "check_one failed for cornerstone %s: %s", cs.id, exc
                )

        return reports

    async def restore(self, cs_id: str) -> "CornerstoneMemory":
        """
        Restore a cornerstone from its golden copy.

        Delegates to storage.restore_cornerstone() which resets content
        to golden_copy and zeroes out drift_score.
        Logs the restore action to drift_log.
        """
        restored = await self.storage.restore_cornerstone(cs_id)

        try:
            await self.storage.log_drift(
                cornerstone_id=cs_id,
                drift_score=0.0,
                semantic_sim=1.0,
                length_ratio=1.0,
                phrase_preservation=1.0,
                action_taken="restored",
                current_text=restored.content,
                golden_text=restored.golden_copy,
                checked_at=datetime.now(timezone.utc).isoformat(),
            )
        except Exception as exc:
            logger.warning("Failed to log restore action: %s", exc)

        # Changelog: cornerstone restored from golden copy
        try:
            history = await self.storage.get_changelog(cs_id)
            next_ver = len(history) + 1
            await self.storage.log_changelog(
                claim_id=cs_id,
                version=next_ver,
                content_snapshot=restored.content,
                action="restored",
                reason="drift detected — restored from golden copy",
                changed_by="drift_restore",
                diff_summary=None,
            )
        except Exception as exc:
            logger.warning("Failed to log changelog for cornerstone restore %s: %s", cs_id, exc)

        logger.info(
            "Cornerstone %r (id=%s) restored from golden copy.", restored.label, cs_id
        )
        return restored
