"""Piggie — Python SDK for Apache AGE + pgvector on PostgreSQL."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from piggie._cypher.builder import Create, Match, Merge
from piggie._util import validate_label
from piggie._version import __version__

if TYPE_CHECKING:
    from piggie.driver import PiggieDriver, PiggiePool

__all__ = [
    "__version__",
    "connect",
    "connect_pool",
    "validate_label",
    "Match",
    "Create",
    "Merge",
]


def connect(
    dsn: str = "",
    *,
    graph: str | None = None,
    **kwargs: Any,
) -> PiggieDriver:
    """Connect to PostgreSQL with AGE and pgvector support.

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

    Returns:
        A connected PiggieDriver instance.
    """
    from piggie.driver import PiggieDriver

    return PiggieDriver(dsn, graph=graph, **kwargs)


def connect_pool(
    dsn: str = "",
    *,
    graph: str | None = None,
    min_size: int = 2,
    max_size: int = 10,
    **kwargs: Any,
) -> PiggiePool:
    """Create a connection pool for multi-threaded applications.

    Args:
        dsn: PostgreSQL connection string.
        graph: Default graph name.
        min_size: Minimum pool connections.
        max_size: Maximum pool connections.
        **kwargs: Additional arguments passed to psycopg.connect().

    Returns:
        A PiggiePool instance.
    """
    from piggie.driver import PiggiePool

    return PiggiePool(dsn, graph=graph, min_size=min_size, max_size=max_size, **kwargs)
