"""
Unit tests for PgStore directly.
Requires PostgreSQL + pgvector + ParadeDB.
"""
import pytest
import sqlalchemy as sa
from conftest import DB_URL
from dot_search.store import PgStore, SubSearch


@pytest.fixture
async def store():
    s = PgStore(db_url=DB_URL)
    yield s
    await s.close()


@pytest.fixture
async def table_with_data(store):
    """Creates a bare table and populates it for store-level tests."""
    async with store._engine.begin() as conn:
        await conn.execute(sa.text(
            """
            CREATE TABLE IF NOT EXISTS store_test (
                id TEXT PRIMARY KEY,
                body TEXT
            )
            """
        ))
        await conn.execute(sa.text(
            """
            INSERT INTO store_test (id, body) VALUES
            ('doc1', 'machine à laver en panne'),
            ('doc2', 'réparation de la chaudière'),
            ('doc3', 'inspection annuelle du système hydraulique')
            ON CONFLICT DO NOTHING
            """
        ))
    yield
    async with store._engine.begin() as conn:
        await conn.execute(sa.text("DROP TABLE IF EXISTS store_test"))


async def test_add_embedding_column_creates_column(store, table_with_data):
    await store.add_embedding_column("store_test", "ds_body_emb", dimension=4)
    async with store._engine.connect() as conn:
        result = await conn.execute(sa.text(
            "SELECT column_name FROM information_schema.columns "
            "WHERE table_name = 'store_test' AND column_name = 'ds_body_emb'"
        ))
        assert result.fetchone() is not None


async def test_add_embedding_column_is_idempotent(store, table_with_data):
    await store.add_embedding_column("store_test", "ds_body_emb", dimension=4)
    await store.add_embedding_column("store_test", "ds_body_emb", dimension=4)  # no error


async def test_write_embeddings(store, table_with_data):
    await store.add_embedding_column("store_test", "ds_body_emb", dimension=4)
    rows = [
        ("doc1", [1.0, 0.0, 0.0, 0.0]),
        ("doc2", [0.0, 1.0, 0.0, 0.0]),
        ("doc3", [0.0, 0.0, 1.0, 0.0]),
    ]
    await store.write_embeddings("store_test", "ds_body_emb", "id", rows)

    async with store._engine.connect() as conn:
        result = await conn.execute(sa.text(
            "SELECT id FROM store_test WHERE ds_body_emb IS NOT NULL"
        ))
        ids = {r[0] for r in result.fetchall()}
    assert ids == {"doc1", "doc2", "doc3"}


async def test_search_returns_results(store, table_with_data):
    await store.add_embedding_column("store_test", "ds_body_emb", dimension=4)
    await store.write_embeddings(
        "store_test", "ds_body_emb", "id",
        [
            ("doc1", [1.0, 0.0, 0.0, 0.0]),
            ("doc2", [0.0, 1.0, 0.0, 0.0]),
            ("doc3", [0.0, 0.0, 1.0, 0.0]),
        ],
    )
    subsearches = [
        SubSearch(
            embedding_column="ds_body_emb",
            query_embedding=[1.0, 0.0, 0.0, 0.0],  # closest to doc1
            weight=1.0,
        )
    ]
    results_per_subsearch = await store.search(
        table="store_test",
        pk_column="id",
        subsearches=subsearches,
        bm25searches=[],
        exact_searches=[],
        where=None,
        limit=3,
    )
    assert len(results_per_subsearch) == 1
    results = results_per_subsearch[0]
    assert len(results) > 0
    assert results[0].id == "doc1"


async def test_search_with_where_filters(store, table_with_data):
    await store.add_embedding_column("store_test", "ds_body_emb", dimension=4)
    await store.write_embeddings(
        "store_test", "ds_body_emb", "id",
        [
            ("doc1", [1.0, 0.0, 0.0, 0.0]),
            ("doc2", [1.0, 0.0, 0.0, 0.0]),
            ("doc3", [1.0, 0.0, 0.0, 0.0]),
        ],
    )
    subsearches = [
        SubSearch(
            embedding_column="ds_body_emb",
            query_embedding=[1.0, 0.0, 0.0, 0.0],
            weight=1.0,
        )
    ]
    results_per_subsearch = await store.search(
        table="store_test",
        pk_column="id",
        subsearches=subsearches,
        bm25searches=[],
        exact_searches=[],
        where="id = 'doc1'",
        limit=10,
    )
    results = results_per_subsearch[0]
    assert all(r.id == "doc1" for r in results)
