"""Tests for storage and migration fixes: connection tracking, store race guard, orphaned vectors.

TDD tests written BEFORE implementation to verify:
- Fix 1: Thread-local connection leak (all connections tracked and closed)
- Fix 2: Race in _store_with_embedding (id=0 vector guard)
- Fix 3: Migration orphaned vectors cleanup
"""

from __future__ import annotations

import sqlite3
import struct
import threading
from collections.abc import Generator
from pathlib import Path
from unittest.mock import patch

import pytest

from claude_engram.migration import migrate_schema
from claude_engram.storage import MemoryStorage


@pytest.fixture()
def storage(tmp_path: Path) -> Generator[MemoryStorage]:
    """Create a MemoryStorage with a temp database."""
    db = tmp_path / "test.db"
    s = MemoryStorage(db_path=db)
    yield s
    s.close()


# ============================================================================
# Fix 1: Thread-local connection leak — all connections must be tracked/closed
# ============================================================================


class TestConnectionTracking:
    """Fix 1: MemoryStorage must track all thread-local connections and close them all."""

    def test_storage_has_connection_lock(self, storage: MemoryStorage) -> None:
        """Storage should have a threading.Lock for connection tracking."""
        assert hasattr(storage, "_conn_lock")
        assert isinstance(storage._conn_lock, type(threading.Lock()))

    def test_storage_has_connection_set(self, storage: MemoryStorage) -> None:
        """Storage should have a set to track all connections."""
        assert hasattr(storage, "_all_conns")
        assert isinstance(storage._all_conns, set)

    def test_get_conn_registers_connection(self, storage: MemoryStorage) -> None:
        """_get_conn() should add the connection to the tracking set."""
        conn = storage._get_conn()
        assert conn in storage._all_conns

    def test_worker_thread_connection_is_tracked(self, storage: MemoryStorage) -> None:
        """A connection created in a worker thread should be tracked."""
        worker_conn = None
        barrier = threading.Barrier(2)

        def worker():
            nonlocal worker_conn
            worker_conn = storage._get_conn()
            barrier.wait()

        t = threading.Thread(target=worker)
        t.start()
        barrier.wait()
        t.join()

        assert worker_conn is not None
        assert worker_conn in storage._all_conns

    def test_close_closes_all_tracked_connections(self, storage: MemoryStorage) -> None:
        """close() should close connections from all threads, not just the calling thread."""
        worker_conns = []
        barrier = threading.Barrier(4)  # 3 workers + main

        def worker():
            conn = storage._get_conn()
            worker_conns.append(conn)
            barrier.wait()

        threads = []
        for _ in range(3):
            t = threading.Thread(target=worker)
            t.start()
            threads.append(t)

        barrier.wait()  # Wait for all workers to create connections
        for t in threads:
            t.join()

        # Main thread also has a connection (from __init__)
        storage._get_conn()

        # All should be tracked
        assert len(storage._all_conns) >= 4  # 3 workers + main

        storage.close()

        # After close, the tracking set should be empty
        assert len(storage._all_conns) == 0

    def test_close_clears_tracking_set(self, storage: MemoryStorage) -> None:
        """close() should clear the _all_conns set after closing."""
        storage._get_conn()
        storage.close()
        assert len(storage._all_conns) == 0

    def test_main_thread_connection_tracked(self, storage: MemoryStorage) -> None:
        """The connection created during __init__ should be tracked."""
        # __init__ calls _init_db which calls _get_conn, so at least 1 connection
        assert len(storage._all_conns) >= 1


# ============================================================================
# Fix 1b: WAL checkpoint for long-running daemon (audit 2026-02-05)
# ============================================================================


class TestWalCheckpoint:
    """checkpoint_wal() should truncate WAL file to prevent unbounded growth."""

    def test_checkpoint_wal_returns_int(self, storage: MemoryStorage) -> None:
        """checkpoint_wal() should return number of pages checkpointed."""
        result = storage.checkpoint_wal()
        assert isinstance(result, int)
        assert result >= 0

    def test_checkpoint_wal_after_writes(self, storage: MemoryStorage) -> None:
        """checkpoint_wal() should work after write operations."""
        # Do some writes to create WAL entries
        storage.store(content="checkpoint test 1", category="lesson")
        storage.store(content="checkpoint test 2", category="lesson")

        # Checkpoint should succeed
        result = storage.checkpoint_wal()
        assert isinstance(result, int)

    def test_checkpoint_wal_idempotent(self, storage: MemoryStorage) -> None:
        """Multiple checkpoint calls should be safe (idempotent)."""
        storage.store(content="idempotent test", category="lesson")

        # First checkpoint
        result1 = storage.checkpoint_wal()
        # Second checkpoint (should be 0 since nothing new to checkpoint)
        result2 = storage.checkpoint_wal()

        assert isinstance(result1, int)
        assert isinstance(result2, int)
        # Second call should checkpoint 0 pages (nothing new)
        assert result2 == 0


