"""
Tests for MODULE 11: DriftCalculator

All storage and LLM operations use mock adapters — no real DB or API calls.
Tests verify: composite scoring, action thresholds, check_one, check_all,
restore, key phrase extraction, and edge cases.
"""

from __future__ import annotations

import math
import pytest
from dataclasses import dataclass, field
from typing import Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch

from cortex.identity.drift_calculator import (
    DriftCalculator,
    DriftReport,
    _cosine_similarity,
    _length_ratio_score,
    _phrase_preservation_score,
    _determine_action,
)


# ---------------------------------------------------------------------------
# Helpers / fixtures
# ---------------------------------------------------------------------------

def _make_cornerstone(
    cs_id: str = "cs-0001",
    label: str = "test-cs",
    content: str = "Eva is an AI companion who loves building systems.",
    golden_copy: Optional[str] = None,
    owner_id: str = "default",
):
    """Return a minimal CornerstoneMemory-like object."""
    from cortex.identity.cornerstone import CornerstoneMemory
    gc = golden_copy if golden_copy is not None else content
    return CornerstoneMemory(
        id=cs_id,
        owner_id=owner_id,
        label=label,
        content=content,
        golden_copy=gc,
        golden_hash="abc",
        current_hash="abc",
        drift_score=0.0,
        last_drift_check=None,
        token_count=10,
        injection_order=0,
    )


@dataclass
class MockConfig:
    embedding_model: str = "text-embedding-3-small"
    extraction_model: str = "gpt-4.1-nano-2025-04-14"
    drift_warn_threshold: float = 0.1
    drift_check_threshold: float = 0.3    # restore threshold
    drift_restore_threshold: float = 0.5  # alert threshold


class MockStorageAdapter:
    def __init__(self, cornerstones=None):
        self._cornerstones: Dict[str, object] = {}
        for cs in (cornerstones or []):
            self._cornerstones[cs.id] = cs
        self.log_drift_calls: List[dict] = []
        self.update_drift_calls: List[tuple] = []
        self.restore_calls: List[str] = []

    async def get_cornerstones(self, owner_id="default", status="active"):
        if owner_id == "":
            return list(self._cornerstones.values())
        return [c for c in self._cornerstones.values() if c.owner_id == owner_id]

    async def update_cornerstone_drift(self, cs_id: str, drift_score: float):
        self.update_drift_calls.append((cs_id, drift_score))
        if cs_id in self._cornerstones:
            self._cornerstones[cs_id].drift_score = drift_score

    async def restore_cornerstone(self, cs_id: str):
        cs = self._cornerstones[cs_id]
        cs.content = cs.golden_copy
        cs.drift_score = 0.0
        self.restore_calls.append(cs_id)
        return cs

    async def log_drift(self, **kwargs):
        self.log_drift_calls.append(kwargs)


class MockLLMProvider:
    def __init__(self, embedding_value: float = 0.99):
        """
        embedding_value: the value used in the identity embedding vector.
        Matching vectors → cosine sim = 1.0 (no drift).
        """
        self._embedding_value = embedding_value
        self.embed_calls: List = []
        self.complete_calls: List = []

    async def embed(self, request=None, **kwargs):
        self.embed_calls.append(request or kwargs)
        # Return a simple unit vector
        vec = [self._embedding_value] * 4
        # Normalise
        norm = math.sqrt(sum(x * x for x in vec))
        vec = [x / norm for x in vec]

        # Return two identical embeddings for [golden, current] or one
        class FakeEmbedResponse:
            embeddings = [vec]
            model = "text-embedding-3-small"
            tokens_in = 10
            tokens_out = 0
            cost_estimate = 0.0
            latency_ms = 1.0

        return FakeEmbedResponse()

    async def complete(self, request=None, **kwargs):
        self.complete_calls.append(request or kwargs)

        class FakeCompletionResponse:
            text = "Eva is an AI companion\nloves building systems\nloyal sharp direct"
            model = "gpt-4.1-nano-2025-04-14"
            tokens_in = 50
            tokens_out = 20
            cost_estimate = 0.0
            latency_ms = 1.0

        return FakeCompletionResponse()


