import os
import time

import psycopg


def connect() -> psycopg.Connection:
    dsn = os.environ["DATABASE_URL"]
    return psycopg.connect(dsn)


def wait_for_db(timeout: float = 60.0, interval: float = 2.0) -> psycopg.Connection:
    """Block until the database is up and the `Job` table exists.

    Repeatedly opens a connection and probes for the `Job` table's presence.
    Tolerates the DB not being reachable yet (`psycopg.OperationalError`) and
    migrations not having run yet (`psycopg.errors.UndefinedTable`), retrying
    with a fixed `interval` until `timeout` elapses. Returns the open, usable
    connection on success; re-raises the last error on timeout.
    """
    deadline = time.monotonic() + timeout
    printed_waiting = False
    last_error: Exception | None = None
    while time.monotonic() < deadline:
        conn: psycopg.Connection | None = None
        try:
            conn = connect()
            with conn.cursor() as cur:
                cur.execute('SELECT 1 FROM "Job" LIMIT 0')
            return conn
        except (psycopg.OperationalError, psycopg.errors.UndefinedTable) as exc:
            last_error = exc
            if conn is not None:
                conn.close()
            if not printed_waiting:
                print("worker: waiting for database schema...", flush=True)
                printed_waiting = True
            time.sleep(interval)
    assert last_error is not None
    raise last_error