#!/usr/bin/env python3

import time

import psycopg2
import psycopg2.extras

OCF_SUCCESS = 0
OCF_ERR_GENERIC = 1
OCF_NOT_RUNNING = 7
OCF_RUNNING_PROMOTED = 8


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


def parse_replication_peers(value):
    peers = []
    for item in (value or "").split(","):
        item = item.strip()
        if not item:
            continue
        name, sep, endpoint = item.partition("=")
        if not sep:
            raise ValueError(f"invalid replication peer entry: {item!r}")
        host, sep, port = endpoint.rpartition(":")
        if not sep or not host or not port:
            raise ValueError(f"invalid replication peer endpoint: {endpoint!r}")
        peers.append({
            "name": name.strip(),
            "host": host.strip(),
            "port": int(port),
        })
    return peers


def probe_replication_peers(config, password):
    results = []
    writable_primaries = []

    for peer in parse_replication_peers(config["PG_REPLICATION_PEERS"]):
        result = dict(peer)
        try:
            conn = psycopg2.connect(
                host=peer["host"],
                port=peer["port"],
                user=config["PGUSER"],
                dbname=config["PGDATABASE"],
                password=password,
                connect_timeout=2,
                application_name="managed-containerd-peer-probe",
            )
            try:
                conn.autocommit = True
                with conn.cursor() as cur:
                    cur.execute("SELECT pg_is_in_recovery()")
                    in_recovery = bool(cur.fetchone()[0])
                result.update({
                    "reachable": True,
                    "in_recovery": in_recovery,
                    "writable_primary": not in_recovery,
                })
                if not in_recovery:
                    writable_primaries.append(peer["name"])
            finally:
                conn.close()
        except psycopg2.Error as exc:
            result.update({
                "reachable": False,
                "in_recovery": None,
                "writable_primary": False,
                "error": str(exc).splitlines()[0],
            })
        results.append(result)

    return {
        "peers": results,
        "writable_primaries": writable_primaries,
        "writable_primary_count": len(writable_primaries),
    }


def reconnect_grace_seconds(config):
    raw = config.get("PG_RECONNECT_GRACE_SECONDS") or "15"
    try:
        value = float(raw)
    except (TypeError, ValueError):
        raise ValueError(
            f"invalid PG_RECONNECT_GRACE_SECONDS: {raw!r}"
        )

    if value < 0:
        raise ValueError(
            "PG_RECONNECT_GRACE_SECONDS must be >= 0"
        )

    return value


def _postgres_runtime_state(runtime_state):
    return runtime_state.setdefault("postgresql", {})


def _set_transition_state(runtime_state, state):
    pg_state = _postgres_runtime_state(runtime_state)
    previous = pg_state.get("state")

    if previous != state:
        pg_state["state"] = state
        pg_state["state_changed_at"] = time.monotonic()

    return pg_state


def _clear_reconnect_grace(runtime_state):
    pg_state = _postgres_runtime_state(runtime_state)
    pg_state.pop("reconnect_started_at", None)


