"""Watch pg_stat_progress_repack live while VACUUM FULL / CLUSTER / REPACK run.

usage: progress_live.py <port> <label>

Creates a table with a TOAST table and three indexes, then for each command
polls the progress view from a second connection every few milliseconds and
prints every change of (phase, index_rebuild_count) with its time.
"""
import sys
import threading
import time

import psycopg

PORT, LABEL = sys.argv[1], sys.argv[2]
POLL = float(sys.argv[3]) if len(sys.argv) > 3 else 0.005
DSN = f"host=/tmp port={PORT} user=postgres dbname=postgres"
ROWS = 1_500_000

COMMANDS = [
    "VACUUM FULL progress_t",
    "CLUSTER progress_t USING progress_t_pkey",
    "REPACK progress_t",
    "REPACK (CONCURRENTLY) progress_t",
]


def setup():
    with psycopg.connect(DSN, autocommit=True) as c:
        c.execute("DROP TABLE IF EXISTS progress_t")
        c.execute("CREATE TABLE progress_t (id int PRIMARY KEY, a int, b text, t text)")
        c.execute(f"""
            INSERT INTO progress_t
            SELECT g, g % 9973, md5(g::text),
                   CASE WHEN g % 500 = 0
                        THEN (SELECT string_agg(md5((g * x)::text), '')
                              FROM generate_series(1, 400) x)
                   END
            FROM generate_series(1, {ROWS}) g""")
        c.execute("CREATE INDEX progress_t_a ON progress_t (a)")
        c.execute("CREATE INDEX progress_t_b ON progress_t (b)")
        c.execute("ANALYZE progress_t")
        n_idx = c.execute("SELECT count(*) FROM pg_index WHERE indrelid = 'progress_t'::regclass").fetchone()[0]
        toast = c.execute("""SELECT t.relname, pg_size_pretty(pg_relation_size(t.oid))
                             FROM pg_class c JOIN pg_class t ON t.oid = c.reltoastrelid
                             WHERE c.relname = 'progress_t'""").fetchone()
        ver = c.execute("SELECT version()").fetchone()[0].split(" on ")[0]
    print(f"### {LABEL}: {ver}")
    print(f"table progress_t: {ROWS} rows, {n_idx} indexes, TOAST {toast[0]} = {toast[1]}\n")
    return n_idx


def watch(command):
    done = threading.Event()
    error = []

    def run():
        try:
            with psycopg.connect(DSN, autocommit=True) as c:
                c.execute(command)
        except Exception as e:  # reported below, not swallowed
            error.append(e)
        finally:
            done.set()

    samples, last = 0, None
    timeline = []
    with psycopg.connect(DSN, autocommit=True) as mon:
        worker = threading.Thread(target=run)
        t0 = time.monotonic()
        worker.start()
        while not done.is_set():
            row = mon.execute(
                "SELECT phase, index_rebuild_count FROM pg_stat_progress_repack "
                "WHERE relid = 'progress_t'::regclass").fetchone()
            samples += 1
            if row is not None and row != last:
                timeline.append((int((time.monotonic() - t0) * 1000), row[0], row[1]))
                last = row
            time.sleep(POLL)
        worker.join()
    elapsed = int((time.monotonic() - t0) * 1000)

    print(f"-- {command}   ({elapsed} ms, {samples} samples)")
    if error:
        print(f"   ERROR: {error[0]}")
    for ms, phase, count in timeline:
        print(f"   {ms:>6} ms  {phase:<28} index_rebuild_count = {count}")
    print()


n_idx = setup()
for cmd in COMMANDS:
    watch(cmd)
print(f"(the table has {n_idx} indexes: a correct count never goes above {n_idx})")