#!/usr/bin/env python3
"""Repeatedly kill blockers holding DROP INDEX CONCURRENTLY for stalled index."""

from __future__ import annotations

import os
import time
from pathlib import Path

import psycopg2


ENV_FILE = Path("/etc/default/domaincheck-worker")
INDEX_NAME = "idx_detect_job_items_stalled_job_activity"
POLL_SECONDS = 3
MAX_LOOPS = 1200


def load_env(path: Path) -> None:
    if not path.exists():
        return
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        os.environ[key] = value


def log(message: str) -> None:
    stamp = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{stamp}] {message}", flush=True)


def main() -> int:
    load_env(ENV_FILE)
    conn = psycopg2.connect(
        host="127.0.0.1",
        port=5432,
        user=os.getenv("DB_USER", "postgres"),
        password=os.getenv("DB_PASSWORD"),
        dbname="domain",
    )
    conn.autocommit = True
    cur = conn.cursor()
    try:
        for loop in range(MAX_LOOPS):
            cur.execute(
                """
                SELECT pid
                FROM pg_stat_activity
                WHERE pid <> pg_backend_pid()
                  AND state <> 'idle'
                  AND query ILIKE %s
                ORDER BY pid
                """,
                (f"%DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME}%",),
            )
            drop_pids = [r[0] for r in cur.fetchall()]
            if not drop_pids:
                log("no_drop_pid")
                break

            blockers: list[int] = []
            for pid in drop_pids:
                cur.execute("SELECT pg_blocking_pids(%s)", (pid,))
                row = cur.fetchone()
                if row and row[0]:
                    blockers.extend(int(v) for v in row[0])

            blockers = sorted(set(blockers))
            log(f"loop={loop} drop_pids={drop_pids} blockers={blockers}")
            if not blockers:
                time.sleep(POLL_SECONDS)
                continue

            for pid in blockers:
                cur.execute("SELECT pg_cancel_backend(%s), pg_terminate_backend(%s)", (pid, pid))
                log(f"kill pid={pid} result={cur.fetchone()}")
            time.sleep(POLL_SECONDS)
    finally:
        cur.close()
        conn.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())