"""Vector similarity search using pgvector."""

from __future__ import annotations

from typing import Any

import psycopg

from piggie.exceptions import VectorError
from piggie.result import Result

# pgvector operator mapping
_METRIC_OPS = {
    "cosine": "<=>",
    "l2": "<->",
    "inner_product": "<#>",
}


def vector_search(
    conn: psycopg.Connection[Any],
    table: str,
    column: str,
    query_vector: list[float],
    *,
    k: int = 10,
    metric: str = "cosine",
    where: str | None = None,
    columns: list[str] | None = None,
) -> Result:
    """Execute a k-NN vector search using pgvector.

    Args:
        conn: psycopg connection.
        table: Table containing the vector column.
        column: Name of the vector column.
        query_vector: Query vector.
        k: Number of nearest neighbors.
        metric: Distance metric: "cosine", "l2", or "inner_product".
        where: Optional WHERE clause (without the WHERE keyword).
        columns: Columns to select. Defaults to all.

    Returns:
        Result with rows ordered by distance, including a _distance column.
    """
    op = _METRIC_OPS.get(metric)
    if op is None:
        raise VectorError(
            f"Unknown metric {metric!r}. Use one of: {list(_METRIC_OPS.keys())}"
        )

    vec_str = "[" + ",".join(str(v) for v in query_vector) + "]"

    select_cols = ", ".join(columns) if columns else "*"
    distance_expr = f"{column} {op} %s::vector"

    parts = [
        f"SELECT {select_cols}, {distance_expr} AS _distance",
        f"FROM {table}",
    ]
    if where:
        # Escape literal % in the WHERE clause so ClientCursor doesn't
        # interpret them as parameter markers.
        parts.append(f"WHERE {where.replace('%', '%%')}")
    parts.append(f"ORDER BY {distance_expr}")
    parts.append(f"LIMIT {int(k)}")

    query = " ".join(parts)

    cursor = conn.cursor()
    try:
        # Pass vec_str twice: once for SELECT, once for ORDER BY
        cursor.execute(query, (vec_str, vec_str))
        if cursor.description:
            result_columns = [desc[0] for desc in cursor.description]
            rows = cursor.fetchall()
        else:
            result_columns = []
            rows = []
    except Exception as exc:
        conn.rollback()
        raise VectorError(f"Vector search failed: {exc}") from exc

    return Result(rows, result_columns)
