"""AGE connection setup: extension loading, type registration, graph management."""

from __future__ import annotations

from typing import Any

import psycopg
from psycopg import sql
from psycopg.client_cursor import ClientCursor
from psycopg.types import TypeInfo

from piggie._agtype.loader import AgtypeLoader
from piggie.exceptions import ConnectionError


def connect(dsn: str = "", *, graph: str | None = None, **kwargs: Any) -> psycopg.Connection[Any]:
    """Connect to PostgreSQL and set up AGE.

    Creates a psycopg3 connection with ClientCursor (for mogrify support),
    loads the AGE extension, registers the agtype type adapter, and
    optionally creates the specified graph.

    Args:
        dsn: PostgreSQL connection string.
        graph: Default graph name. Created if it doesn't exist.
        **kwargs: Additional arguments passed to psycopg.connect().

    Returns:
        A configured psycopg connection.
    """
    try:
        conn = psycopg.connect(dsn, cursor_factory=ClientCursor, **kwargs)
    except Exception as exc:
        raise ConnectionError(f"Failed to connect to PostgreSQL: {exc}") from exc

    try:
        _setup_age(conn)
    except Exception as exc:
        conn.close()
        raise ConnectionError(f"Failed to set up AGE: {exc}") from exc

    if graph is not None:
        ensure_graph(conn, graph)

    return conn


def _setup_age(conn: psycopg.Connection[Any]) -> None:
    """Load AGE extension and register agtype adapter."""
    with conn.cursor() as cur:
        cur.execute("LOAD 'age';")
        cur.execute("SET search_path = ag_catalog, '$user', public;")

    ag_info = TypeInfo.fetch(conn, "agtype")
    if not ag_info:
        raise ConnectionError(
            "AGE agtype type not found. Ensure the AGE extension is installed "
            "and loaded. Run: CREATE EXTENSION age;"
        )

    conn.adapters.register_loader(ag_info.oid, AgtypeLoader)
    conn.adapters.register_loader(ag_info.array_oid, AgtypeLoader)


def ensure_graph(conn: psycopg.Connection[Any], graph_name: str) -> None:
    """Create the graph if it doesn't exist."""
    with conn.cursor() as cur:
        cur.execute(
            sql.SQL("SELECT count(*) FROM ag_graph WHERE name={name}").format(
                name=sql.Literal(graph_name)
            )
        )
        row = cur.fetchone()
        if row is not None and row[0] == 0:
            cur.execute(
                sql.SQL("SELECT create_graph({name})").format(
                    name=sql.Literal(graph_name)
                )
            )
            conn.commit()


def drop_graph(conn: psycopg.Connection[Any], graph_name: str) -> None:
    """Drop a graph and all its data."""
    with conn.cursor() as cur:
        cur.execute(
            sql.SQL("SELECT drop_graph({name}, true)").format(
                name=sql.Literal(graph_name)
            )
        )
        conn.commit()


def list_graphs(conn: psycopg.Connection[Any]) -> list[str]:
    """List all AGE graphs in the database."""
    with conn.cursor() as cur:
        cur.execute("SELECT name FROM ag_graph ORDER BY name")
        return [row[0] for row in cur.fetchall()]
