"""Vector index management for pgvector."""

from __future__ import annotations

from typing import Any

import psycopg

from piggie.exceptions import VectorError

# Maps (method, metric) to pgvector operator class
_OPS_CLASS = {
    ("hnsw", "cosine"): "vector_cosine_ops",
    ("hnsw", "l2"): "vector_l2_ops",
    ("hnsw", "inner_product"): "vector_ip_ops",
    ("ivfflat", "cosine"): "vector_cosine_ops",
    ("ivfflat", "l2"): "vector_l2_ops",
    ("ivfflat", "inner_product"): "vector_ip_ops",
}


def create_vector_index(
    conn: psycopg.Connection[Any],
    table: str,
    column: str,
    *,
    method: str = "hnsw",
    metric: str = "cosine",
    name: str | None = None,
    **kwargs: Any,
) -> None:
    """Create a vector index on a table column.

    Args:
        conn: psycopg connection.
        table: Table name.
        column: Vector column name.
        method: "hnsw" or "ivfflat".
        metric: "cosine", "l2", or "inner_product".
        name: Optional index name. Auto-generated if None.
        **kwargs: Index parameters (m, ef_construction for HNSW; lists for IVFFlat).
    """
    ops_class = _OPS_CLASS.get((method, metric))
    if ops_class is None:
        raise VectorError(
            f"Unknown method/metric combination: {method}/{metric}"
        )

    index_name = name or f"idx_{table}_{column}_{method}"

    with_params: list[str] = []
    if method == "hnsw":
        if "m" in kwargs:
            with_params.append(f"m = {int(kwargs['m'])}")
        if "ef_construction" in kwargs:
            with_params.append(f"ef_construction = {int(kwargs['ef_construction'])}")
    elif method == "ivfflat":
        if "lists" in kwargs:
            with_params.append(f"lists = {int(kwargs['lists'])}")

    with_clause = f" WITH ({', '.join(with_params)})" if with_params else ""

    stmt = (
        f"CREATE INDEX IF NOT EXISTS {index_name} ON {table} "
        f"USING {method} ({column} {ops_class}){with_clause}"
    )

    cursor = conn.cursor()
    try:
        cursor.execute(stmt)
        conn.commit()
    except Exception as exc:
        conn.rollback()
        raise VectorError(f"Failed to create vector index: {exc}") from exc


def drop_vector_index(
    conn: psycopg.Connection[Any],
    name: str,
) -> None:
    """Drop a vector index by name."""
    cursor = conn.cursor()
    try:
        cursor.execute(f"DROP INDEX IF EXISTS {name}")
        conn.commit()
    except Exception as exc:
        conn.rollback()
        raise VectorError(f"Failed to drop vector index: {exc}") from exc


def list_vector_indexes(
    conn: psycopg.Connection[Any],
    table: str | None = None,
) -> list[dict[str, Any]]:
    """List vector indexes, optionally filtered by table.

    Returns:
        List of dicts with keys: name, table, column, method, ops_class.
    """
    query = (
        "SELECT indexname, tablename, indexdef "
        "FROM pg_indexes "
        "WHERE indexdef LIKE %s"
    )
    params: list[str] = ["%vector_%_ops%"]
    if table:
        query += " AND tablename = %s"
        params.append(table)

    cursor = conn.cursor()
    cursor.execute(query, params)

    indexes: list[dict[str, Any]] = []
    for row in cursor.fetchall():
        indexes.append({
            "name": row[0],
            "table": row[1],
            "definition": row[2],
        })
    return indexes
