from __future__ import annotations

import argparse
import json
import os
import platform
import subprocess
import time
import urllib.request
from pathlib import Path
from typing import Any


CANDIDATE_RELEASE = "6.1.31-matrix-axp313a1"
MARKER_PATH = Path("/boot/matrix-kernel-good")
CPUFREQ_ROOT = Path("/sys/devices/system/cpu/cpufreq")


class HealthError(RuntimeError):
    pass


def _read(path: Path) -> str:
    return path.read_text(encoding="ascii", errors="strict").strip()


def check_axp313a_binding(sys_root: Path = Path("/sys")) -> None:
    matches = list((sys_root / "bus/i2c/devices").glob("*-0036/of_node/compatible"))
    if not matches:
        raise HealthError("PMIC device-tree node at I2C address 0x36 is missing")
    compatible = matches[0].read_bytes().replace(b"\0", b"\n").decode("ascii", errors="replace")
    if "x-powers,axp313a" not in compatible.splitlines():
        raise HealthError(f"PMIC compatible is not x-powers,axp313a: {compatible!r}")
    driver_link = matches[0].parents[1] / "driver"
    if not driver_link.exists() or Path(os.path.realpath(driver_link)).name != "axp20x-i2c":
        raise HealthError("AXP313A is not bound to axp20x-i2c")


def check_cpufreq(
    cpufreq_root: Path = CPUFREQ_ROOT,
    debugfs_deferred: Path = Path("/sys/kernel/debug/devices_deferred"),
) -> None:
    policy = cpufreq_root / "policy0"
    if not policy.is_dir():
        raise HealthError("cpufreq policy0 is missing")
    affected = set(_read(policy / "affected_cpus").split())
    if affected != {"0", "1", "2", "3"}:
        raise HealthError(f"policy0 does not cover CPU0-3: {sorted(affected)}")
    available = set(_read(policy / "scaling_available_governors").split())
    if "performance" not in available:
        raise HealthError("performance governor is unavailable")
    if debugfs_deferred.is_file() and "cpufreq-dt" in debugfs_deferred.read_text(
        encoding="utf-8", errors="replace"
    ):
        raise HealthError("cpufreq-dt is still deferred")


def check_cpu_regulator(sys_root: Path = Path("/sys")) -> None:
    regulator_root = sys_root / "class/regulator"
    for directory in regulator_root.glob("regulator.*"):
        name_path = directory / "name"
        if not name_path.is_file() or _read(name_path) != "vdd-cpu":
            continue
        voltage = int(_read(directory / "microvolts"))
        if not 810_000 <= voltage <= 1_080_000:
            raise HealthError(f"vdd-cpu is outside the approved range: {voltage}uV")
        return
    raise HealthError("vdd-cpu regulator is missing")


def check_kernel_messages(messages: str) -> None:
    forbidden = (
        "DCDC frequency on AXP313a is fixed to 3 MHz",
        "Error setting dcdc frequency",
    )
    for message in forbidden:
        if message in messages:
            raise HealthError(f"AXP313A regulator logged a boot error: {message}")


def check_kernel_log() -> None:
    result = subprocess.run(
        ["/bin/dmesg"],
        check=True,
        capture_output=True,
        text=True,
        encoding="utf-8",
        errors="replace",
    )
    check_kernel_messages(result.stdout)


def _fetch_status(url: str) -> dict[str, Any]:
    # The candidate kernel is first validated against the already-installed
    # production application.  That older build can take several seconds to
    # answer while its display compositor is busy, which must not masquerade
    # as a kernel/cpufreq failure.  The stricter 500 ms HTTP performance gate
    # is measured separately after the candidate application is staged.
    with urllib.request.urlopen(url, timeout=10.0) as response:
        return json.loads(response.read().decode("utf-8"))


def check_driver_pair(first: dict[str, Any], second: dict[str, Any]) -> None:
    screen = second.get("screen") or {}
    if screen.get("driver") != "walnutpi-h618-hub75":
        raise HealthError("production H618 driver is not active")
    driver = screen.get("driver_status") or {}
    previous = (first.get("screen") or {}).get("driver_status") or {}
    for field in ("actual_refresh_rate_hz", "panel_scan_rate_hz"):
        if float(driver.get(field) or 0.0) < 95.0:
            raise HealthError(f"{field} is below 95Hz")
    if driver.get("scans_per_frame") != 1:
        raise HealthError("scans_per_frame is not 1")
    frame_delta = int(driver.get("completed_frames") or 0) - int(previous.get("completed_frames") or 0)
    scan_delta = int(driver.get("completed_scans") or 0) - int(previous.get("completed_scans") or 0)
    miss_delta = int(driver.get("deadline_misses") or 0) - int(previous.get("deadline_misses") or 0)
    if frame_delta <= 0 or scan_delta != frame_delta:
        raise HealthError("HUB75 completed frame/scan counters are stalled or unequal")
    if miss_delta < 0 or miss_delta / frame_delta > 0.001:
        raise HealthError("HUB75 deadline miss rate exceeds 0.1%")
    for field in ("oe_pulse_faults", "oe_forced_blanks"):
        if int(driver.get(field) or 0) != int(previous.get(field) or 0):
            raise HealthError(f"{field} increased during health check")
    for field in ("driver_error", "oe_timing_error", "safeoff_error"):
        if driver.get(field):
            raise HealthError(f"{field} is not empty: {driver[field]}")
    if driver.get("cpu_affinity") != 3 or driver.get("cpu_affinity_active") is not True:
        raise HealthError("CPU3 affinity is not active")
    if driver.get("realtime_priority") != 50 or driver.get("realtime_priority_active") is not True:
        raise HealthError("SCHED_FIFO 50 is not active")
    if driver.get("memory_locked") is not True:
        raise HealthError("refresh memory is not locked")


def wait_for_driver(url: str, timeout_seconds: int = 60) -> None:
    deadline = time.monotonic() + timeout_seconds
    last_error: Exception | None = None
    while time.monotonic() < deadline:
        try:
            first = _fetch_status(url)
            time.sleep(2)
            second = _fetch_status(url)
            check_driver_pair(first, second)
            return
        except Exception as exc:
            last_error = exc
            time.sleep(1)
    raise HealthError(f"application/GPIO health timed out: {last_error}")


def write_marker(marker: Path = MARKER_PATH) -> None:
    temporary = marker.with_name(f".{marker.name}.tmp")
    with temporary.open("w", encoding="ascii", newline="\n") as handle:
        handle.write(f"{CANDIDATE_RELEASE}\n")
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temporary, marker)
    os.sync()


def main() -> int:
    parser = argparse.ArgumentParser(description="Confirm the one-shot WalnutPi AXP313A kernel")
    parser.add_argument("--status-url", default="http://127.0.0.1:8080/api/status")
    parser.add_argument("--marker", type=Path, default=MARKER_PATH)
    args = parser.parse_args()
    if platform.release() != CANDIDATE_RELEASE:
        raise HealthError(f"unexpected candidate release: {platform.release()}")
    check_axp313a_binding()
    check_cpufreq()
    check_cpu_regulator()
    check_kernel_log()
    wait_for_driver(args.status_url)
    write_marker(args.marker)
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except HealthError as exc:
        print(f"AXP313A candidate health failed: {exc}", flush=True)
        raise SystemExit(1)