"""
Amazon Aurora PostgreSQL connection pool — the structured data layer.

This sits ALONGSIDE the RAG path (ChromaDB + Sheet + Claude), which is
untouched. Aurora holds the relational product data: organizations, users,
opportunities, curators, bookings, conversations/messages, inquiries — and
the analytics views the admin dashboard reads. See schema.sql for the model.

One pool per process (avoids serverless connection exhaustion). The pool is
opened lazily so the app still boots when the database is not configured —
the data-layer routes then return 503 instead of crashing the whole Space.

----------------------------------------------------------------------------
AUTH — two supported modes
----------------------------------------------------------------------------
This cluster (studentcompanionai) authenticates with IAM, not a static
password (Secrets Manager holds no retrievable secret). So we support both:

  1. PASSWORD mode — set DATABASE_URL to a full libpq URL:
       postgresql://USER:PASSWORD@<endpoint>:5432/postgres?sslmode=require
     Used as-is. Simplest, if a master password ever gets set.

  2. IAM mode (this cluster) — set these instead of DATABASE_URL:
       DB_HOST   = studentcompanionai.cluster-xxxx.us-east-1.rds.amazonaws.com
       DB_USER   = postgres
       DB_NAME   = postgres            (optional, defaults to 'postgres')
       DB_REGION = us-east-1           (optional, defaults to AWS_REGION)
       AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY  (HF Space secrets)
     A fresh 15-min RDS auth token is generated as the password EACH TIME a
     physical connection is opened (pool `configure` hook), so token expiry
     never breaks the long-lived pool. Requires, once, on the DB:
         GRANT rds_iam TO postgres;

Mode is auto-selected: if DATABASE_URL is set we use PASSWORD mode, else if
DB_HOST is set we use IAM mode, else the data layer is disabled.
"""

import os
import logging

logger = logging.getLogger("db")

_pool = None  # lazily-initialised psycopg_pool.ConnectionPool
_init_attempted = False


def _iam_configured() -> bool:
    return bool(os.environ.get("DB_HOST", "").strip())


def _password_configured() -> bool:
    return bool(os.environ.get("DATABASE_URL", "").strip())


def _region() -> str:
    return (
        os.environ.get("DB_REGION")
        or os.environ.get("AWS_REGION")
        or os.environ.get("AWS_DEFAULT_REGION")
        or "us-east-1"
    ).strip()


def _build_password_pool():
    """PASSWORD mode: DATABASE_URL used verbatim as the connection string."""
    from psycopg_pool import ConnectionPool
    from psycopg.rows import dict_row

    dsn = os.environ["DATABASE_URL"].strip()
    pool = ConnectionPool(
        dsn,
        min_size=1,
        max_size=10,
        kwargs={"row_factory": dict_row},
        open=True,
    )
    pool.wait(timeout=10.0)
    return pool


def _build_iam_pool():
    """
    IAM mode: connect with a freshly-minted RDS auth token as the password.

    psycopg_pool has no per-connection password hook, so we subclass
    psycopg.Connection and override `connect` to inject a fresh 15-min RDS
    auth token at the moment each physical connection is opened. The pool
    calls this on initial fill AND on every reconnect, so a token never
    outlives the connection that used it — long-lived pools keep working.
    """
    import boto3
    from psycopg_pool import ConnectionPool
    from psycopg.rows import dict_row
    import psycopg

    host = os.environ["DB_HOST"].strip()
    user = os.environ.get("DB_USER", "postgres").strip()
    dbname = os.environ.get("DB_NAME", "postgres").strip()
    port = int(os.environ.get("DB_PORT", "5432"))
    region = _region()

    rds = boto3.client("rds", region_name=region)

    class _IAMConnection(psycopg.Connection):
        @classmethod
        def connect(cls, conninfo="", **kwargs):
            token = rds.generate_db_auth_token(
                DBHostname=host, Port=port, DBUsername=user, Region=region
            )
            kwargs.update(
                host=host,
                port=port,
                dbname=dbname,
                user=user,
                password=token,
                sslmode="require",
            )
            return super().connect("", **kwargs)

    pool = ConnectionPool(
        connection_class=_IAMConnection,
        kwargs={"row_factory": dict_row},
        min_size=1,
        max_size=10,
        # This Serverless v2 cluster auto-pauses (min 0 ACUs); a cold resume
        # can take 30s+. Allow generous time for the first connection and a
        # long checkout window before a request gives up.
        timeout=60.0,
        open=True,
    )
    # Try to warm the pool, but DON'T fail the whole data layer if the cluster
    # is mid-resume — keep the pool so later requests succeed once it's warm.
    try:
        pool.wait(timeout=45.0)
        logger.info("[db] IAM pool warm")
    except Exception as e:
        logger.warning("[db] pool not warm yet (cluster resuming?): %s", e)
    return pool


def _build_pool():
    """Create the connection pool, or return None if it can't be built."""
    if not _password_configured() and not _iam_configured():
        logger.error(
            "[db] No DB config — set DATABASE_URL (password) or DB_HOST (IAM). "
            "Aurora data-layer routes will 503."
        )
        return None
    try:
        if _password_configured():
            pool = _build_password_pool()
            logger.info("[db] Aurora pool ready (password mode)")
        else:
            pool = _build_iam_pool()
            logger.info("[db] Aurora pool ready (IAM-token mode)")
        return pool
    except ImportError as e:
        logger.error('[db] Missing dependency: %s — need "psycopg[binary,pool]" and boto3', e)
        return None
    except Exception as e:
        logger.error("[db] Failed to open Aurora pool: %s", e)
        return None


def get_pool():
    """
    Return the process-wide connection pool, building it on first use.

    Returns None if no DB is configured or the pool can't be opened; callers
    should treat None as "data layer unavailable" and raise 503.
    """
    global _pool, _init_attempted
    if _pool is None and not _init_attempted:
        _init_attempted = True
        _pool = _build_pool()
    return _pool


def db_enabled() -> bool:
    return get_pool() is not None


def get_status() -> dict:
    """
    Health snapshot for /api/services/status. Actually probes a connection so
    the status reflects whether queries work right now — not just whether the
    pool object exists (it may be mid-resume on this auto-pausing cluster).
    """
    mode = "password" if _password_configured() else ("iam" if _iam_configured() else "unconfigured")
    pool = get_pool()
    if pool is None:
        return {"enabled": False, "mode": mode, "reachable": False}

    reachable = False
    detail = None
    try:
        with pool.connection(timeout=60.0) as conn:
            conn.execute("SELECT 1")
        reachable = True
    except Exception as e:
        detail = str(e)[:200]

    out = {"enabled": True, "mode": mode, "reachable": reachable}
    if detail:
        out["error"] = detail
    return out