"""
tests/test_migrate_mem0.py — Unit tests for the Mem0 → Cortex migration tool.

Tests cover:
    - Mem0Memory field parsing from raw API dicts
    - Field mapping (subject, predicate, object_value, created_at)
    - Cornerstone detection and sealing
    - Idempotency (re-run doesn't duplicate)
    - Dry-run mode (no writes)
    - Entity extraction heuristic
    - MigrationReport summary
    - Mem0Client pagination (mocked)
    - Full end-to-end run with mock adapter

All network calls are mocked via pytest + unittest.mock; no real HTTP or SQLite
connections are made in these tests.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from cortex.tools.migrate_mem0 import (
    DEFAULT_PREDICATE,
    MEM0_SOURCE_PREFIX,
    Mem0Client,
    Mem0Memory,
    Mem0Migrator,
    MigrationReport,
    _extract_entities,
    _map_created_at,
    _map_predicate,
    _map_subject,
)


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

def _make_raw_memory(
    mem_id: str = "abc123",
    memory: str = "Eva is an AI assistant.",
    agent_id: str = "eva",
    metadata: Optional[Dict[str, Any]] = None,
    created_at: str = "2025-01-01T00:00:00Z",
) -> Dict[str, Any]:
    """Build a minimal raw Mem0 API memory dict."""
    return {
        "id": mem_id,
        "memory": memory,
        "agent_id": agent_id,
        "user_id": None,
        "created_at": created_at,
        "updated_at": created_at,
        "metadata": metadata or {},
    }


# ---------------------------------------------------------------------------
# Minimal in-memory mock adapter
# ---------------------------------------------------------------------------


class _MockAdapter:
    """
    Minimal mock of SQLiteAdapter for migration tests.

    Records all write calls so tests can assert what was stored.
    """

    def __init__(self) -> None:
        self.claims: List[Dict[str, Any]] = []
        self.cornerstones: List[Dict[str, Any]] = []
        self.entities: List[Dict[str, Any]] = []
        self.entity_aliases: List[Dict[str, Any]] = []
        # source_session → list of claim dicts (mirrors get_session_claims)
        self._by_source_session: Dict[str, List] = {}
        self._counter = 0

    def _next_id(self) -> str:
        self._counter += 1
        return f"id-{self._counter:04d}"

    # --- Claim operations ---

    async def create_claim(
        self,
        subject: str,
        predicate: str,
        object_value: str,
        owner_id: str = "default",
        **kwargs,
    ):
        rec = {
            "id": self._next_id(),
            "subject": subject,
            "predicate": predicate,
            "object_value": object_value,
            "owner_id": owner_id,
            **kwargs,
        }
        self.claims.append(rec)
        ss = kwargs.get("source_session")
        if ss:
            self._by_source_session.setdefault(ss, []).append(rec)
        # Return a simple namespace-like object with .id
        return _ClaimRecord(**rec)

    async def get_session_claims(self, session_id: str) -> List:
        return [_ClaimRecord(**c) for c in self._by_source_session.get(session_id, [])]

    # --- Cornerstone operations ---

    async def seal_cornerstone(
        self,
        label: str,
        content: str,
        owner_id: str = "default",
        **kwargs,
    ):
        rec = {
            "id": self._next_id(),
            "label": label,
            "content": content,
            "owner_id": owner_id,
            "golden_copy": content,
            "golden_hash": hashlib.sha256(content.encode()).hexdigest(),
            "current_hash": hashlib.sha256(content.encode()).hexdigest(),
            **kwargs,
        }
        self.cornerstones.append(rec)
        return _CornerstoneRecord(**rec)

    async def get_cornerstones(self, owner_id: str = "default", status: str = "active"):
        return [_CornerstoneRecord(**c) for c in self.cornerstones if c.get("owner_id") == owner_id]

    # --- Entity operations ---

    async def create_entity(
        self,
        canonical_name: str,
        entity_type: str = "unknown",
        owner_id: str = "default",
        **kwargs,
    ):
        rec = {
            "id": self._next_id(),
            "canonical_name": canonical_name,
            "entity_type": entity_type,
            "owner_id": owner_id,
        }
        self.entities.append(rec)
        return _EntityRecord(**rec)

    async def find_entity_by_alias(self, alias: str, owner_id: str = "default"):
        # Return existing if already added
        for a in self.entity_aliases:
            if a["alias"].lower() == alias.lower() and a["owner_id"] == owner_id:
                return _EntityRecord(id=a["entity_id"], canonical_name=alias,
                                     entity_type="unknown", owner_id=owner_id)
        return None

    async def add_entity_alias(
        self,
        entity_id: str,
        alias: str,
        alias_type: str = "name",
        **kwargs,
    ):
        rec = {
            "entity_id": entity_id,
            "alias": alias,
            "alias_type": alias_type,
            "owner_id": kwargs.get("owner_id", "default"),
        }
        self.entity_aliases.append(rec)
        return rec


@dataclass
class _ClaimRecord:
    id: str
    subject: str
    predicate: str
    object_value: str
    owner_id: str = "default"
    source_session: Optional[str] = None

    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)


@dataclass
class _CornerstoneRecord:
    id: str
    label: str
    content: str
    owner_id: str = "default"

    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)


@dataclass
class _EntityRecord:
    id: str
    canonical_name: str
    entity_type: str = "unknown"
    owner_id: str = "default"


# ---------------------------------------------------------------------------
# 1. Mem0Memory parsing
# ---------------------------------------------------------------------------


class TestMem0MemoryParsing:
    def test_basic_fields(self):
        raw = _make_raw_memory(mem_id="x1", memory="Eva uses Python.")
        m = Mem0Memory.from_api_dict(raw)
        assert m.id == "x1"
        assert m.memory == "Eva uses Python."
        assert m.agent_id == "eva"

    def test_metadata_preserved(self):
        raw = _make_raw_memory(metadata={"kind": "fact", "is_cornerstone": "true"})
        m = Mem0Memory.from_api_dict(raw)
        assert m.metadata["kind"] == "fact"
        assert m.is_cornerstone is True

    def test_not_cornerstone_by_default(self):
        raw = _make_raw_memory()
        m = Mem0Memory.from_api_dict(raw)
        assert m.is_cornerstone is False

    def test_cornerstone_false_string(self):
        raw = _make_raw_memory(metadata={"is_cornerstone": "false"})
        m = Mem0Memory.from_api_dict(raw)
        assert m.is_cornerstone is False

    def test_source_session_key(self):
        raw = _make_raw_memory(mem_id="test999")
        m = Mem0Memory.from_api_dict(raw)
        assert m.source_session_key == f"{MEM0_SOURCE_PREFIX}test999"

    def test_missing_metadata_defaults_to_empty(self):
        raw = _make_raw_memory()
        raw.pop("metadata")
        m = Mem0Memory.from_api_dict(raw)
        assert m.metadata == {}


# ---------------------------------------------------------------------------
# 2. Field mapping helpers
# ---------------------------------------------------------------------------


class TestFieldMapping:
    def test_map_subject_from_metadata(self):
        raw = _make_raw_memory(metadata={"subject": "Andrew"})
        m = Mem0Memory.from_api_dict(raw)
        assert _map_subject(m, "fallback") == "Andrew"

    def test_map_subject_from_agent_id(self):
        raw = _make_raw_memory(agent_id="eva")
        m = Mem0Memory.from_api_dict(raw)
        assert _map_subject(m, "fallback") == "eva"

    def test_map_subject_fallback(self):
        raw = _make_raw_memory(agent_id="")
        m = Mem0Memory.from_api_dict(raw)
        m.agent_id = None
        m.user_id = None
        assert _map_subject(m, "fallback") == "fallback"

    def test_map_predicate_from_metadata_kind(self):
        raw = _make_raw_memory(metadata={"kind": "preference"})
        m = Mem0Memory.from_api_dict(raw)
        assert _map_predicate(m) == "preference"

    def test_map_predicate_default(self):
        raw = _make_raw_memory()
        m = Mem0Memory.from_api_dict(raw)
        assert _map_predicate(m) == DEFAULT_PREDICATE

    def test_map_created_at_utc(self):
        raw = _make_raw_memory(created_at="2025-06-01T12:00:00Z")
        m = Mem0Memory.from_api_dict(raw)
        ts = _map_created_at(m)
        assert ts is not None
        assert "2025-06-01" in ts

    def test_map_created_at_none(self):
        raw = _make_raw_memory()
        m = Mem0Memory.from_api_dict(raw)
        m.created_at = None
        assert _map_created_at(m) is None


# ---------------------------------------------------------------------------
# 3. Entity extraction heuristic
# ---------------------------------------------------------------------------


class TestEntityExtraction:
    def test_extracts_proper_nouns(self):
        text = "Andrew uses Python and Django for web development."
        entities = _extract_entities(text)
        names = [e[0] for e in entities]
        assert "Andrew" in names

    def test_skips_short_tokens(self):
        entities = _extract_entities("Hi Eva!")
        # "Hi" is 2 chars and single — should be skipped by length filter
        names = [e[0] for e in entities]
        # "Eva" is 3 chars < 4, should also be skipped
        assert not any(n in ("Hi", "Eva") for n in names)

    def test_no_duplicates(self):
        text = "Andrew is great. Andrew built this."
        entities = _extract_entities(text)
        names = [e[0] for e in entities]
        assert len(names) == len(set(names))

    def test_empty_text(self):
        assert _extract_entities("") == []


# ---------------------------------------------------------------------------
# 4. Migrator — basic import
# ---------------------------------------------------------------------------


class TestMem0MigratorBasic:
    @pytest.fixture
    def adapter(self):
        return _MockAdapter()

    @pytest.fixture
    def memories(self):
        return [
            Mem0Memory.from_api_dict(_make_raw_memory("m1", "Eva is an AI assistant.")),
            Mem0Memory.from_api_dict(_make_raw_memory("m2", "Andrew is Eva's creator.")),
        ]

    @pytest.mark.asyncio
    async def test_imports_claims(self, adapter, memories):
        migrator = Mem0Migrator(adapter, owner_id="default", dry_run=False)
        report = await migrator.run(memories)
        assert report.imported_claims == 2
        assert report.skipped_duplicates == 0
        assert len(adapter.claims) == 2

    @pytest.mark.asyncio
    async def test_claim_fields_mapped_correctly(self, adapter, memories):
        migrator = Mem0Migrator(adapter, owner_id="default", dry_run=False)
        await migrator.run(memories)
        claim = adapter.claims[0]
        assert claim["object_value"] == "Eva is an AI assistant."
        assert claim["source_session"] == f"{MEM0_SOURCE_PREFIX}m1"
        assert claim["owner_id"] == "default"

    @pytest.mark.asyncio
    async def test_report_counts(self, adapter, memories):
        migrator = Mem0Migrator(adapter, owner_id="default", dry_run=False)
        report = await migrator.run(memories)
        assert report.total_fetched == 2
        assert report.imported_claims == 2
        assert len(report.errors) == 0


# ---------------------------------------------------------------------------
# 5. Cornerstone detection and sealing
# ---------------------------------------------------------------------------


class TestCornerstoneImport:
    @pytest.fixture
    def adapter(self):
        return _MockAdapter()

    def _cornerstone_memory(self, mem_id: str = "cs1") -> Mem0Memory:
        raw = _make_raw_memory(
            mem_id=mem_id,
            memory="Eva's core identity: loyal, inventive, engineering-focused.",
            metadata={"is_cornerstone": "true", "kind": "identity", "label": "Core Identity"},
        )
        return Mem0Memory.from_api_dict(raw)

    @pytest.mark.asyncio
    async def test_cornerstone_is_sealed(self, adapter):
        mem = self._cornerstone_memory()
        migrator = Mem0Migrator(adapter, dry_run=False)
        report = await migrator.run([mem])
        assert report.imported_cornerstones == 1
        assert len(adapter.cornerstones) == 1

    @pytest.mark.asyncio
    async def test_cornerstone_label_from_metadata(self, adapter):
        mem = self._cornerstone_memory()
        migrator = Mem0Migrator(adapter, dry_run=False)
        await migrator.run([mem])
        cs = adapter.cornerstones[0]
        assert cs["label"] == "Core Identity"

    @pytest.mark.asyncio
    async def test_non_cornerstone_not_sealed(self, adapter):
        mem = Mem0Memory.from_api_dict(_make_raw_memory("nc1"))
        migrator = Mem0Migrator(adapter, dry_run=False)
        report = await migrator.run([mem])
        assert report.imported_cornerstones == 0
        assert len(adapter.cornerstones) == 0

    @pytest.mark.asyncio
    async def test_cornerstone_content_matches_memory_text(self, adapter):
        mem = self._cornerstone_memory()
        migrator = Mem0Migrator(adapter, dry_run=False)
        await migrator.run([mem])
        cs = adapter.cornerstones[0]
        assert cs["content"] == mem.memory

    @pytest.mark.asyncio
    async def test_cornerstone_seal_failure_logged_as_error(self, adapter):
        """If seal raises, it should appear in report.errors (not crash migrator)."""
        mem = self._cornerstone_memory()

        async def _boom(*a, **kw):
            raise RuntimeError("seal failed")

        adapter.seal_cornerstone = _boom
        migrator = Mem0Migrator(adapter, dry_run=False)
        report = await migrator.run([mem])
        # Claim is still imported; cornerstone failure goes to errors
        assert report.imported_claims == 1
        assert report.imported_cornerstones == 0
        assert any("Cornerstone seal failed" in e for e in report.errors)


# ---------------------------------------------------------------------------
# 6. Idempotency (re-run doesn't duplicate)
# ---------------------------------------------------------------------------


class TestIdempotency:
    @pytest.fixture
    def adapter(self):
        return _MockAdapter()

    @pytest.mark.asyncio
    async def test_skip_on_second_run(self, adapter):
        mem = Mem0Memory.from_api_dict(_make_raw_memory("dup1"))
        migrator = Mem0Migrator(adapter, dry_run=False)

        # First run
        r1 = await migrator.run([mem])
        assert r1.imported_claims == 1
        assert r1.skipped_duplicates == 0

        # Second run — same memory
        r2 = await migrator.run([mem])
        assert r2.imported_claims == 0
        assert r2.skipped_duplicates == 1

        # Only 1 claim in DB total
        assert len(adapter.claims) == 1

    @pytest.mark.asyncio
    async def test_new_memories_still_imported_on_rerun(self, adapter):
        mem1 = Mem0Memory.from_api_dict(_make_raw_memory("a1"))
        mem2 = Mem0Memory.from_api_dict(_make_raw_memory("a2"))
        migrator = Mem0Migrator(adapter, dry_run=False)

        await migrator.run([mem1])
        r2 = await migrator.run([mem1, mem2])

        assert r2.imported_claims == 1  # only mem2 is new
        assert r2.skipped_duplicates == 1
        assert len(adapter.claims) == 2


# ---------------------------------------------------------------------------
# 7. Dry-run mode
# ---------------------------------------------------------------------------


class TestDryRunMode:
    @pytest.fixture
    def adapter(self):
        return _MockAdapter()

    @pytest.mark.asyncio
    async def test_dry_run_writes_nothing(self, adapter):
        memories = [
            Mem0Memory.from_api_dict(_make_raw_memory("dr1")),
            Mem0Memory.from_api_dict(_make_raw_memory(
                "dr2",
                metadata={"is_cornerstone": "true"},
            )),
        ]
        migrator = Mem0Migrator(adapter, dry_run=True)
        report = await migrator.run(memories)

        assert len(adapter.claims) == 0
        assert len(adapter.cornerstones) == 0
        assert len(adapter.entities) == 0

    @pytest.mark.asyncio
    async def test_dry_run_counts_are_correct(self, adapter):
        memories = [
            Mem0Memory.from_api_dict(_make_raw_memory("dr3")),
            Mem0Memory.from_api_dict(_make_raw_memory("dr4", metadata={"is_cornerstone": "true"})),
        ]
        migrator = Mem0Migrator(adapter, dry_run=True)
        report = await migrator.run(memories)

        assert report.dry_run is True
        assert report.imported_claims == 2
        assert report.imported_cornerstones == 1

    @pytest.mark.asyncio
    async def test_dry_run_does_not_prevent_duplicate_on_real_run(self, adapter):
        """Dry-run should NOT mark things as imported (idempotency key not written)."""
        mem = Mem0Memory.from_api_dict(_make_raw_memory("dr5"))

        dry_migrator = Mem0Migrator(adapter, dry_run=True)
        await dry_migrator.run([mem])
        assert len(adapter.claims) == 0

        real_migrator = Mem0Migrator(adapter, dry_run=False)
        r = await real_migrator.run([mem])
        assert r.imported_claims == 1
        assert len(adapter.claims) == 1


# ---------------------------------------------------------------------------
# 8. Mem0Client (mocked HTTP)
# ---------------------------------------------------------------------------


class TestMem0Client:
    @pytest.mark.asyncio
    async def test_fetch_memories_list_response(self):
        """Mem0 API returns a raw list — parsed correctly."""
        raw = [_make_raw_memory("a"), _make_raw_memory("b")]

        mock_resp = MagicMock()
        mock_resp.raise_for_status = MagicMock()
        mock_resp.json.return_value = raw

        with patch("httpx.AsyncClient") as MockClient:
            ctx = MockClient.return_value.__aenter__.return_value
            ctx.get = AsyncMock(return_value=mock_resp)

            client = Mem0Client(api_key="test-key")
            memories = await client.fetch_memories(agent_id="eva")

        assert len(memories) == 2
        assert memories[0].id == "a"
        assert memories[1].id == "b"

    @pytest.mark.asyncio
    async def test_fetch_memories_paginated_response(self):
        """Mem0 API returns paginated envelope — last page handled correctly."""
        page1 = {"results": [_make_raw_memory("p1"), _make_raw_memory("p2")], "next": "url"}
        page2 = {"results": [_make_raw_memory("p3")], "next": None}

        mock_resp1 = MagicMock()
        mock_resp1.raise_for_status = MagicMock()
        mock_resp1.json.return_value = page1

        mock_resp2 = MagicMock()
        mock_resp2.raise_for_status = MagicMock()
        mock_resp2.json.return_value = page2

        with patch("httpx.AsyncClient") as MockClient:
            ctx = MockClient.return_value.__aenter__.return_value
            ctx.get = AsyncMock(side_effect=[mock_resp1, mock_resp2])

            client = Mem0Client(api_key="test-key")
            memories = await client.fetch_memories(agent_id="eva", page_size=2)

        assert len(memories) == 3

    @pytest.mark.asyncio
    async def test_fetch_requires_agent_or_user_id(self):
        client = Mem0Client(api_key="test-key")
        with pytest.raises(ValueError, match="Either agent_id or user_id"):
            await client.fetch_memories()

    @pytest.mark.asyncio
    async def test_auth_header_set(self):
        """Token auth header is added correctly."""
        raw = [_make_raw_memory("z")]

        mock_resp = MagicMock()
        mock_resp.raise_for_status = MagicMock()
        mock_resp.json.return_value = raw

        with patch("httpx.AsyncClient") as MockClient:
            ctx = MockClient.return_value.__aenter__.return_value
            ctx.get = AsyncMock(return_value=mock_resp)

            client = Mem0Client(api_key="secret-key", org_id="org1", project_id="proj1")
            await client.fetch_memories(agent_id="eva")

            call_kwargs = ctx.get.call_args
            headers = call_kwargs.kwargs.get("headers") or call_kwargs[1].get("headers", {})
            assert headers.get("Authorization") == "Token secret-key"
            assert headers.get("Mem0-Org-Id") == "org1"
            assert headers.get("Mem0-Project-Id") == "proj1"


# ---------------------------------------------------------------------------
# 9. MigrationReport __str__
# ---------------------------------------------------------------------------


class TestMigrationReport:
    def test_str_includes_counts(self):
        r = MigrationReport(
            total_fetched=10,
            imported_claims=8,
            skipped_duplicates=2,
            imported_cornerstones=1,
            imported_entities=3,
        )
        text = str(r)
        assert "10" in text
        assert "8" in text
        assert "2" in text
        assert "1" in text

    def test_dry_run_prefix(self):
        r = MigrationReport(dry_run=True)
        assert "[DRY-RUN]" in str(r)

    def test_no_dry_run_prefix_in_real_run(self):
        r = MigrationReport(dry_run=False)
        assert "[DRY-RUN]" not in str(r)


# ---------------------------------------------------------------------------
# 10. No-entities flag
# ---------------------------------------------------------------------------


class TestNoEntitiesFlag:
    @pytest.mark.asyncio
    async def test_no_entities_skips_entity_creation(self):
        adapter = _MockAdapter()
        mem = Mem0Memory.from_api_dict(
            _make_raw_memory("e1", memory="Andrew built Eva in Bangkok.")
        )
        migrator = Mem0Migrator(adapter, create_entities=False, dry_run=False)
        report = await migrator.run([mem])
        assert report.imported_entities == 0
        assert len(adapter.entities) == 0