def normalize_ha_state(
    *,
    config,
    runtime_state,
    in_recovery,
    streaming,
    receiver,
    recovery,
    peer_state=None,
):
    """Translate only validated PostgreSQL states into OCF semantics."""
    if not in_recovery:
        _clear_reconnect_grace(runtime_state)
        _set_transition_state(runtime_state, "primary")
        return {
            "health": "healthy",
            "ha_role": "promoted",
            "state": "primary",
            "ocf_rc": OCF_RUNNING_PROMOTED,
            "promotion_eligible": True,
        }

    replay_paused = bool((recovery or {}).get("replay_paused"))
    if streaming and not replay_paused:
        _clear_reconnect_grace(runtime_state)
        _set_transition_state(runtime_state, "streaming")
        return {
            "health": "healthy",
            "ha_role": "unpromoted",
            "state": "streaming",
            "ocf_rc": OCF_SUCCESS,
            "promotion_eligible": True,
        }

    receiver_status = (receiver or {}).get("status")
    state = receiver_status or "not_streaming"
    writable_primary_count = None

    if peer_state is not None:
        writable_primary_count = peer_state.get(
            "writable_primary_count"
        )

    # Validated against the authoritative RA on 2026-09-17: if the local
    # server remains a reachable standby, its WAL receiver is not streaming,
    # and no writable primary is reachable, the RA returns OCF_SUCCESS and
    # keeps the node as a reduced-priority election candidate.
    if state == "not_streaming" and writable_primary_count == 0:
        _clear_reconnect_grace(runtime_state)
        _set_transition_state(runtime_state, "disconnected_candidate")
        return {
            "health": "degraded",
            "ha_role": "unpromoted",
            "state": "disconnected_candidate",
            "ocf_rc": OCF_SUCCESS,
            "promotion_eligible": True,
        }

    # Reconnect grace is PostgreSQL-specific policy layered on generic daemon
    # runtime state. It is entered only after this daemon has actually observed
    # disconnected_candidate and then sees exactly one writable primary return
    # while the local WAL receiver has not yet resumed streaming.
    if state == "not_streaming" and writable_primary_count == 1:
        pg_state = _postgres_runtime_state(runtime_state)
        previous_state = pg_state.get("state")
        started_at = pg_state.get("reconnect_started_at")
        now = time.monotonic()
        grace = reconnect_grace_seconds(config)

        if previous_state == "disconnected_candidate":
            started_at = now
            pg_state["reconnect_started_at"] = started_at

        if started_at is not None:
            elapsed = max(0.0, now - float(started_at))

            if elapsed < grace:
                _set_transition_state(runtime_state, "reconnecting")
                return {
                    "health": "degraded",
                    "ha_role": "unpromoted",
                    "state": "reconnecting",
                    "ocf_rc": OCF_SUCCESS,
                    "promotion_eligible": False,
                    "grace": {
                        "seconds": grace,
                        "elapsed_seconds": round(elapsed, 3),
                        "remaining_seconds": round(
                            max(0.0, grace - elapsed), 3
                        ),
                    },
                }

            _clear_reconnect_grace(runtime_state)

        # Already validated for the steady-state failure case: a reachable
        # standby whose WAL receiver is not streaming while exactly one remote
        # writable primary is available returns OCF_ERR_GENERIC.
        _set_transition_state(runtime_state, "not_streaming")
        return {
            "health": "degraded",
            "ha_role": "unpromoted",
            "state": "not_streaming",
            "ocf_rc": OCF_ERR_GENERIC,
            "promotion_eligible": False,
        }

    # Multiple writable primaries and other degraded states remain deliberately
    # unclassified until independently validated.
    _clear_reconnect_grace(runtime_state)
    _set_transition_state(runtime_state, state)
    return {
        "health": "degraded",
        "ha_role": "unpromoted",
        "state": state,
        "ocf_rc": None,
        "promotion_eligible": False,
    }


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

    try:
        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",
        )
    except psycopg2.OperationalError as exc:
        message = str(exc)

        # Validated against the authoritative RA on 2026-09-17:
        # an explicitly refused local PostgreSQL connection corresponds to
        # OCF_NOT_RUNNING.  Other OperationalError classes remain unclassified.
        if "connection refused" in message.lower():
            _clear_reconnect_grace(runtime_state)
            _set_transition_state(runtime_state, "not_running")
            return {
                "ha": {
                    "health": "down",
                    "ha_role": "unknown",
                    "state": "not_running",
                    "ocf_rc": OCF_NOT_RUNNING,
                    "promotion_eligible": False,
                },
                "role": "unknown",
                "in_recovery": None,
                "streaming": False,
                "wal_receiver": None,
                "recovery": None,
                "primary": None,
                "server": None,
                "error": {
                    "kind": "connection_refused",
                    "message": message,
                },
            }

        raise

    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

        peer_state = None
        if in_recovery and not streaming:
            peer_state = probe_replication_peers(config, password)

        ha = normalize_ha_state(
            config=config,
            runtime_state=runtime_state,
            in_recovery=in_recovery,
            streaming=streaming,
            receiver=receiver,
            recovery=recovery,
            peer_state=peer_state,
        )

        return {
            "ha": ha,
            "peer_state": peer_state,
            "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()


NOTIFY_FIELDS = (
    "type",
    "operation",
    "promote_uname",
    "demote_uname",
    "promoted_uname",
    "unpromoted_uname",
)


def _notify_value(parameters, field):
    value = parameters.get(field, "")

    if not isinstance(value, str):
        raise ValueError(
            f"notify field {field} must be a string"
        )

    if any(char in value for char in ("\n", "\r", "\x00")):
        raise ValueError(
            f"notify field {field} contains invalid characters"
        )

    return value.strip()


def notify(config, runtime_state, request_parameters):
    """
    Record Pacemaker notification metadata in daemon runtime state.

    The daemon remains the local authority for this resource instance.
    Node-target interpretation is deliberately deferred until the state
    persistence and demotion action are migrated together.
    """
    if not isinstance(request_parameters, dict):
        raise ValueError(
            "notify parameters must be a JSON object"
        )

    state = {
        field: _notify_value(request_parameters, field)
        for field in NOTIFY_FIELDS
    }

    state["received_at"] = time.time()

    postgresql_state = _postgres_runtime_state(runtime_state)
    postgresql_state["notification"] = state

    return {
        "ocf_rc": OCF_SUCCESS,
        "state": dict(state),
    }


NOTIFY_TARGET_MAX_AGE_SECONDS = 120


def notify_target(config, runtime_state, request_parameters):
    """
    Report whether the local daemon is the target of the last notification.

    This is deliberately read-only. Promotion and demotion execution remain
    separate actions until target selection is proven in production.
    """
    del request_parameters

    node_name = str(
        config.get("PACEMAKER_NODE_NAME", "")
    ).strip().lower()

    if not node_name:
        raise ValueError(
            "PACEMAKER_NODE_NAME is required for notify.target"
        )

    notification = _postgres_runtime_state(
        runtime_state
    ).get("notification")

    if not isinstance(notification, dict):
        return {
            "ocf_rc": 1,
            "target": False,
            "reason": "notification_missing",
            "node": node_name,
        }

    received_at = notification.get("received_at")

    if not isinstance(received_at, (int, float)):
        return {
            "ocf_rc": 1,
            "target": False,
            "reason": "notification_timestamp_invalid",
            "node": node_name,
        }

    age = time.time() - received_at

    if age < 0 or age > NOTIFY_TARGET_MAX_AGE_SECONDS:
        return {
            "ocf_rc": 1,
            "target": False,
            "reason": "notification_stale",
            "age_seconds": age,
            "node": node_name,
        }

    operation = notification.get("operation")

    if operation == "demote":
        target_node = notification.get(
            "demote_uname",
            "",
        ).strip().lower()
        successor_node = notification.get(
            "promote_uname",
            "",
        ).strip().lower()
    elif operation == "promote":
        target_node = notification.get(
            "promote_uname",
            "",
        ).strip().lower()
        successor_node = notification.get(
            "demote_uname",
            "",
        ).strip().lower()
    else:
        return {
            "ocf_rc": 1,
            "target": False,
            "reason": "operation_unsupported",
            "operation": operation,
            "node": node_name,
        }

    return {
        "ocf_rc": 0,
        "target": target_node == node_name,
        "operation": operation,
        "target_node": target_node,
        "successor_node": successor_node,
        "node": node_name,
        "age_seconds": age,
    }


def demote_prepare(config, runtime_state, request_parameters):
    """
    Read-only demotion preparation.

    Validates the latest Pacemaker demotion notification against this local
    daemon and reports the successor. PostgreSQL mutation remains outside this
    action until candidate and WAL checks are migrated.
    """
    del request_parameters

    node_name = str(
        config.get("PACEMAKER_NODE_NAME", "")
    ).strip().lower()

    if not node_name:
        raise ValueError(
            "PACEMAKER_NODE_NAME is required for demote.prepare"
        )

    notification = _postgres_runtime_state(
        runtime_state
    ).get("notification")

    if not isinstance(notification, dict):
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "notification_missing",
            "node": node_name,
        }

    if notification.get("type") != "pre":
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "notification_type_invalid",
            "node": node_name,
        }

    if notification.get("operation") != "demote":
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "notification_operation_invalid",
            "node": node_name,
        }

    received_at = notification.get("received_at")

    if not isinstance(received_at, (int, float)):
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "notification_timestamp_invalid",
            "node": node_name,
        }

    age = time.time() - received_at

    if age < 0 or age > NOTIFY_TARGET_MAX_AGE_SECONDS:
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "notification_stale",
            "age_seconds": age,
            "node": node_name,
        }

    target_node = str(
        notification.get("demote_uname", "")
    ).strip().lower()

    successor_node = str(
        notification.get("promote_uname", "")
    ).strip().lower()

    if target_node != node_name:
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "local_node_not_target",
            "node": node_name,
            "target_node": target_node,
            "successor_node": successor_node,
        }

    if not successor_node:
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "successor_missing",
            "node": node_name,
        }

    try:
        peers = parse_replication_peers(
            config.get("PG_REPLICATION_PEERS", "")
        )
    except (TypeError, ValueError):
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "peer_map_invalid",
            "node": node_name,
            "successor_node": successor_node,
        }

    successor_peer = next(
        (
            peer for peer in peers
            if peer["name"].strip().lower() == successor_node
        ),
        None,
    )

    if successor_peer is None:
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "successor_not_configured",
            "node": node_name,
            "successor_node": successor_node,
        }

    return {
        "ocf_rc": 0,
        "ready": False,
        "identity_validated": True,
        "reason": "database_checks_pending",
        "mutation": False,
        "node": node_name,
        "target_node": target_node,
        "successor_node": successor_node,
        "successor_host": successor_peer["host"],
        "successor_port": successor_peer["port"],
        "age_seconds": age,
        "checks": {
            "peer_map": True,
            "notification": True,
            "local_target": True,
            "successor_selected": True,
            "postgresql_mutated": False,
            "database_checks": False,
        },
    }