# ============================================================================
# Fix 2: Race in _store_with_embedding (id=0 vector guard)
# ============================================================================


class TestStoreRaceGuard:
    """Fix 2: _store_with_embedding must not insert a vector with id=0."""

    def test_store_with_embedding_no_zero_id_vector(self, storage: MemoryStorage) -> None:
        """If INSERT OR IGNORE returns lastrowid=0 and re-SELECT finds nothing,
        no vector should be inserted with id=0."""
        # After a normal store, there should be no vector with id=0
        storage.store(content="race test content", category="lesson")

        conn = storage._get_conn()
        # Check no vector with id=0 exists
        try:
            row = conn.execute("SELECT COUNT(*) FROM memory_vectors WHERE id = 0").fetchone()
            assert row[0] == 0, "No vector should exist with id=0"
        except sqlite3.OperationalError:
            pass  # Table might be empty, which is fine

    def test_store_graceful_degradation_returns_dict_fields(self, storage: MemoryStorage) -> None:
        """When store succeeds normally, it returns a valid memory ID > 0."""
        mid = storage.store(content="normal store test", category="lesson")
        assert mid > 0

    def test_no_vector_id_zero_after_dedup_race(self, storage: MemoryStorage) -> None:
        """Simulate: INSERT OR IGNORE fails (dedup), re-SELECT also fails (race).
        The method should return 0 but NOT insert a vector with id=0."""
        conn = storage._get_conn()

        # First, verify the direct code path by checking database state
        # Store something normal first
        storage.store(content="baseline content", category="lesson")

        # Now check there is no id=0 vector
        row = conn.execute("SELECT COUNT(*) FROM memory_vectors WHERE id = 0").fetchone()
        assert row[0] == 0, "No vector should exist with id=0 after normal operations"

    def test_store_with_embedding_zero_id_guard(self, storage: MemoryStorage) -> None:
        """Directly test the _store_with_embedding method with a simulated race.

        When both INSERT OR IGNORE and the fallback SELECT fail to produce an id > 0,
        the method should log a warning and NOT insert a vector with id=0.
        """
        import numpy as np

        dummy_embedding = np.zeros(384, dtype=np.float32)
        conn = storage._get_conn()

        content = "race condition content"

        # We'll use a FakeConn that intercepts calls to simulate the race:
        # - First SELECT (dedup check): returns None (no existing match)
        # - INSERT OR IGNORE: returns lastrowid=0, rowcount=0 (constraint violation)
        # - Second SELECT (re-SELECT): also returns None (race lost)

        class FakeCursor:
            def __init__(self, lastrowid=0, rowcount=0, fetchone_val=None):
                self.lastrowid = lastrowid
                self.rowcount = rowcount
                self._fetchone_val = fetchone_val

            def fetchone(self):
                return self._fetchone_val

        class FakeConn:
            """Wraps the real connection, intercepting specific calls."""

            def __init__(self, real_conn):
                self._real = real_conn
                self._select_count = 0

            def execute(self, sql, params=None):
                # Intercept SELECT id FROM memories WHERE content_hash (dedup checks)
                if "SELECT id FROM memories WHERE content_hash" in sql:
                    self._select_count += 1
                    # Both the initial check AND the re-SELECT return nothing
                    return FakeCursor(fetchone_val=None)

                # Intercept the INSERT OR IGNORE
                if "INSERT OR IGNORE INTO memories" in sql:
                    return FakeCursor(lastrowid=0, rowcount=0)

                # Let everything else through
                result = self._real.execute(sql, params) if params else self._real.execute(sql)
                return result

            def commit(self):
                return self._real.commit()

            def __getattr__(self, name):
                return getattr(self._real, name)

        fake_conn = FakeConn(conn)

        with (
            patch.object(storage, "_get_conn", return_value=fake_conn),
            patch("claude_engram.storage.logger") as mock_logger,
        ):
            result = storage._store_with_embedding(
                content=content,
                category="lesson",
                metadata=None,
                proj_tag=None,
                embedding=dummy_embedding,
            )

        # The result should be 0 (no valid ID)
        assert result == 0

        # Logger warning should have been called
        mock_logger.warning.assert_called_once()

        # No vector with id=0 should exist
        row = conn.execute("SELECT COUNT(*) FROM memory_vectors WHERE id = 0").fetchone()
        assert row[0] == 0, "Vector with id=0 must not be inserted"


