#!/usr/bin/env python3

import psycopg2
import psycopg2.extras


def read_secret(path):
    with open(path, "r", encoding="utf-8") as f:
        return f.read().strip()


def monitor(config):
    password = read_secret(
        config["PGPASSWORD_FILE"]
    )

    conn = psycopg2.connect(
        host=config["PGHOST"],
        port=int(config["PGPORT"]),
        user=config["PGUSER"],
        dbname=config["PGDATABASE"],
        password=password,
        connect_timeout=5,
        application_name="managed-containerd",
    )

    try:
        conn.autocommit = True

        with conn.cursor(
            cursor_factory=psycopg2.extras.RealDictCursor
        ) as cur:

            cur.execute("""
                SELECT pg_is_in_recovery() AS in_recovery
            """)

            in_recovery = bool(
                cur.fetchone()["in_recovery"]
            )

            if not in_recovery:
                cur.execute("""
                    SELECT
                        pg_current_wal_lsn()::text
                            AS current_wal_lsn
                """)

                primary = cur.fetchone()

            else:
                primary = None

            if in_recovery:
                cur.execute("""
                    SELECT
                        pg_last_wal_receive_lsn()::text
                            AS receive_lsn,
                        pg_last_wal_replay_lsn()::text
                            AS replay_lsn,
                        pg_is_wal_replay_paused()
                            AS replay_paused
                """)

                recovery = cur.fetchone()

                cur.execute("""
                    SELECT
                        status,
                        sender_host,
                        sender_port,
                        slot_name,
                        latest_end_lsn::text
                            AS latest_end_lsn
                    FROM pg_stat_wal_receiver
                    LIMIT 1
                """)

                receiver = cur.fetchone()

            else:
                recovery = {
                    "receive_lsn": None,
                    "replay_lsn": None,
                    "replay_paused": False,
                }

                receiver = None

            cur.execute("""
                SELECT
                    current_setting('server_version')
                        AS server_version,
                    current_setting('cluster_name', true)
                        AS cluster_name,
                    current_user AS current_user,
                    current_database()
                        AS current_database,
                    inet_server_addr()::text
                        AS server_address,
                    inet_server_port()
                        AS server_port
            """)

            server = cur.fetchone()

        if in_recovery:
            role = "standby"

            streaming = bool(
                receiver
                and receiver.get("status")
                    == "streaming"
            )

        else:
            role = "primary"
            streaming = False

        return {
            "role": role,
            "in_recovery": in_recovery,
            "streaming": streaming,
            "wal_receiver": (
                dict(receiver)
                if receiver
                else None
            ),
            "recovery": dict(recovery),
            "primary": (
                dict(primary)
                if primary
                else None
            ),
            "server": dict(server),
        }

    finally:
        conn.close()