#!/usr/bin/env python3
"""Restore this RK3528 terminal's verified HDMI timing after X11 hotplug."""

import argparse
import fcntl
import os
from pathlib import Path
import re
import signal
import subprocess
import sys
import time
from dataclasses import dataclass

OUTPUT = "HDMI-1"
MODE = "2560x1440_60_fixed297"
# The vendor PHY rejects the EDID's 241.70 MHz timing. Keep the verified PLL.
TIMING = "297.00 2560 2608 2640 2720 1440 1443 1448 1820 +HSync -VSync".split()
POLL_SECONDS = 2
SETTLE_SECONDS = 2


def log(message):
    print(time.strftime("%Y-%m-%d %H:%M:%S"), message, flush=True)


def xrandr(*args):
    result = subprocess.run(["xrandr", *args], capture_output=True, text=True,
                            timeout=10, check=False)
    if result.returncode:
        raise RuntimeError((result.stderr or result.stdout).strip()[:400])
    return result.stdout


@dataclass
class DisplayState:
    connected: bool
    geometry: str
    active_mode: str
    modes: set
    defined: bool

    @property
    def healthy(self):
        return (self.connected and self.geometry == "2560x1440+0+0"
                and self.active_mode == MODE)


def parse_state(query):
    connected, geometry, active = False, "", ""
    modes, defined, in_output = set(), False, False
    for line in query.splitlines():
        fields = line.split()
        if not fields:
            continue
        if fields[0] == MODE:
            defined = True  # Includes a globally defined but detached mode.
        if not line[0].isspace():
            in_output = fields[0] == OUTPUT
            if in_output:
                connected = len(fields) > 1 and fields[1] == "connected"
                geometry = next((f for f in fields[2:]
                                 if re.fullmatch(r"\d+x\d+\+\d+\+\d+", f)), "")
        elif in_output and len(fields) > 1 and re.match(r"^\d+(?:\.\d+)?[+*]*$", fields[1]):
            modes.add(fields[0])
            if "*" in " ".join(fields[1:]):
                active = fields[0]
    return DisplayState(connected, geometry, active, modes, defined)


class Guard:
    def __init__(self, runner=xrandr, reporter=log):
        self.run = runner
        self.log = reporter
        self.bad_since = None
        self.next_retry = 0
        self.failures = 0
        self.last_status = None

    def report(self, status):
        if status != self.last_status:
            self.log(status)
            self.last_status = status

    def step(self, now):
        try:
            # Xorg already receives DRM hotplug events. Read its current state
            # without probing every EDID/mode: the vendor driver logs clock
            # warnings for those probes, adding unnecessary work while streaming.
            state = parse_state(self.run("--current"))
        except (OSError, RuntimeError, subprocess.TimeoutExpired) as error:
            self.bad_since = None
            self.report("X11 unavailable; waiting: " + str(error)[:400])
            return False
        if not state.connected:
            self.bad_since = None
            self.next_retry = 0
            self.failures = 0
            self.report(OUTPUT + " disconnected; no mode changes")
            return False
        if state.healthy:
            self.bad_since = None
            self.next_retry = 0
            self.failures = 0
            self.report("healthy: " + MODE)
            return True
        self.report("mode drift: " + (state.active_mode or "off") + " " + state.geometry)
        if self.bad_since is None:
            self.bad_since = now
        if now - self.bad_since < SETTLE_SECONDS or now < self.next_retry:
            return False
        try:
            if not state.defined:
                self.run("--newmode", MODE, *TIMING)
            if MODE not in state.modes:
                self.run("--addmode", OUTPUT, MODE)
            self.run("--output", OUTPUT, "--mode", MODE, "--pos", "0x0")
            if not parse_state(self.run("--current")).healthy:
                raise RuntimeError("mode was not retained after apply")
        except (OSError, RuntimeError, subprocess.TimeoutExpired) as error:
            self.failures = min(self.failures + 1, 6)
            delay = min(60, 2 ** self.failures)
            # Base backoff on completion, including a slow/timed-out modeset.
            self.next_retry = time.monotonic() + delay
            self.log("restore failed; retry in %ds: %s" % (delay, str(error)[:400]))
            return False
        self.bad_since = None
        self.next_retry = 0
        self.failures = 0
        self.report("restored: " + MODE + " (stream process untouched)")
        return True


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--once", action="store_true",
                        help="initial display recovery with a 20s retry window, then exit")
    args = parser.parse_args()
    os.environ.setdefault("DISPLAY", ":0")
    os.environ.setdefault("XAUTHORITY", str(Path.home() / ".Xauthority"))
    runtime = Path(os.environ.get("XDG_RUNTIME_DIR", str(Path.home() / ".cache/vdi-terminal")))
    runtime.mkdir(parents=True, exist_ok=True)
    # Same lock for startup and the systemd service. Never run competing guards.
    with (runtime / "vdi-display-guard.lock").open("a") as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            log("display guard already running")
            return 0
        stop = False

        def request_stop(signum, frame):
            nonlocal stop
            stop = True

        signal.signal(signal.SIGTERM, request_stop)
        signal.signal(signal.SIGINT, request_stop)
        guard = Guard()
        deadline = time.monotonic() + 20
        while not stop:
            if guard.step(time.monotonic()) and args.once:
                return 0
            if args.once and time.monotonic() >= deadline:
                return 1
            time.sleep(POLL_SECONDS)
    return 0


if __name__ == "__main__":
    sys.exit(main())