"""PostgreSQL connection helpers for Polleo Demand.

Two flavours of connection:
  - get_connection() returns a raw psycopg2 connection for executing
    SQL DDL / one-off statements (used by init_db.py, ad-hoc maintenance).
  - get_engine() returns a SQLAlchemy engine, cached at module level so
    pandas read_sql / to_sql share a single pool across the app.

is_db_available() is a no-raise probe — call it before doing anything
that assumes the DB is up. The CSV fallback layer (db/csv_fallback.py)
uses it to decide whether to attempt a DB query at all.
"""
from __future__ import annotations

import os
from typing import Optional
from urllib.parse import unquote, urlparse

import psycopg2
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine


def _config() -> dict:
    """Read connection params from env vars, with dev defaults that match
    docker-compose.yml.

    DATABASE_URL takes precedence when set — this is what the FastAPI app
    (backend/config.py) and docker-compose.prod.yml configure. Without this,
    a deployment that sets only DATABASE_URL (e.g. the prod compose points the
    api container at the `db` service) would fall back to DB_HOST's default of
    `localhost`, which has no Postgres inside the container — the stock-upload
    DB reload then fails with "connection refused on localhost:5432". The
    discrete DB_* vars remain a fallback for CLI scripts run outside the app.
    """
    url = os.environ.get("DATABASE_URL", "").strip()
    if url:
        p = urlparse(url)
        return {
            "host": p.hostname or "localhost",
            "port": p.port or 5432,
            "dbname": (p.path or "").lstrip("/") or "polleo_demand",
            "user": unquote(p.username) if p.username else "polleo",
            "password": unquote(p.password) if p.password else "polleo_dev",
        }
    return {
        "host": os.environ.get("DB_HOST", "localhost"),
        "port": int(os.environ.get("DB_PORT", "5432")),
        "dbname": os.environ.get("DB_NAME", "polleo_demand"),
        "user": os.environ.get("DB_USER", "polleo"),
        "password": os.environ.get("DB_PASSWORD", "polleo_dev"),
    }


def get_connection():
    """Return a raw psycopg2 connection. Caller is responsible for closing.

    Raises psycopg2.OperationalError if the DB is unreachable — wrap in a
    try/except or guard with is_db_available() first.
    """
    cfg = _config()
    return psycopg2.connect(
        host=cfg["host"],
        port=cfg["port"],
        dbname=cfg["dbname"],
        user=cfg["user"],
        password=cfg["password"],
    )


_engine: Optional[Engine] = None


def get_engine() -> Engine:
    """Return a process-wide SQLAlchemy engine with a small connection pool.

    pool_size=5 / max_overflow=10 is sized for a single-user Streamlit
    session running pandas read_sql calls. Bump if multiple concurrent
    page renders start blocking on connection checkout.
    """
    global _engine
    if _engine is None:
        cfg = _config()
        url = (
            f"postgresql+psycopg2://{cfg['user']}:{cfg['password']}"
            f"@{cfg['host']}:{cfg['port']}/{cfg['dbname']}"
        )
        _engine = create_engine(
            url,
            pool_size=5,
            max_overflow=10,
            pool_pre_ping=True,
            future=True,
        )
    return _engine


def is_db_available() -> bool:
    """Probe the DB with a short-lived connection. Never raises.

    Returns True iff a connection can be opened and closed cleanly.
    Used by db/csv_fallback.py before each query so the app stays usable
    when Postgres is not running.
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute("SELECT 1")
            cur.fetchone()
        return True
    except Exception:
        return False
    finally:
        if conn is not None:
            try:
                conn.close()
            except Exception:
                pass