# ============================================================================
# Fix 3: Migration orphaned vectors cleanup
# ============================================================================


@pytest.fixture()
def v1_db_with_vectors(tmp_path) -> sqlite3.Connection:
    """Create a v1 database with memories AND a memory_vectors table with data."""
    db_path = tmp_path / "test_vectors.db"
    conn = sqlite3.connect(str(db_path))

    # Load sqlite-vec extension BEFORE creating vec0 virtual table
    conn.enable_load_extension(True)
    import sqlite_vec

    sqlite_vec.load(conn)
    conn.enable_load_extension(False)

    conn.executescript("""
        CREATE TABLE memories (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            content TEXT NOT NULL,
            category TEXT DEFAULT 'general',
            metadata TEXT DEFAULT '{}',
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL
        );
        CREATE INDEX idx_memories_category ON memories(category);
        CREATE INDEX idx_memories_created ON memories(created_at);

        CREATE VIRTUAL TABLE IF NOT EXISTS memory_vectors USING vec0(
            id INTEGER PRIMARY KEY,
            embedding FLOAT[384]
        );
    """)
    conn.commit()
    return conn


class TestMigrationOrphanedVectors:
    """Fix 3: v1→v2 migration should clean up orphaned vectors after dedup."""

    def test_orphaned_vectors_cleaned_after_dedup(
        self, v1_db_with_vectors: sqlite3.Connection
    ) -> None:
        """When duplicate memories are removed, their vectors should also be removed."""
        conn = v1_db_with_vectors
        dummy_vec = struct.pack("384f", *([0.1] * 384))

        # Insert 3 memories with same content (will be deduped)
        for _i in range(3):
            conn.execute(
                "INSERT INTO memories (content, category, created_at, updated_at) "
                "VALUES (?, ?, ?, ?)",
                ("duplicate content", "lesson", "2024-01-01T00:00:00", "2024-01-01T00:00:00"),
            )
        conn.commit()

        # Get the IDs
        rows = conn.execute("SELECT id FROM memories ORDER BY id").fetchall()
        assert len(rows) == 3

        # Insert vectors for all 3
        for row in rows:
            conn.execute(
                "INSERT INTO memory_vectors (id, embedding) VALUES (?, ?)",
                (row[0], dummy_vec),
            )
        conn.commit()

        # Verify 3 vectors exist
        vec_count = conn.execute("SELECT COUNT(*) FROM memory_vectors").fetchone()[0]
        assert vec_count == 3

        # Run migration — should dedup to 1 memory AND clean orphaned vectors
        migrate_schema(conn)

        # Only 1 memory should remain
        mem_count = conn.execute("SELECT COUNT(*) FROM memories").fetchone()[0]
        assert mem_count == 1

        # Only 1 vector should remain (orphans cleaned)
        vec_count = conn.execute("SELECT COUNT(*) FROM memory_vectors").fetchone()[0]
        assert vec_count == 1

    def test_no_orphaned_vectors_without_duplicates(
        self, v1_db_with_vectors: sqlite3.Connection
    ) -> None:
        """When there are no duplicates, vectors should not be affected."""
        conn = v1_db_with_vectors
        dummy_vec = struct.pack("384f", *([0.1] * 384))

        # Insert 2 unique memories
        conn.execute(
            "INSERT INTO memories (content, category, created_at, updated_at) VALUES (?, ?, ?, ?)",
            ("unique A", "lesson", "2024-01-01T00:00:00", "2024-01-01T00:00:00"),
        )
        conn.execute(
            "INSERT INTO memories (content, category, created_at, updated_at) VALUES (?, ?, ?, ?)",
            ("unique B", "lesson", "2024-01-02T00:00:00", "2024-01-02T00:00:00"),
        )
        conn.commit()

        rows = conn.execute("SELECT id FROM memories ORDER BY id").fetchall()
        for row in rows:
            conn.execute(
                "INSERT INTO memory_vectors (id, embedding) VALUES (?, ?)",
                (row[0], dummy_vec),
            )
        conn.commit()

        migrate_schema(conn)

        # Both memories should remain
        mem_count = conn.execute("SELECT COUNT(*) FROM memories").fetchone()[0]
        assert mem_count == 2

        # Both vectors should remain
        vec_count = conn.execute("SELECT COUNT(*) FROM memory_vectors").fetchone()[0]
        assert vec_count == 2

    def test_orphan_cleanup_skipped_when_no_vectors_table(self, tmp_path) -> None:
        """If memory_vectors table doesn't exist, migration should still succeed."""
        db_path = tmp_path / "no_vectors.db"
        conn = sqlite3.connect(str(db_path))
        conn.executescript("""
            CREATE TABLE memories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                content TEXT NOT NULL,
                category TEXT DEFAULT 'general',
                metadata TEXT DEFAULT '{}',
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            );
        """)
        conn.commit()

        # Should not raise even without memory_vectors table
        migrate_schema(conn)

        # Migration completed
        from claude_engram.migration import get_schema_version

        assert get_schema_version(conn) == 10

    def test_orphaned_vectors_cleaned_mixed_scenario(
        self, v1_db_with_vectors: sqlite3.Connection
    ) -> None:
        """Mixed: some memories have duplicates (with vectors), some don't."""
        conn = v1_db_with_vectors
        dummy_vec = struct.pack("384f", *([0.1] * 384))

        # 2 unique + 2 duplicates of "dup content"
        conn.execute(
            "INSERT INTO memories (content, category, created_at, updated_at) VALUES (?, ?, ?, ?)",
            ("unique one", "lesson", "2024-01-01T00:00:00", "2024-01-01T00:00:00"),
        )
        conn.execute(
            "INSERT INTO memories (content, category, created_at, updated_at) VALUES (?, ?, ?, ?)",
            ("dup content", "lesson", "2024-01-02T00:00:00", "2024-01-02T00:00:00"),
        )
        conn.execute(
            "INSERT INTO memories (content, category, created_at, updated_at) VALUES (?, ?, ?, ?)",
            ("dup content", "lesson", "2024-01-03T00:00:00", "2024-01-03T00:00:00"),
        )
        conn.execute(
            "INSERT INTO memories (content, category, created_at, updated_at) VALUES (?, ?, ?, ?)",
            ("unique two", "decision", "2024-01-04T00:00:00", "2024-01-04T00:00:00"),
        )
        conn.commit()

        # Insert vectors for all 4
        rows = conn.execute("SELECT id FROM memories ORDER BY id").fetchall()
        assert len(rows) == 4
        for row in rows:
            conn.execute(
                "INSERT INTO memory_vectors (id, embedding) VALUES (?, ?)",
                (row[0], dummy_vec),
            )
        conn.commit()

        migrate_schema(conn)

        # 3 memories should remain (unique one, dup content [earliest], unique two)
        mem_count = conn.execute("SELECT COUNT(*) FROM memories").fetchone()[0]
        assert mem_count == 3

        # 3 vectors should remain (orphan from deleted dup cleaned)
        vec_count = conn.execute("SELECT COUNT(*) FROM memory_vectors").fetchone()[0]
        assert vec_count == 3