def _make_calculator(cornerstones=None, embedding_value=0.99):
    storage = MockStorageAdapter(cornerstones)
    config = MockConfig()
    llm = MockLLMProvider(embedding_value=embedding_value)
    return DriftCalculator(storage, config, llm), storage, llm


# ---------------------------------------------------------------------------
# Unit tests — pure helper functions
# ---------------------------------------------------------------------------

class TestCosimeSimilarity:
    def test_identical_vectors(self):
        a = [1.0, 0.0, 0.0]
        assert _cosine_similarity(a, a) == pytest.approx(1.0)

    def test_orthogonal_vectors(self):
        a = [1.0, 0.0]
        b = [0.0, 1.0]
        assert _cosine_similarity(a, b) == pytest.approx(0.0)

    def test_opposite_vectors(self):
        a = [1.0, 0.0]
        b = [-1.0, 0.0]
        assert _cosine_similarity(a, b) == pytest.approx(-1.0)

    def test_zero_vector_returns_zero(self):
        a = [0.0, 0.0]
        b = [1.0, 0.0]
        assert _cosine_similarity(a, b) == 0.0

    def test_dimension_mismatch(self):
        with pytest.raises(ValueError):
            _cosine_similarity([1.0], [1.0, 2.0])


class TestLengthRatioScore:
    def test_identical_lengths(self):
        s = "hello world"
        assert _length_ratio_score(s, s) == pytest.approx(1.0)

    def test_empty_current(self):
        assert _length_ratio_score("", "hello world") == pytest.approx(0.0)

    def test_empty_golden(self):
        assert _length_ratio_score("", "") == pytest.approx(1.0)
        assert _length_ratio_score("hello", "") == pytest.approx(0.0)

    def test_half_length(self):
        # current is half golden → ratio=0.5 → score = 1 - 0.5 = 0.5
        golden = "a" * 100
        current = "a" * 50
        assert _length_ratio_score(current, golden) == pytest.approx(0.5)

    def test_double_length(self):
        # current is double golden → ratio=2.0 → score = 1 - 1.0 = 0.0
        golden = "a" * 50
        current = "a" * 100
        assert _length_ratio_score(current, golden) == pytest.approx(0.0)


class TestPhrasePreservationScore:
    def test_all_preserved(self):
        phrases = ["Eva is an AI", "loves building systems"]
        text = "Eva is an AI companion who loves building systems and more."
        assert _phrase_preservation_score(text, phrases) == pytest.approx(1.0)

    def test_none_preserved(self):
        phrases = ["completely different phrase", "another missing thing"]
        text = "hello world"
        assert _phrase_preservation_score(text, phrases) == pytest.approx(0.0)

    def test_partial_preservation(self):
        phrases = ["Eva is an AI", "missing phrase"]
        text = "Eva is an AI companion."
        assert _phrase_preservation_score(text, phrases) == pytest.approx(0.5)

    def test_empty_phrases(self):
        assert _phrase_preservation_score("any text", []) == pytest.approx(1.0)

    def test_case_insensitive(self):
        phrases = ["EVA IS AN AI"]
        text = "eva is an ai companion"
        assert _phrase_preservation_score(text, phrases) == pytest.approx(1.0)


class TestDetermineAction:
    def test_none_zone(self):
        assert _determine_action(0.0) == "none"
        assert _determine_action(0.05) == "none"
        assert _determine_action(0.09) == "none"

    def test_warn_zone(self):
        assert _determine_action(0.1) == "warn"
        assert _determine_action(0.2) == "warn"
        assert _determine_action(0.29) == "warn"

    def test_restore_zone(self):
        assert _determine_action(0.3) == "restore"
        assert _determine_action(0.4) == "restore"
        assert _determine_action(0.49) == "restore"

    def test_alert_zone(self):
        assert _determine_action(0.5) == "alert"
        assert _determine_action(0.75) == "alert"
        assert _determine_action(1.0) == "alert"

    def test_custom_thresholds(self):
        # with threshold (0.05, 0.2, 0.4)
        assert _determine_action(0.03, warn_threshold=0.05) == "none"
        assert _determine_action(0.06, warn_threshold=0.05, restore_threshold=0.2) == "warn"
        assert _determine_action(0.25, warn_threshold=0.05, restore_threshold=0.2, alert_threshold=0.4) == "restore"
        assert _determine_action(0.45, warn_threshold=0.05, restore_threshold=0.2, alert_threshold=0.4) == "alert"


