"""Tests for VectorSearch Protocol and SqliteVecSearch implementation.

PgVectorSearch tests use mocks since they require a running CockroachDB.
SqliteVecSearch tests use real sqlite-vec extension.
"""
import sqlite3
from unittest.mock import MagicMock, patch

import pytest
import sqlite_vec

from greymatter_core.vector import (
    VectorSearch,
    SqliteVecSearch,
    PgVectorSearch,
    serialize_f32,
)


# --- Fixtures ---


@pytest.fixture
def vec_conn():
    """SQLite connection with vec extension and a test table."""
    conn = sqlite3.connect(":memory:")
    conn.enable_load_extension(True)
    sqlite_vec.load(conn)
    conn.enable_load_extension(False)
    conn.execute("""
        CREATE VIRTUAL TABLE vec_tests USING vec0(
            test_id TEXT PRIMARY KEY,
            embedding float[4]
        )
    """)
    conn.commit()
    return conn


@pytest.fixture
def sqlite_vs(vec_conn):
    """SqliteVecSearch instance with 4-dimensional test table."""
    return SqliteVecSearch(vec_conn)


def _vec4(x: float, y: float, z: float, w: float) -> list[float]:
    """Helper to create a 4-dimensional vector."""
    return [x, y, z, w]


# --- serialize_f32 ---


class TestSerializeF32:
    def test_roundtrip(self):
        vec = [1.0, 2.0, 3.0]
        blob = serialize_f32(vec)
        assert len(blob) == 12  # 3 * 4 bytes

    def test_empty_vector(self):
        blob = serialize_f32([])
        assert len(blob) == 0


# --- Protocol compliance ---


class TestProtocolCompliance:
    def test_sqlite_vec_is_vector_search(self, sqlite_vs):
        assert isinstance(sqlite_vs, VectorSearch)

    def test_pgvector_is_vector_search(self):
        pvs = PgVectorSearch(lambda: MagicMock())
        assert isinstance(pvs, VectorSearch)


# --- SqliteVecSearch ---


class TestSqliteVecSearch:
    def test_upsert_and_count(self, sqlite_vs):
        sqlite_vs.upsert("vec_tests", "id-1", _vec4(1, 0, 0, 0))
        assert sqlite_vs.count("vec_tests") == 1

    def test_upsert_replaces(self, sqlite_vs):
        sqlite_vs.upsert("vec_tests", "id-1", _vec4(1, 0, 0, 0))
        sqlite_vs.upsert("vec_tests", "id-1", _vec4(0, 1, 0, 0))
        assert sqlite_vs.count("vec_tests") == 1

    def test_search_returns_results(self, sqlite_vs):
        sqlite_vs.upsert("vec_tests", "a", _vec4(1, 0, 0, 0))
        sqlite_vs.upsert("vec_tests", "b", _vec4(0, 1, 0, 0))
        sqlite_vs.upsert("vec_tests", "c", _vec4(0.9, 0.1, 0, 0))

        results = sqlite_vs.search("vec_tests", _vec4(1, 0, 0, 0), limit=3)
        assert len(results) == 3
        # Closest to [1,0,0,0] should be "a" (exact match)
        assert results[0]["id"] == "a"
        assert results[0]["distance"] < 0.01

    def test_search_respects_limit(self, sqlite_vs):
        for i in range(10):
            sqlite_vs.upsert("vec_tests", f"id-{i}", _vec4(float(i), 0, 0, 0))

        results = sqlite_vs.search("vec_tests", _vec4(5, 0, 0, 0), limit=3)
        assert len(results) == 3

    def test_search_empty_table(self, sqlite_vs):
        results = sqlite_vs.search("vec_tests", _vec4(1, 0, 0, 0), limit=5)
        assert results == []

    def test_delete(self, sqlite_vs):
        sqlite_vs.upsert("vec_tests", "id-1", _vec4(1, 0, 0, 0))
        assert sqlite_vs.count("vec_tests") == 1
        sqlite_vs.delete("vec_tests", "id-1")
        assert sqlite_vs.count("vec_tests") == 0

    def test_delete_nonexistent_is_noop(self, sqlite_vs):
        sqlite_vs.delete("vec_tests", "nonexistent")
        assert sqlite_vs.count("vec_tests") == 0

    def test_pk_col_derivation(self):
        assert SqliteVecSearch._get_pk_col("vec_observations") == "observation_id"
        assert SqliteVecSearch._get_pk_col("vec_documents") == "document_id"
        assert SqliteVecSearch._get_pk_col("vec_knowledge") == "knowledge_id"
        assert SqliteVecSearch._get_pk_col("vec_code_symbols") == "code_symbol_id"
        assert SqliteVecSearch._get_pk_col("vec_chunks") == "chunk_id"
        # edge: table without vec_ prefix
        assert SqliteVecSearch._get_pk_col("tests") == "test_id"


