"""Tests for MemoryConsolidator — two-phase conflict resolution."""

import json
import sqlite3
import time
from unittest.mock import AsyncMock

import pytest

from neo_cortex.dedup import DedupStore
from neo_cortex.graph import ConceptGraph
from neo_cortex.memory_index import MemoryIndex
from neo_cortex.merge_consolidator import MemoryConsolidator, _UnionFind
from neo_cortex.models import Activity, MemoryRecord, StructuredFields
from neo_cortex.store import MemoryStore


FAKE_EMBEDDING = [0.1] * 1024
FAKE_EMBEDDING_2 = [0.2] * 1024


def _insert_memory(store, index, mid, topic="test", project="neo-cortex",
                    content="Test content", concepts=None, embedding=None):
    """Insert a memory into both store and index."""
    ts = time.time()
    record = MemoryRecord(
        id=mid, session_id="s1", timestamp=ts,
        question=topic, answer_preview=content[:500],
        document=content, project=project, topic=topic,
        activity=Activity.FEATURE, energy=1.0, source="distiller",
    )
    structured = StructuredFields(
        title=topic, summary=content[:300],
        facts=[content], concepts=concepts or [],
    )
    store.insert(record, embedding or FAKE_EMBEDDING)
    index.insert(record, structured)
    return record


class TestUnionFind:
    def test_basic_union(self):
        uf = _UnionFind()
        uf.union("a", "b")
        assert uf.find("a") == uf.find("b")

    def test_separate_groups(self):
        uf = _UnionFind()
        uf.union("a", "b")
        uf.union("c", "d")
        assert uf.find("a") == uf.find("b")
        assert uf.find("c") == uf.find("d")
        assert uf.find("a") != uf.find("c")

    def test_transitive_union(self):
        uf = _UnionFind()
        uf.union("a", "b")
        uf.union("b", "c")
        assert uf.find("a") == uf.find("c")

    def test_groups(self):
        uf = _UnionFind()
        uf.union("a", "b")
        uf.union("b", "c")
        uf.union("x", "y")
        groups = uf.groups()
        assert len(groups) == 2
        sizes = sorted(len(v) for v in groups.values())
        assert sizes == [2, 3]

    def test_single_element(self):
        uf = _UnionFind()
        uf.find("solo")
        groups = uf.groups()
        assert len(groups) == 1
        assert len(groups[list(groups.keys())[0]]) == 1


class TestPendingMergesCRUD:
    def test_add_and_get(self, tmp_path):
        index = MemoryIndex(str(tmp_path / "idx.db"))
        index.add_pending_merge("mem_a", "mem_b", 0.75)
        pairs = index.get_pending_merges()
        assert len(pairs) == 1
        id_a, id_b, sim, row_id = pairs[0]
        assert {id_a, id_b} == {"mem_a", "mem_b"}
        assert abs(sim - 0.75) < 0.01

    def test_sorted_pair_dedup(self, tmp_path):
        """Same pair in different order → only 1 row."""
        index = MemoryIndex(str(tmp_path / "idx.db"))
        index.add_pending_merge("mem_b", "mem_a", 0.75)
        index.add_pending_merge("mem_a", "mem_b", 0.80)
        pairs = index.get_pending_merges()
        assert len(pairs) == 1

    def test_mark_processed(self, tmp_path):
        index = MemoryIndex(str(tmp_path / "idx.db"))
        index.add_pending_merge("a", "b", 0.7)
        pairs = index.get_pending_merges()
        assert len(pairs) == 1
        index.mark_merges_processed([pairs[0][3]])
        assert len(index.get_pending_merges()) == 0

    def test_clear_pending_merges(self, tmp_path):
        index = MemoryIndex(str(tmp_path / "idx.db"))
        index.add_pending_merge("a", "b", 0.7)
        index.add_pending_merge("c", "d", 0.8)
        n = index.clear_pending_merges()
        assert n == 2
        assert len(index.get_pending_merges()) == 0

    def test_empty_get(self, tmp_path):
        index = MemoryIndex(str(tmp_path / "idx.db"))
        assert index.get_pending_merges() == []

    def test_mark_empty_list(self, tmp_path):
        """mark_merges_processed with empty list → no crash."""
        index = MemoryIndex(str(tmp_path / "idx.db"))
        index.mark_merges_processed([])  # no crash


class TestConsolidateNoWork:
    @pytest.mark.asyncio
    async def test_no_pending_returns_zero(self, tmp_path):
        store = MemoryStore(str(tmp_path / "chroma"), collection_name="test_con_noop")
        index = MemoryIndex(str(tmp_path / "idx.db"))
        embedder = AsyncMock()
        llm = AsyncMock()
        c = MemoryConsolidator(store, index, embedder, llm)
        assert await c.consolidate() == 0