# ---------------------------------------------------------------------------
# DriftCalculator — integration-style tests
# ---------------------------------------------------------------------------

class TestDriftCalculatorCheckOne:
    @pytest.mark.asyncio
    async def test_no_drift_identical_content(self):
        """When current == golden, drift score should be near 0."""
        text = "Eva is an AI companion who loves building systems."
        cs = _make_cornerstone(content=text, golden_copy=text)
        calc, storage, llm = _make_calculator(cornerstones=[cs], embedding_value=0.99)

        report = await calc.check_one(cs.id)

        assert isinstance(report, DriftReport)
        assert report.cornerstone_id == cs.id
        assert report.label == cs.label
        assert report.drift_score >= 0.0
        assert report.drift_score <= 1.0
        # With identical embeddings and text, drift should be very low
        # (≤ 0.1 — the boundary itself corresponds to action='none')
        assert report.drift_score <= 0.1

    @pytest.mark.asyncio
    async def test_check_one_logs_drift(self):
        """check_one should log to drift_log exactly once."""
        text = "Eva is an AI companion."
        cs = _make_cornerstone(content=text, golden_copy=text)
        calc, storage, _ = _make_calculator(cornerstones=[cs])

        await calc.check_one(cs.id)

        # At least one drift log entry
        assert len(storage.log_drift_calls) >= 1
        logged = storage.log_drift_calls[0]
        assert logged["cornerstone_id"] == cs.id
        assert "drift_score" in logged

    @pytest.mark.asyncio
    async def test_check_one_updates_cornerstone_drift_score(self):
        """check_one should call update_cornerstone_drift."""
        text = "Eva is an AI companion."
        cs = _make_cornerstone(content=text, golden_copy=text)
        calc, storage, _ = _make_calculator(cornerstones=[cs])

        await calc.check_one(cs.id)

        assert len(storage.update_drift_calls) == 1
        cs_id, drift_score = storage.update_drift_calls[0]
        assert cs_id == cs.id
        assert isinstance(drift_score, float)

    @pytest.mark.asyncio
    async def test_check_one_raises_for_unknown_cs(self):
        """check_one should raise ValueError for unknown cornerstone id."""
        calc, _, _ = _make_calculator(cornerstones=[])
        with pytest.raises(ValueError, match="Cornerstone not found"):
            await calc.check_one("nonexistent-id")

    @pytest.mark.asyncio
    async def test_report_contains_sub_scores(self):
        """DriftReport must contain semantic_sim, length_ratio, phrase_preservation."""
        text = "Eva is an AI companion who loves building systems."
        cs = _make_cornerstone(content=text, golden_copy=text)
        calc, _, _ = _make_calculator(cornerstones=[cs])

        report = await calc.check_one(cs.id)

        assert 0.0 <= report.semantic_sim <= 1.0
        assert 0.0 <= report.length_ratio <= 1.0
        assert 0.0 <= report.phrase_preservation <= 1.0

    @pytest.mark.asyncio
    async def test_drift_detected_when_content_empty(self):
        """Empty current content against golden should produce high drift."""
        golden = "Eva is an AI companion who loves building systems and is fiercely loyal."
        cs = _make_cornerstone(content="", golden_copy=golden)
        # Use embedding value that makes embeddings very different
        calc, storage, _ = _make_calculator(cornerstones=[cs], embedding_value=0.0)

        report = await calc.check_one(cs.id)

        # Length component: empty vs golden → score=0 → drift+=0.3
        # Phrase component: no phrases in empty string → drift+=0.3
        # Total >= 0.3 in most cases
        assert report.drift_score >= 0.3

    @pytest.mark.asyncio
    async def test_restore_called_on_restore_action(self):
        """Cornerstone should be restored when action is 'restore'."""
        golden = "Eva is a loyal AI companion."
        current = ""  # Empty content → high drift
        cs = _make_cornerstone(content=current, golden_copy=golden)

        class ForceRestoreConfig(MockConfig):
            drift_warn_threshold: float = 0.0   # everything triggers at least warn
            drift_check_threshold: float = 0.0   # everything triggers restore
            drift_restore_threshold: float = 0.0  # everything triggers alert

        calc = DriftCalculator(
            MockStorageAdapter([cs]),
            ForceRestoreConfig(),
            MockLLMProvider(embedding_value=0.0),
        )

        report = await calc.check_one(cs.id)

        assert report.action in ("restore", "alert")
        # Storage.restore_cornerstone should have been called
        assert cs.id in calc.storage.restore_calls  # type: ignore[attr-defined]

    @pytest.mark.asyncio
    async def test_embedding_failure_degrades_gracefully(self):
        """If embedding call fails, drift should fall back to 0 similarity for that component."""

        class FailingLLM(MockLLMProvider):
            async def embed(self, request=None, **kwargs):
                raise RuntimeError("API down")

        text = "Eva is an AI companion."
        cs = _make_cornerstone(content=text, golden_copy=text)
        storage = MockStorageAdapter([cs])
        calc = DriftCalculator(storage, MockConfig(), FailingLLM())

        report = await calc.check_one(cs.id)

        # semantic_sim should be 0.0 (failed)
        assert report.semantic_sim == 0.0
        # But report should still be generated
        assert isinstance(report, DriftReport)


