"""Schema introspection for AGE graphs and pgvector indexes."""

from __future__ import annotations

from typing import Any

import psycopg
from psycopg import sql

from piggie._vector.index import list_vector_indexes
from piggie.exceptions import SchemaError


def get_schema(conn: psycopg.Connection[Any], graph_name: str) -> dict[str, Any]:
    """Get schema information for a graph.

    Returns:
        Dict with keys:
        - vertex_labels: list of vertex label dicts (name, estimated_count)
        - edge_labels: list of edge label dicts (name, estimated_count)
        - vector_indexes: list of vector index dicts
    """
    try:
        vertex_labels = _get_labels(conn, graph_name, "v", "_ag_label_vertex")
        edge_labels = _get_labels(conn, graph_name, "e", "_ag_label_edge")
        vector_indexes = list_vector_indexes(conn)
    except SchemaError:
        raise
    except Exception as exc:
        raise SchemaError(f"Schema introspection failed: {exc}") from exc

    return {
        "graph": graph_name,
        "vertex_labels": vertex_labels,
        "edge_labels": edge_labels,
        "vector_indexes": vector_indexes,
    }


def _get_labels(
    conn: psycopg.Connection[Any],
    graph_name: str,
    kind: str,
    exclude: str,
) -> list[dict[str, Any]]:
    """Get labels and their approximate row counts.

    The ag_label.relation column is a regclass. We cast it to oid
    to look up pg_class.reltuples for approximate counts.
    """
    cursor = conn.cursor()
    cursor.execute(
        sql.SQL(
            "SELECT l.name, c.reltuples::bigint "
            "FROM ag_label l "
            "JOIN pg_class c ON c.oid = l.relation::oid "
            "WHERE l.graph = (SELECT graphid FROM ag_graph WHERE name = {graph}) "
            "AND l.kind = {kind} AND l.name != {exclude} "
            "ORDER BY l.name"
        ).format(
            graph=sql.Literal(graph_name),
            kind=sql.Literal(kind),
            exclude=sql.Literal(exclude),
        )
    )

    labels: list[dict[str, Any]] = []
    for row in cursor.fetchall():
        count = int(row[1]) if row[1] >= 0 else 0
        labels.append({"name": row[0], "estimated_count": count})

    return labels
