"""Integration tests for hybrid search.

These tests require a running PostgreSQL instance with AGE and pgvector.
"""

import pytest

from piggie.connection import drop_graph, ensure_graph

pytestmark = pytest.mark.integration

HYBRID_GRAPH = "piggie_hybrid_test"


@pytest.fixture
def hybrid_setup(db):
    """Set up graph + vector data for hybrid tests."""
    # Ensure clean transaction state
    db.rollback()

    # Clean up
    try:
        drop_graph(db.connection, HYBRID_GRAPH)
    except Exception:
        db.rollback()
    try:
        db.sql("DROP TABLE IF EXISTS hybrid_papers")
        db.commit()
    except Exception:
        db.rollback()

    # Create graph with papers and citations
    ensure_graph(db.connection, HYBRID_GRAPH)
    db.set_graph(HYBRID_GRAPH)

    db.cypher("CREATE (:Paper {id: 1, title: 'Graph Neural Networks'})")
    db.cypher("CREATE (:Paper {id: 2, title: 'Attention Is All You Need'})")
    db.cypher("CREATE (:Paper {id: 3, title: 'BERT Pre-training'})")
    db.commit()

    db.cypher(
        "MATCH (a:Paper {id: 1}), (b:Paper {id: 2}) "
        "CREATE (a)-[:CITES]->(b)"
    )
    db.cypher(
        "MATCH (a:Paper {id: 1}), (b:Paper {id: 3}) "
        "CREATE (a)-[:CITES]->(b)"
    )
    db.commit()

    # Create vector table
    db.sql(
        "CREATE TABLE hybrid_papers ("
        "  id integer PRIMARY KEY,"
        "  title text,"
        "  embedding vector(3)"
        ")"
    )
    db.commit()
    db.sql(
        "INSERT INTO hybrid_papers (id, title, embedding) VALUES "
        "(1, 'Graph Neural Networks', '[0.9, 0.1, 0.0]'),"
        "(2, 'Attention Is All You Need', '[0.1, 0.9, 0.0]'),"
        "(3, 'BERT Pre-training', '[0.1, 0.8, 0.2]')"
    )
    db.commit()

    yield db

    # Cleanup
    try:
        drop_graph(db.connection, HYBRID_GRAPH)
    except Exception:
        pass
    db.sql("DROP TABLE IF EXISTS hybrid_papers")
    db.commit()


def test_graph_then_vector(hybrid_setup):
    db = hybrid_setup
    result = db.hybrid_search(
        cypher="MATCH (:Paper {id: 1})-[:CITES]->(p:Paper) RETURN p.id AS id",
        vector_table="hybrid_papers",
        vector_column="embedding",
        query_vector=[0.1, 0.9, 0.1],
        k=10,
        metric="cosine",
        join_column="id",
        mode="graph_then_vector",
    )
    assert len(result) > 0
    df = result.to_df()
    assert "_distance" in df.columns
    # Results should be from the papers cited by paper 1 (papers 2 and 3)
    result_ids = set(df["id"].tolist())
    assert result_ids.issubset({2, 3})


def test_vector_then_graph(hybrid_setup):
    db = hybrid_setup
    # vector_then_graph: find similar vectors, then return vector results
    # (without a Cypher expansion step, just vector results)
    result = db.hybrid_search(
        vector_table="hybrid_papers",
        vector_column="embedding",
        query_vector=[0.1, 0.9, 0.1],
        k=2,
        metric="cosine",
        join_column="id",
        mode="vector_then_graph",
    )
    # Should return the 2 most similar papers
    assert len(result) == 2