def _lsn_to_int(value):
    if not value:
        return None

    upper, lower = value.split("/", 1)
    return (int(upper, 16) << 32) + int(lower, 16)


def demote_database(config, runtime_state, request_parameters):
    """
    Read-only PostgreSQL successor validation for demotion.
    """
    del request_parameters

    preparation = demote_prepare(
        config,
        runtime_state,
        {},
    )

    if preparation.get("reason") not in (
        "database_checks_pending",
    ):
        return preparation

    local = monitor(config, runtime_state)

    if local["ha"]["ocf_rc"] != OCF_RUNNING_PROMOTED:
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "local_not_primary",
            "local": local["ha"],
        }

    successor_node = preparation["successor_node"]
    peers = parse_replication_peers(
        config.get("PG_REPLICATION_PEERS", "")
    )

    successor = next(
        (
            peer for peer in peers
            if peer["name"].strip().lower() == successor_node
        ),
        None,
    )

    if successor is None:
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "successor_not_configured",
            "successor_node": successor_node,
        }

    try:
        password = read_secret(config["PGPASSWORD_FILE"])

        conn = psycopg2.connect(
            host=successor["host"],
            port=successor["port"],
            user=config["PGUSER"],
            dbname=config["PGDATABASE"],
            password=password,
            connect_timeout=3,
            application_name="managed-containerd-demote-check",
        )

        try:
            conn.autocommit = True

            with conn.cursor(
                cursor_factory=psycopg2.extras.RealDictCursor
            ) as cur:
                cur.execute("""
                    SELECT
                        pg_is_in_recovery() AS in_recovery,
                        pg_is_wal_replay_paused() AS replay_paused,
                        pg_last_wal_receive_lsn()::text AS receive_lsn,
                        pg_last_wal_replay_lsn()::text AS replay_lsn,
                        (
                            SELECT status
                            FROM pg_stat_wal_receiver
                            LIMIT 1
                        ) AS receiver_status,
                        (
                            SELECT slot_name
                            FROM pg_stat_wal_receiver
                            LIMIT 1
                        ) AS slot_name
                """)
                successor_state = dict(cur.fetchone())
        finally:
            conn.close()

    except psycopg2.Error as exc:
        return {
            "ocf_rc": 1,
            "ready": False,
            "reason": "successor_unreachable",
            "successor_node": successor_node,
            "successor_host": successor["host"],
            "successor_port": successor["port"],
            "error": str(exc).splitlines()[0],
        }

    primary_lsn = local["primary"]["current_wal_lsn"]
    replay_lsn = successor_state["replay_lsn"]

    primary_value = _lsn_to_int(primary_lsn)
    replay_value = _lsn_to_int(replay_lsn)

    wal_lag_bytes = None
    if primary_value is not None and replay_value is not None:
        wal_lag_bytes = max(0, primary_value - replay_value)

    checks = {
        "local_primary": True,
        "successor_reachable": True,
        "successor_in_recovery": (
            successor_state["in_recovery"] is True
        ),
        "successor_streaming": (
            successor_state["receiver_status"] == "streaming"
        ),
        "successor_replay_not_paused": (
            successor_state["replay_paused"] is False
        ),
        "successor_slot_present": bool(
            successor_state["slot_name"]
        ),
    }

    ready = all(checks.values())

    return {
        "ocf_rc": 0 if ready else 1,
        "ready": ready,
        "mutation": False,
        "final_handoff_required": True,
        "node": preparation["node"],
        "target_node": preparation["target_node"],
        "successor_node": successor_node,
        "successor_host": successor["host"],
        "successor_port": successor["port"],
        "primary_wal_lsn": primary_lsn,
        "successor_receive_lsn": successor_state["receive_lsn"],
        "successor_replay_lsn": replay_lsn,
        "wal_lag_bytes": wal_lag_bytes,
        "successor_slot_name": successor_state["slot_name"],
        "checks": checks,
    }