class TestConsolidateMerge:
    @pytest.mark.asyncio
    async def test_merge_reduces_memories(self, tmp_path):
        """Two similar memories → LLM says merge → 1 merged memory."""
        store = MemoryStore(str(tmp_path / "chroma"), collection_name="test_con_merge")
        index = MemoryIndex(str(tmp_path / "idx.db"))
        embedder = AsyncMock()
        embedder.embed_passage = AsyncMock(return_value=FAKE_EMBEDDING)
        llm = AsyncMock()
        llm.call = AsyncMock(return_value=json.dumps({
            "action": "merge",
            "merged_text": "Combined: encoding + BOM fix using UTF8Encoding(false)",
            "merged_title": "Encoding BOM fix",
            "reason": "same topic",
        }))

        _insert_memory(store, index, "kn_a", topic="Encoding fix",
                       content="Use UTF8Encoding(false) to avoid BOM", concepts=["encoding"])
        _insert_memory(store, index, "kn_b", topic="BOM issue",
                       content="BOM header corrupts JSON stdin", concepts=["encoding", "bom"],
                       embedding=FAKE_EMBEDDING_2)
        assert store.count() == 2
        assert index.count() == 2

        index.add_pending_merge("kn_a", "kn_b", 0.72)

        c = MemoryConsolidator(store, index, embedder, llm)
        merged = await c.consolidate()

        assert merged == 1
        assert store.count() == 1
        assert index.count() == 1
        # Verify merged record
        record = index.get_by_ids(["cn_%d" % (int(time.time()) % 100000)])[0] if index.count() == 1 else None
        # Can't predict exact ID, but count is right
        assert len(index.get_pending_merges()) == 0  # all processed

    @pytest.mark.asyncio
    async def test_keep_separate(self, tmp_path):
        """Two memories → LLM says keep_separate → both remain."""
        store = MemoryStore(str(tmp_path / "chroma"), collection_name="test_con_keep")
        index = MemoryIndex(str(tmp_path / "idx.db"))
        embedder = AsyncMock()
        llm = AsyncMock()
        llm.call = AsyncMock(return_value=json.dumps({
            "action": "keep_separate",
            "reason": "different topics",
        }))

        _insert_memory(store, index, "kn_x", topic="Encoding fix", content="BOM fix")
        _insert_memory(store, index, "kn_y", topic="Deploy fix", content="Systemd",
                       embedding=FAKE_EMBEDDING_2)

        index.add_pending_merge("kn_x", "kn_y", 0.68)

        c = MemoryConsolidator(store, index, embedder, llm)
        merged = await c.consolidate()

        assert merged == 0
        assert store.count() == 2
        assert index.count() == 2

    @pytest.mark.asyncio
    async def test_llm_error_keeps_all(self, tmp_path):
        """LLM returns garbage → no merge, no crash."""
        store = MemoryStore(str(tmp_path / "chroma"), collection_name="test_con_err")
        index = MemoryIndex(str(tmp_path / "idx.db"))
        embedder = AsyncMock()
        llm = AsyncMock()
        llm.call = AsyncMock(return_value="not json at all")

        _insert_memory(store, index, "kn_e1", topic="A", content="content a")
        _insert_memory(store, index, "kn_e2", topic="B", content="content b",
                       embedding=FAKE_EMBEDDING_2)
        index.add_pending_merge("kn_e1", "kn_e2", 0.7)

        c = MemoryConsolidator(store, index, embedder, llm)
        merged = await c.consolidate()

        assert merged == 0
        assert store.count() == 2

    @pytest.mark.asyncio
    async def test_merge_with_graph(self, tmp_path):
        """Merge correctly removes old concepts and adds new ones to graph."""
        store = MemoryStore(str(tmp_path / "chroma"), collection_name="test_con_graph")
        index = MemoryIndex(str(tmp_path / "idx.db"))
        graph = ConceptGraph(index.conn)
        embedder = AsyncMock()
        embedder.embed_passage = AsyncMock(return_value=FAKE_EMBEDDING)
        llm = AsyncMock()
        llm.call = AsyncMock(return_value=json.dumps({
            "action": "merge",
            "merged_text": "Merged encoding content",
            "merged_title": "Encoding merged",
            "reason": "same",
        }))

        _insert_memory(store, index, "kn_g1", topic="Enc1", content="c1", concepts=["encoding"])
        _insert_memory(store, index, "kn_g2", topic="Enc2", content="c2", concepts=["encoding", "bom"],
                       embedding=FAKE_EMBEDDING_2)
        graph.add_memory("kn_g1", ["encoding"])
        graph.add_memory("kn_g2", ["encoding", "bom"])

        index.add_pending_merge("kn_g1", "kn_g2", 0.75)

        c = MemoryConsolidator(store, index, embedder, llm, graph=graph)
        merged = await c.consolidate()

        assert merged == 1
        # Old memories removed from graph
        assert "kn_g1" not in graph.memories_for("encoding")
        assert "kn_g2" not in graph.memories_for("encoding")

    @pytest.mark.asyncio
    async def test_already_deleted_memories_skip(self, tmp_path):
        """If a memory was already deleted (e.g., by previous group), skip group."""
        store = MemoryStore(str(tmp_path / "chroma"), collection_name="test_con_del")
        index = MemoryIndex(str(tmp_path / "idx.db"))
        embedder = AsyncMock()
        llm = AsyncMock()

        # Only insert one of the two memories
        _insert_memory(store, index, "kn_exists", topic="Exists", content="I exist")
        # kn_gone is NOT inserted — simulates already-deleted

        index.add_pending_merge("kn_exists", "kn_gone", 0.7)

        c = MemoryConsolidator(store, index, embedder, llm)
        merged = await c.consolidate()

        assert merged == 0  # skipped — <2 records found
        assert store.count() == 1  # untouched
        llm.call.assert_not_called()


class TestBuildGroups:
    def test_multiple_groups(self, tmp_path):
        store = MemoryStore(str(tmp_path / "chroma"), collection_name="test_groups")
        index = MemoryIndex(str(tmp_path / "idx.db"))
        embedder = AsyncMock()
        llm = AsyncMock()
        c = MemoryConsolidator(store, index, embedder, llm)

        pairs = [
            ("a", "b", 0.7, 1),
            ("b", "c", 0.8, 2),
            ("x", "y", 0.9, 3),
        ]
        groups = c._build_groups(pairs)
        assert len(groups) == 2
        sizes = sorted(len(g) for g in groups)
        assert sizes == [2, 3]