# ============================================================================
# Fix 4: validate_category called on store()
# ============================================================================


# ============================================================================
# Fix P3: Duplicate cleanup — remove pre-existing duplicate memories
# ============================================================================


class TestCleanupDuplicates:
    """P3: cleanup_duplicates should remove pre-existing dupes keeping newest."""

    def test_cleanup_duplicates_removes_older_copies(self, storage: MemoryStorage) -> None:
        """Given 3 copies with same content_hash, keep newest, delete 2 older."""
        conn = storage._get_conn()
        now = "2024-01-03T00:00:00+00:00"
        old = "2024-01-01T00:00:00+00:00"
        mid = "2024-01-02T00:00:00+00:00"
        content_hash = "abc123duplicate"
        dummy_vec = struct.pack("384f", *([0.1] * 384))

        # Temporarily drop UNIQUE index to allow inserting dupes
        conn.execute("DROP INDEX IF EXISTS idx_memories_dedup")

        for ts in [old, mid, now]:
            conn.execute(
                "INSERT INTO memories (content, category, metadata, created_at, updated_at, "
                "proj_tag, ttl_tier, content_hash, token_count) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                ("dup content", "lesson", "{}", ts, ts, "", "permanent", content_hash, 3),
            )
        conn.commit()

        # Insert vectors for all 3
        rows = conn.execute("SELECT id FROM memories ORDER BY id").fetchall()
        for row in rows:
            conn.execute(
                "INSERT INTO memory_vectors (id, embedding) VALUES (?, ?)",
                (row["id"], dummy_vec),
            )
        conn.commit()

        assert conn.execute("SELECT COUNT(*) FROM memories").fetchone()[0] == 3

        removed = storage.cleanup_duplicates()
        assert removed == 2

        # Only 1 memory should remain — the newest
        remaining = conn.execute("SELECT created_at FROM memories").fetchall()
        assert len(remaining) == 1
        assert remaining[0]["created_at"] == now

        # Only 1 vector should remain
        vec_count = conn.execute("SELECT COUNT(*) FROM memory_vectors").fetchone()[0]
        assert vec_count == 1

    def test_cleanup_duplicates_no_dupes_returns_zero(self, storage: MemoryStorage) -> None:
        """When no duplicates exist, cleanup should return 0."""
        storage.store(content="unique a", category="lesson")
        storage.store(content="unique b", category="decision")
        removed = storage.cleanup_duplicates()
        assert removed == 0

    def test_cleanup_duplicates_different_proj_tag_not_removed(
        self, storage: MemoryStorage
    ) -> None:
        """Same content_hash but different proj_tag should NOT be treated as dupes."""
        storage.store(content="cross-proj content", category="lesson", proj_tag="proj:A")
        storage.store(content="cross-proj content", category="lesson", proj_tag="proj:B")
        removed = storage.cleanup_duplicates()
        assert removed == 0

    def test_cleanup_duplicates_recreates_unique_index(self, storage: MemoryStorage) -> None:
        """After cleanup, the UNIQUE dedup index should exist."""
        storage.cleanup_duplicates()
        conn = storage._get_conn()
        indexes = conn.execute("PRAGMA index_list(memories)").fetchall()
        index_names = [idx[1] for idx in indexes]
        assert "idx_memories_dedup" in index_names


