#!/usr/bin/env python3
"""Local HTTPS health -> BIRD static protocol control for the CDN tutorial.

Run under the supplied systemd unit. This controls only cdn_prefix on a
dedicated example edge; it does not measure global reachability or capacity.
"""

import ipaddress
import os
import socket
import subprocess
import time
from dataclasses import dataclass


PROTOCOL = "cdn_prefix"
INTERVAL = 5
FAILURES = 3
HEALTHY_SECONDS = 30
HOLD_DOWN = 60


@dataclass
class Policy:
    withdrawn_at: float
    advertised: bool = False
    failures: int = 0
    healthy_since: float | None = None

    def update(self, healthy, now):
        if healthy:
            self.failures = 0
            if self.healthy_since is None:
                self.healthy_since = now
            if (
                not self.advertised
                and now - self.healthy_since >= HEALTHY_SECONDS
                and now - self.withdrawn_at >= HOLD_DOWN
            ):
                self.advertised = True
        else:
            self.healthy_since = None
            self.failures += 1
            if self.advertised and self.failures >= FAILURES:
                self.advertised = False
                self.withdrawn_at = now
        return self.advertised


def run(args, timeout=5):
    return subprocess.run(
        args, check=True, capture_output=True, text=True, timeout=timeout
    ).stdout


def route_state():
    output = run(["/usr/sbin/birdc", "show", "protocols", PROTOCOL])
    for line in output.splitlines():
        fields = line.split()
        if len(fields) >= 4 and fields[0] == PROTOCOL:
            return fields[3]
    raise RuntimeError("BIRD did not report cdn_prefix")


def set_route(advertise):
    expected = "up" if advertise else "down"
    if route_state() == expected:
        return
    run(["/usr/sbin/birdc", "enable" if advertise else "disable", PROTOCOL])
    if route_state() != expected:
        raise RuntimeError("BIRD did not reach the requested protocol state")
    print("cdn_prefix " + expected, flush=True)


def probe(host, address):
    try:
        body = run([
            "/usr/bin/curl", "--fail", "--silent", "--show-error",
            "--noproxy", "*", "--max-time", "3", "--max-filesize", "64",
            "--resolve", f"{host}:443:{address}",
            f"https://{host}/__edge/health",
        ])
        return body == "edge-ok\n"
    except (subprocess.SubprocessError, OSError):
        return False


def notify(message):
    address = os.environ["NOTIFY_SOCKET"]
    if address.startswith("@"):
        address = "\0" + address[1:]
    with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sock:
        sock.settimeout(2)
        sock.connect(address)
        sock.sendall(message.encode())


def main():
    # Start withdrawn, including after a restart or configuration change.
    set_route(False)
    host = os.environ["CDN_HOST"]
    address = str(ipaddress.IPv4Address(os.environ["CDN_IP"]))
    policy = Policy(withdrawn_at=time.monotonic())
    notify("READY=1")
    try:
        while True:
            healthy = probe(host, address)
            advertise = policy.update(healthy, time.monotonic())
            # Read back BIRD state on every pass, including after BIRD restarts.
            set_route(advertise)
            notify("WATCHDOG=1")
            time.sleep(INTERVAL)
    finally:
        set_route(False)


if __name__ == "__main__":
    main()