"""Trino connections via the Andlake gateway (notebook-service).

The gateway proxies Trino queries on port 8082, enforcing tenant isolation
and attaching the JWT access token for authentication.

Usage::

    from andlake import get_trino_connection
    import pandas as pd

    conn = get_trino_connection()
    df = pd.read_sql("SELECT * FROM lake.silver.transactions LIMIT 10", conn)

For SQLAlchemy (e.g. with ``pd.read_sql`` engine mode)::

    from andlake import get_trino_engine
    engine = get_trino_engine()
    df = pd.read_sql("SELECT 1", engine)
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from andlake.config import get_config

if TYPE_CHECKING:
    from sqlalchemy.engine import Engine
    from trino.dbapi import Connection


def get_trino_connection(
    *,
    catalog: str | None = None,
    schema: str | None = None,
) -> Connection:
    """Create a ``trino.dbapi.Connection`` pointing at the Andlake gateway.

    Parameters
    ----------
    catalog:
        Trino catalog name. Defaults to ``ANDLAKE_DEFAULT_CATALOG`` (``lake``).
    schema:
        Trino schema (Iceberg namespace). Defaults to ``None``.

    Returns
    -------
    trino.dbapi.Connection
    """
    import trino
    import trino.dbapi

    cfg = get_config()

    connect_args = dict(
        host=cfg.trino_host,
        port=cfg.trino_port,
        catalog=catalog or cfg.default_catalog,
        schema=schema,
        http_scheme="http",
    )

    if cfg.tenant_id:
        connect_args["user"] = cfg.tenant_id

    if cfg.access_token:
        import trino.auth

        connect_args["auth"] = trino.auth.JWTAuthentication(cfg.access_token)

    return trino.dbapi.connect(**connect_args)


def get_trino_engine(
    *,
    catalog: str | None = None,
    schema: str | None = None,
) -> Engine:
    """Create a SQLAlchemy ``Engine`` for Trino via the Andlake gateway.

    Useful with ``pd.read_sql(sql, engine)`` and other SQLAlchemy-aware tools.

    Parameters
    ----------
    catalog:
        Trino catalog name. Defaults to ``ANDLAKE_DEFAULT_CATALOG`` (``lake``).
    schema:
        Trino schema (Iceberg namespace). Defaults to ``None``.

    Returns
    -------
    sqlalchemy.engine.Engine
    """
    from sqlalchemy import create_engine

    cfg = get_config()
    cat = catalog or cfg.default_catalog
    schema_part = f"/{schema}" if schema else ""

    user_part = f"{cfg.tenant_id}@" if cfg.tenant_id else ""
    url = (
        f"trino://{user_part}{cfg.trino_host}:{cfg.trino_port}"
        f"/{cat}{schema_part}"
    )

    connect_args: dict = {"http_scheme": "http"}

    if cfg.access_token:
        import trino.auth

        connect_args["auth"] = trino.auth.JWTAuthentication(cfg.access_token)

    return create_engine(url, connect_args=connect_args)