class TestStoreCategoryValidation:
    """Fix 4: store() must call validate_category to reject invalid categories."""

    def test_store_rejects_invalid_category(self, storage: MemoryStorage) -> None:
        """store() should raise ValueError for an unknown category name."""
        with pytest.raises(ValueError, match="Unknown category"):
            storage.store(content="test content", category="nonexistent_category")

    def test_store_rejects_typo_category(self, storage: MemoryStorage) -> None:
        """store() should reject common typos like 'lessson' or 'decison'."""
        with pytest.raises(ValueError, match="Unknown category"):
            storage.store(content="test content", category="lessson")

    def test_store_accepts_valid_category(self, storage: MemoryStorage) -> None:
        """store() should accept all valid MemoryCategory values."""
        mid = storage.store(content="valid category test", category="lesson")
        assert mid > 0

    def test_store_accepts_general_legacy(self, storage: MemoryStorage) -> None:
        """store() should accept the legacy 'general' category."""
        mid = storage.store(content="general category test", category="general")
        assert mid > 0

    def test_store_accepts_all_valid_categories(self, storage: MemoryStorage) -> None:
        """store() should accept all 7 defined categories."""
        valid_categories = [
            "file_ownership",
            "interface",
            "query_cache",
            "lesson",
            "checkpoint",
            "session_summary",
            "decision",
        ]
        for cat in valid_categories:
            mid = storage.store(content=f"test {cat}", category=cat)
            assert mid > 0, f"Category '{cat}' should be accepted"

    def test_store_validation_happens_before_embedding(self, storage: MemoryStorage) -> None:
        """Invalid category should fail fast, before computing the embedding."""
        with (
            patch("claude_engram.storage._embed_text") as mock_embed,
            pytest.raises(ValueError, match="Unknown category"),
        ):
            storage.store(content="test content", category="bogus")
        # _embed_text should NOT have been called
        mock_embed.assert_not_called()