# --- PgVectorSearch (mocked) ---


class TestPgVectorSearch:
    def test_upsert_calls_execute(self):
        mock_conn = MagicMock()
        mock_cursor = MagicMock()
        mock_conn.cursor.return_value.__enter__ = lambda s: mock_cursor
        mock_conn.cursor.return_value.__exit__ = MagicMock(return_value=False)

        pvs = PgVectorSearch(lambda: mock_conn)
        pvs.upsert("vec_obs", "id-1", [1.0, 2.0, 3.0])

        mock_cursor.execute.assert_called_once()
        call_args = mock_cursor.execute.call_args
        assert "INSERT INTO vec_obs" in call_args[0][0]
        assert "ON CONFLICT" in call_args[0][0]
        mock_conn.commit.assert_called_once()

    def test_search_calls_execute(self):
        mock_conn = MagicMock()
        mock_cursor = MagicMock()
        mock_cursor.fetchall.return_value = [("id-1", 0.05), ("id-2", 0.12)]
        mock_conn.cursor.return_value.__enter__ = lambda s: mock_cursor
        mock_conn.cursor.return_value.__exit__ = MagicMock(return_value=False)

        pvs = PgVectorSearch(lambda: mock_conn)
        results = pvs.search("vec_obs", [1.0, 2.0, 3.0], limit=2)

        assert len(results) == 2
        assert results[0]["id"] == "id-1"
        assert results[0]["distance"] == 0.05
        mock_cursor.execute.assert_called_once()
        assert "<=>" in mock_cursor.execute.call_args[0][0]

    def test_delete_calls_execute(self):
        mock_conn = MagicMock()
        mock_cursor = MagicMock()
        mock_conn.cursor.return_value.__enter__ = lambda s: mock_cursor
        mock_conn.cursor.return_value.__exit__ = MagicMock(return_value=False)

        pvs = PgVectorSearch(lambda: mock_conn)
        pvs.delete("vec_obs", "id-1")

        mock_cursor.execute.assert_called_once()
        assert "DELETE FROM vec_obs" in mock_cursor.execute.call_args[0][0]
        mock_conn.commit.assert_called_once()

    def test_count_calls_execute(self):
        mock_conn = MagicMock()
        mock_cursor = MagicMock()
        mock_cursor.fetchone.return_value = (42,)
        mock_conn.cursor.return_value.__enter__ = lambda s: mock_cursor
        mock_conn.cursor.return_value.__exit__ = MagicMock(return_value=False)

        pvs = PgVectorSearch(lambda: mock_conn)
        assert pvs.count("vec_obs") == 42

    def test_upsert_rollback_on_error(self):
        mock_conn = MagicMock()
        mock_cursor = MagicMock()
        mock_cursor.execute.side_effect = Exception("connection lost")
        mock_conn.cursor.return_value.__enter__ = lambda s: mock_cursor
        mock_conn.cursor.return_value.__exit__ = MagicMock(return_value=False)

        pvs = PgVectorSearch(lambda: mock_conn)
        with pytest.raises(Exception, match="connection lost"):
            pvs.upsert("vec_obs", "id-1", [1.0, 2.0])

        mock_conn.rollback.assert_called_once()
