"""Connection and query primitives.

Every notebook starts with:
    from utils.connection import init, query_df, execute
    conn = init()
"""

import os
from datetime import datetime, timezone

import pandas as pd
import psycopg
from dotenv import load_dotenv


def init(env_filename: str = ".prod.env") -> psycopg.Connection:
    """Load env, open connection, return it.

    Caller is responsible for conn.close() or using it in a `with` block.
    """
    _load_env(env_filename)
    dsn = os.environ["QC_TRACE_DSN"]
    return psycopg.connect(dsn)


def query_df(conn: psycopg.Connection, sql: str, params: tuple = ()) -> pd.DataFrame:
    """Run a SELECT and return a DataFrame."""
    with conn.cursor() as cur:
        cur.execute(sql, params)
        cols = [d.name for d in cur.description] if cur.description else []
        return pd.DataFrame(cur.fetchall(), columns=cols)


def execute(conn: psycopg.Connection, sql: str, params: tuple = ()) -> None:
    """Run a non-SELECT statement (SET, etc.)."""
    with conn.cursor() as cur:
        cur.execute(sql, params)


def scalar(conn: psycopg.Connection, sql: str, params: tuple = ()):
    """Return a single scalar value."""
    with conn.cursor() as cur:
        cur.execute(sql, params)
        row = cur.fetchone()
        return row[0] if row else None


def now_utc() -> datetime:
    """Current UTC datetime (convenience for param building)."""
    return datetime.now(timezone.utc)


# ---------------------------------------------------------------------------
# Internal
# ---------------------------------------------------------------------------

def _load_env(filename: str) -> None:
    repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    env_path = os.path.join(repo_root, filename)
    load_dotenv(env_path, override=True)
    assert os.environ.get("QC_TRACE_DSN"), f"QC_TRACE_DSN not found in {filename}"