class TestDriftCalculatorCheckAll:
    @pytest.mark.asyncio
    async def test_check_all_returns_one_report_per_cornerstone(self):
        cs1 = _make_cornerstone("cs-0001", "label-1", "Content one.", "Content one.")
        cs2 = _make_cornerstone("cs-0002", "label-2", "Content two.", "Content two.")
        calc, _, _ = _make_calculator(cornerstones=[cs1, cs2])

        reports = await calc.check_all(owner_id="default")

        assert len(reports) == 2
        ids = {r.cornerstone_id for r in reports}
        assert ids == {"cs-0001", "cs-0002"}

    @pytest.mark.asyncio
    async def test_check_all_empty_returns_empty_list(self):
        calc, _, _ = _make_calculator(cornerstones=[])

        reports = await calc.check_all(owner_id="default")

        assert reports == []

    @pytest.mark.asyncio
    async def test_check_all_filters_by_owner(self):
        cs_alice = _make_cornerstone("cs-alice", owner_id="alice")
        cs_bob = _make_cornerstone("cs-bob", owner_id="bob")
        calc, storage, _ = _make_calculator(cornerstones=[cs_alice, cs_bob])

        # check_all filters by owner_id
        reports = await calc.check_all(owner_id="alice")

        # Storage returns alice cornerstones only for owner_id="alice"
        assert all(r.cornerstone_id == "cs-alice" for r in reports)

    @pytest.mark.asyncio
    async def test_check_all_continues_on_individual_failure(self):
        """check_all should not abort entirely if one check_one fails."""
        cs1 = _make_cornerstone("cs-0001", "good-cs", "Content one.", "Content one.")
        cs2 = _make_cornerstone("cs-0002", "bad-cs", "Content two.", "Content two.")

        class BrokenCalc(DriftCalculator):
            async def check_one(self, cs_id):
                if cs_id == "cs-0002":
                    raise RuntimeError("Simulated failure")
                return await super().check_one(cs_id)

        storage = MockStorageAdapter([cs1, cs2])
        calc = BrokenCalc(storage, MockConfig(), MockLLMProvider())

        reports = await calc.check_all(owner_id="default")

        # Only cs1 should produce a report; cs2 error is swallowed
        assert len(reports) == 1
        assert reports[0].cornerstone_id == "cs-0001"


