"""Hybrid graph + vector search."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from piggie.exceptions import PiggieError
from piggie.result import Result

if TYPE_CHECKING:
    from piggie.driver import PiggieDriver


def hybrid_search(
    *,
    driver: PiggieDriver,
    cypher: str | None = None,
    cypher_params: tuple[Any, ...] | None = None,
    vector_table: str,
    vector_column: str,
    query_vector: list[float],
    k: int = 10,
    metric: str = "cosine",
    join_column: str = "id",
    graph: str | None = None,
    mode: str = "graph_then_vector",
) -> Result:
    """Execute a hybrid graph + vector search.

    Two modes:
    - graph_then_vector: Run Cypher first, then rank results by vector similarity.
    - vector_then_graph: Find similar vectors first, then expand with Cypher.

    Args:
        driver: PiggieDriver instance.
        cypher: Cypher query for the graph part.
        cypher_params: Parameters for the Cypher query.
        vector_table: Table with vector embeddings.
        vector_column: Vector column name.
        query_vector: Query vector for similarity search.
        k: Number of results.
        metric: Distance metric.
        join_column: Column to join graph and vector results.
        graph: Graph name override.
        mode: "graph_then_vector" or "vector_then_graph".

    Returns:
        Combined Result.
    """
    if mode == "graph_then_vector":
        return _graph_then_vector(
            driver=driver,
            cypher=cypher,
            cypher_params=cypher_params,
            vector_table=vector_table,
            vector_column=vector_column,
            query_vector=query_vector,
            k=k,
            metric=metric,
            join_column=join_column,
            graph=graph,
        )
    elif mode == "vector_then_graph":
        return _vector_then_graph(
            driver=driver,
            cypher=cypher,
            cypher_params=cypher_params,
            vector_table=vector_table,
            vector_column=vector_column,
            query_vector=query_vector,
            k=k,
            metric=metric,
            join_column=join_column,
            graph=graph,
        )
    else:
        raise PiggieError(f"Unknown hybrid search mode: {mode!r}")


def _graph_then_vector(
    *,
    driver: PiggieDriver,
    cypher: str | None,
    cypher_params: tuple[Any, ...] | None,
    vector_table: str,
    vector_column: str,
    query_vector: list[float],
    k: int,
    metric: str,
    join_column: str,
    graph: str | None,
) -> Result:
    """Run Cypher to get candidates, then rank by vector similarity.

    Strategy: Execute the Cypher query to get a set of candidate IDs,
    then query the vector table for those IDs ordered by similarity.
    """
    if not cypher:
        raise PiggieError("cypher query is required for graph_then_vector mode")

    # Step 1: Execute Cypher to get candidate IDs
    graph_result = driver.cypher(cypher, params=cypher_params, graph=graph)
    if not graph_result:
        return Result([], ["id", "_distance"])

    # Extract join column values from graph results
    candidate_ids = []
    for record in graph_result.records:
        val = record.get(join_column)
        if val is not None:
            candidate_ids.append(val)

    if not candidate_ids:
        return Result([], ["id", "_distance"])

    # Step 2: Vector search filtered to candidates
    from piggie._vector.search import _METRIC_OPS
    op = _METRIC_OPS.get(metric, "<=>")
    vec_str = "[" + ",".join(str(v) for v in query_vector) + "]"

    placeholders = ",".join(["%s"] * len(candidate_ids))
    query = (
        f"SELECT *, {vector_column} {op} %s::vector AS _distance "
        f"FROM {vector_table} "
        f"WHERE {join_column} IN ({placeholders}) "
        f"ORDER BY {vector_column} {op} %s::vector "
        f"LIMIT {int(k)}"
    )
    params_list: list[Any] = [vec_str, *candidate_ids, vec_str]

    cursor = driver.connection.cursor()
    try:
        cursor.execute(query, params_list)
        if cursor.description:
            columns = [desc[0] for desc in cursor.description]
            rows = cursor.fetchall()
        else:
            columns = []
            rows = []
    except Exception as exc:
        driver.connection.rollback()
        raise PiggieError(f"Hybrid search vector step failed: {exc}") from exc

    return Result(rows, columns)


def _vector_then_graph(
    *,
    driver: PiggieDriver,
    cypher: str | None,
    cypher_params: tuple[Any, ...] | None,
    vector_table: str,
    vector_column: str,
    query_vector: list[float],
    k: int,
    metric: str,
    join_column: str,
    graph: str | None,
) -> Result:
    """Find similar vectors first, then expand with Cypher.

    Strategy: Run vector search to get top-k similar items, then use their
    IDs in a Cypher query to get graph context.
    """
    # Step 1: Vector search for top-k
    vector_result = driver.vector_search(
        vector_table, vector_column, query_vector,
        k=k, metric=metric, columns=[join_column],
    )

    if not vector_result:
        return Result([], [])

    # Extract IDs from vector results
    ids = [record[join_column] for record in vector_result.records]

    if not ids or not cypher:
        return vector_result

    # Step 2: Execute Cypher with the found IDs
    # The cypher query should use %s for the ID list
    if cypher_params:
        full_params = (*cypher_params, ids)
    else:
        full_params = (ids,)

    graph_result = driver.cypher(cypher, params=full_params, graph=graph)
    return graph_result
