#!/usr/bin/env python3

import importlib.util
import sys
import time
import types
from pathlib import Path

fake_psycopg2 = types.ModuleType("psycopg2")
fake_psycopg2.Error = Exception
fake_psycopg2.OperationalError = OSError
fake_extras = types.ModuleType("psycopg2.extras")
fake_extras.RealDictCursor = object
fake_psycopg2.extras = fake_extras

sys.modules["psycopg2"] = fake_psycopg2
sys.modules["psycopg2.extras"] = fake_extras

spec = importlib.util.spec_from_file_location(
    "postgresql_plugin",
    Path("plugins/postgresql/plugin.py"),
)
plugin = importlib.util.module_from_spec(spec)
spec.loader.exec_module(plugin)


class Cursor:
    def __enter__(self):
        return self

    def __exit__(self, *args):
        return False

    def execute(self, query):
        pass

    def fetchone(self):
        return {
            "in_recovery": True,
            "replay_paused": False,
            "receive_lsn": "14/1C07A910",
            "replay_lsn": "14/1C07A910",
            "receiver_status": "streaming",
            "slot_name": "ha_pg_david_eve",
        }


class Connection:
    autocommit = False

    def cursor(self, cursor_factory=None):
        return Cursor()

    def close(self):
        pass


plugin.monitor = lambda config, runtime_state: {
    "ha": {"ocf_rc": plugin.OCF_RUNNING_PROMOTED},
    "primary": {"current_wal_lsn": "14/1C07A920"},
}

plugin.read_secret = lambda path: "test-password"
plugin.psycopg2.connect = lambda **kwargs: Connection()

runtime_state = {
    "postgresql": {
        "notification": {
            "type": "pre",
            "operation": "demote",
            "demote_uname": "david-prodesk",
            "promote_uname": "david-eve",
            "received_at": time.time(),
        },
    },
}

config = {
    "PACEMAKER_NODE_NAME": "david-prodesk",
    "PG_REPLICATION_PEERS": (
        "david-prodesk=192.168.1.4:5432,"
        "david-eve=192.168.1.3:5432"
    ),
    "PGPASSWORD_FILE": "/run/secrets/postgres_password",
    "PGUSER": "david",
    "PGDATABASE": "david",
}

result = plugin.demote_database(config, runtime_state, {})

assert result["ocf_rc"] == 0
assert result["ready"] is True
assert result["mutation"] is False
assert result["final_handoff_required"] is True
assert result["successor_node"] == "david-eve"
assert result["successor_slot_name"] == "ha_pg_david_eve"
assert result["checks"]["successor_streaming"] is True
assert result["checks"]["successor_replay_not_paused"] is True
assert result["wal_lag_bytes"] == 16

print("valid_database_probe     accepted")
print("wal_lag_reported         accepted")
print()
print("DEMOTE_DATABASE_TEST_PASS=1")


plugin.monitor = lambda config, runtime_state: {
    "ha": {"ocf_rc": plugin.OCF_SUCCESS},
    "primary": None,
}

result = plugin.demote_database(config, runtime_state, {})

assert result["ocf_rc"] == 1
assert result["reason"] == "local_not_primary"
print("local_not_primary        rejected")


plugin.monitor = lambda config, runtime_state: {
    "ha": {"ocf_rc": plugin.OCF_RUNNING_PROMOTED},
    "primary": {"current_wal_lsn": "14/1C07A920"},
}


def unavailable(**kwargs):
    raise plugin.psycopg2.OperationalError(
        "connection refused"
    )


plugin.psycopg2.connect = unavailable

result = plugin.demote_database(config, runtime_state, {})

assert result["ocf_rc"] == 1
assert result["reason"] == "successor_unreachable"
print("successor_unreachable    rejected")

print("DEMOTE_DATABASE_FAILURE_TEST_PASS=1")

print()
print("===== SELF-RESTART DATABASE GATE =====")

plugin.monitor = lambda config, runtime_state: {
    "ha": {"ocf_rc": plugin.OCF_RUNNING_PROMOTED},
    "primary": {"current_wal_lsn": "14/1C07A920"},
}

self_restart_state = {
    "postgresql": {
        "notification": {
            "type": "pre",
            "operation": "demote",
            "demote_uname": "david-prodesk",
            "promote_uname": "david-prodesk",
            "promoted_uname": "david-prodesk",
            "unpromoted_uname": "david-eve",
            "received_at": time.time(),
        },
    },
}

connect_called = False

def must_not_connect(**kwargs):
    global connect_called
    connect_called = True
    raise AssertionError(
        "self-restart demotion must not probe a remote successor"
    )

plugin.psycopg2.connect = must_not_connect

result = plugin.demote_database(
    config,
    self_restart_state,
    {},
)

assert result["ocf_rc"] == 0
assert result["ready"] is True
assert result["mutation"] is False
assert result["final_handoff_required"] is False
assert result["handoff_mode"] == "local_restart"
assert result["reason"] == "local_restart_demotion_ready"
assert result["successor_node"] == "david-prodesk"
assert result["checks"]["local_primary"] is True
assert result["checks"]["remote_handoff_required"] is False
assert connect_called is False

print("self_restart_database     accepted")
print("remote_probe_skipped      accepted")