#!/usr/bin/env python3
"""bty live env: flash-on-boot.

Reads ``bty.*`` parameters from ``/proc/cmdline``, fetches the
assigned image from the bty server, runs ``bty flash``, signals
completion, and reboots.

Driven by ``bty-flash-on-boot.service`` after ``network-online.target``.
If no ``bty.*`` parameters are present in /proc/cmdline (e.g. the
operator is booting the live env interactively for debugging), the
service exits 0 and the live env drops to its console.

Recognised cmdline parameters:
- ``bty.server=URL``       — base URL of the bty-web server
- ``bty.mac=MAC``          — this machine's MAC (xx-xx-...)
- ``bty.image_url=URL``    — direct URL to the image to flash
- ``bty.provisioning=MODE``— 'none' / 'cloud-init' / 'cijoe'
                             (defaults to 'none' if absent)
"""

from __future__ import annotations

import json
import shlex
import subprocess
import sys
import urllib.request
from pathlib import Path

CMDLINE = Path("/proc/cmdline")
LOCAL_IMAGE = Path("/var/tmp/bty-flash-on-boot.image")

REQUIRED = ("bty.server", "bty.mac", "bty.image_url")


def cmdline_args() -> dict[str, str]:
    """Parse /proc/cmdline into a dict of ``bty.*`` -> value tokens."""
    raw = CMDLINE.read_text()
    out: dict[str, str] = {}
    for token in shlex.split(raw):
        if token.startswith("bty.") and "=" in token:
            k, _, v = token.partition("=")
            out[k] = v
    return out


def pick_target() -> str:
    """Phase D-2 placeholder: first non-removable, non-readonly disk.

    The bty server's per-MAC plan will supply an explicit target hint
    in Phase D-3; until then the live env picks whatever the kernel
    enumerates first. For a bare-metal target with one fixed disk
    that's the right answer; multi-disk VMs need the D-3 hint.
    """
    out = subprocess.run(
        ["bty", "list", "disks", "--json"],
        check=True,
        capture_output=True,
        text=True,
    )
    payload = json.loads(out.stdout)
    for disk in payload.get("disks", []):
        if disk.get("removable", False):
            continue
        if disk.get("readonly", False):
            continue
        return disk["path"]
    raise SystemExit("bty-flash-on-boot: no flashable disk found")


def download(url: str, dest: Path) -> None:
    print(f"bty-flash-on-boot: downloading {url} -> {dest}", flush=True)
    dest.parent.mkdir(parents=True, exist_ok=True)
    with urllib.request.urlopen(url) as resp, dest.open("wb") as f:
        while True:
            chunk = resp.read(1 << 20)
            if not chunk:
                break
            f.write(chunk)


def signal_done(server: str, mac: str) -> None:
    """Best-effort completion signal. The endpoint lands in D-3."""
    url = f"{server.rstrip('/')}/pxe/{mac}/done"
    try:
        req = urllib.request.Request(url, data=b"", method="POST")
        with urllib.request.urlopen(req, timeout=10) as _:
            print(f"bty-flash-on-boot: completion signalled to {url}", flush=True)
    except Exception as exc:  # noqa: BLE001
        print(
            f"bty-flash-on-boot: completion signal to {url} failed: {exc}; "
            "continuing to reboot anyway",
            file=sys.stderr,
            flush=True,
        )


def main() -> int:
    args = cmdline_args()
    missing = [k for k in REQUIRED if k not in args]
    if missing:
        print(
            "bty-flash-on-boot: cmdline missing required keys "
            f"({', '.join(missing)}); not flashing - dropping to console",
            file=sys.stderr,
            flush=True,
        )
        return 0  # not an error, just no work

    server = args["bty.server"].rstrip("/")
    mac = args["bty.mac"]
    image_url = args["bty.image_url"]
    provisioning = args.get("bty.provisioning", "none")

    print(
        f"bty-flash-on-boot: server={server} mac={mac} provisioning={provisioning}",
        flush=True,
    )

    download(image_url, LOCAL_IMAGE)
    target = pick_target()
    print(f"bty-flash-on-boot: target disk {target}", flush=True)

    subprocess.run(
        [
            "bty",
            "flash",
            "--image",
            str(LOCAL_IMAGE),
            "--target",
            target,
            "--provision",
            provisioning,
            "--yes",
        ],
        check=True,
    )

    signal_done(server, mac)

    print("bty-flash-on-boot: flash complete; rebooting in 5s", flush=True)
    subprocess.run(["sleep", "5"], check=False)
    subprocess.run(["systemctl", "reboot"], check=False)
    return 0


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