class TestDriftCalculatorRestore:
    @pytest.mark.asyncio
    async def test_restore_resets_content(self):
        golden = "Eva is a loyal AI companion."
        cs = _make_cornerstone(content="drifted content", golden_copy=golden)
        calc, storage, _ = _make_calculator(cornerstones=[cs])

        restored = await calc.restore(cs.id)

        assert restored.content == golden
        assert cs.id in storage.restore_calls

    @pytest.mark.asyncio
    async def test_restore_logs_drift(self):
        golden = "Eva is a loyal AI companion."
        cs = _make_cornerstone(content="drifted content", golden_copy=golden)
        calc, storage, _ = _make_calculator(cornerstones=[cs])

        await calc.restore(cs.id)

        # A log entry with action_taken='restored' should exist
        restore_logs = [
            e for e in storage.log_drift_calls if e.get("action_taken") == "restored"
        ]
        assert len(restore_logs) >= 1
        log = restore_logs[0]
        assert log["cornerstone_id"] == cs.id
        assert log["drift_score"] == 0.0


class TestDriftCalculatorKeyPhrases:
    @pytest.mark.asyncio
    async def test_extract_key_phrases_returns_list(self):
        calc, _, _ = _make_calculator(cornerstones=[])
        text = "Eva is an AI companion. She loves building systems. She is fiercely loyal."
        phrases = await calc.extract_key_phrases(text)

        assert isinstance(phrases, list)
        assert len(phrases) > 0

    @pytest.mark.asyncio
    async def test_extract_key_phrases_max_10(self):
        calc, _, _ = _make_calculator(cornerstones=[])
        long_text = " ".join([f"Phrase number {i} is important." for i in range(20)])
        phrases = await calc.extract_key_phrases(long_text)

        assert len(phrases) <= 10

    @pytest.mark.asyncio
    async def test_extract_key_phrases_fallback_on_llm_failure(self):
        """When LLM fails, regex fallback should still return phrases."""

        class FailingLLM(MockLLMProvider):
            async def complete(self, request=None, **kwargs):
                raise RuntimeError("LLM unavailable")

        calc = DriftCalculator(
            MockStorageAdapter([]),
            MockConfig(),
            FailingLLM(),
        )
        text = (
            "Eva is an AI companion who loves building systems. "
            "She is fiercely loyal and protective. "
            "Eva was built on curiosity and courage."
        )
        phrases = await calc.extract_key_phrases(text)

        # Fallback should still produce something useful
        assert isinstance(phrases, list)


class TestDriftCalculatorWeights:
    def test_weights_sum_to_one(self):
        """similarity_weight + length_weight + phrase_weight must equal 1.0."""
        calc = DriftCalculator(
            MockStorageAdapter([]), MockConfig(), MockLLMProvider()
        )
        total = calc.similarity_weight + calc.length_weight + calc.phrase_weight
        assert total == pytest.approx(1.0)

    def test_default_weights_match_architecture_spec(self):
        calc = DriftCalculator(
            MockStorageAdapter([]), MockConfig(), MockLLMProvider()
        )
        assert calc.similarity_weight == pytest.approx(0.4)
        assert calc.length_weight == pytest.approx(0.3)
        assert calc.phrase_weight == pytest.approx(0.3)


class TestDriftReport:
    def test_drift_report_fields(self):
        report = DriftReport(
            cornerstone_id="cs-0001",
            label="test",
            drift_score=0.15,
            semantic_sim=0.9,
            length_ratio=0.95,
            phrase_preservation=0.8,
            action="warn",
        )
        assert report.cornerstone_id == "cs-0001"
        assert report.drift_score == 0.15
        assert report.action == "warn"
        assert report.checked_at  # auto-populated

    def test_drift_score_range_semantics(self):
        """drift_score 0.0 = identical, 1.0 = completely different."""
        # These are documentation tests, not computation tests
        min_report = DriftReport("id", "label", 0.0, 1.0, 1.0, 1.0, "none")
        max_report = DriftReport("id", "label", 1.0, 0.0, 0.0, 0.0, "alert")

        assert min_report.drift_score == 0.0
        assert max_report.drift_score == 1.0
