#!/usr/bin/env python3
"""
WML OS Installer - OPNsense-style numbered menu installer.

Runs as PID 1's child on tty1 when booting the live USB.  Same look-and-feel
as wml-console (post-install menu) so the product feels consistent.

Architecture (lifted from VyOS image_installer.py):
  - sgdisk hybrid GPT layout (bios_grub + ESP + ext4 root) in ONE call
  - partx -u to refresh kernel partition table (not partprobe - it races)
  - unsquashfs to extract live image (not file-copy)
  - mount by LABEL so fstab survives reboots
  - chroot chpasswd for root password
  - write /etc/wml/admin.json so web UI uses same credential
  - grub-install for both BIOS (i386-pc) and UEFI (x86_64-efi --force-extra-removable)
  - efibootmgr to register NVRAM entry
  - hardchecks for partition device-nodes existence, write-test on target/etc

Strict input validation everywhere.  All confirms use [y/N], never literal YES.
"""
import os
import sys
import json
import time
import shutil
import signal
import subprocess

sys.path.insert(0, "/usr/local/lib/wml")
from tui import (
    header, hr, section, menu, prompt, prompt_int, prompt_yn,
    prompt_password, prompt_hostname, pause, transition,
    Spinner, run_status, cols, BOLD, DIM, PINK, WHITE, YELLOW, RED,
    GREEN, PURPLE, RESET
)


TARGET = "/mnt/installation/disk_dst"
IMAGE_NAME = "wml-os"
IMAGE_VERSION = "1.0.0"
LOG_FILE = "/tmp/wml-install.log"


def log(msg):
    try:
        with open(LOG_FILE, "a") as f:
            f.write(msg + "\n")
    except Exception:
        pass


def run(cmd, capture=False, check=True):
    # ALWAYS capture so sub-command stdout/stderr never leaks red text to the
    # clean install screen. Everything goes to the log; callers that asked for
    # capture=True still get .stdout/.stderr. (Fixes the transient red line
    # flashing during partition/grub steps.)
    log(f"$ {' '.join(cmd) if isinstance(cmd, list) else cmd}")
    if isinstance(cmd, list):
        r = subprocess.run(cmd, capture_output=True, text=True)
    else:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if r.stdout:
        log(f"  stdout: {r.stdout[:500]}")
    if r.stderr:
        log(f"  stderr: {r.stderr[:500]}")
    log(f"  -> exit {r.returncode}")
    if check and r.returncode != 0:
        raise RuntimeError(f"command failed (exit {r.returncode}): {cmd}")
    return r


# ---------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------
def show_header():
    arch = subprocess.check_output(["uname", "-m"], text=True).strip()
    header(f"Installer: WML OS {IMAGE_VERSION} ({arch})")
    body = []
    try:
        for line in open("/proc/cpuinfo"):
            if line.startswith("model name"):
                cpu = line.split(":", 1)[1].strip()
                body.append(f"CPU      : {cpu}")
                break
    except Exception:
        pass
    try:
        # Prefer dmidecode for true physical DIMM size; fallback to /proc/meminfo
        # rounded to nearest GB (kernel reserves ~1-2% so simple truncation underreports).
        ram_gb = None
        try:
            r = subprocess.run(["dmidecode", "-t", "memory"],
                               capture_output=True, text=True, timeout=5)
            total_mb = 0
            for line in r.stdout.split("\n"):
                line = line.strip()
                if line.startswith("Size:") and "No Module Installed" not in line:
                    parts = line.split()
                    if len(parts) >= 3:
                        try:
                            val = int(parts[1])
                            unit = parts[2].upper()
                            if unit == "GB":
                                total_mb += val * 1024
                            elif unit == "MB":
                                total_mb += val
                            elif unit == "TB":
                                total_mb += val * 1024 * 1024
                        except ValueError:
                            pass
            if total_mb > 0:
                ram_gb = total_mb // 1024
        except Exception:
            pass
        if ram_gb is None:
            with open("/proc/meminfo") as f:
                for line in f:
                    if line.startswith("MemTotal:"):
                        kb = int(line.split()[1])
                        # Round to nearest GB (account for kernel reservations)
                        ram_gb = (kb + 512 * 1024) // (1024 * 1024)
                        break
        if ram_gb is not None:
            body.append(f"RAM      : {ram_gb} GB")
    except Exception:
        pass
    # NIC list
    nics = list_nics()
    body.append(f"NICs     : {', '.join(n[0] for n in nics) if nics else 'none'}")
    section("Detected hardware", body)


# ---------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------
def list_disks():
    """List installable disks (excludes USB/live medium itself)."""
    result = run(["lsblk", "-Jbp"], capture=True, check=False)
    data = json.loads(result.stdout)
    out = []
    live_dev = ""
    try:
        live_dev = subprocess.check_output(
            ["findmnt", "-n", "-o", "SOURCE", "/run/live/medium"],
            text=True).strip()
        live_dev = live_dev.rstrip("0123456789").rstrip("p")
    except Exception:
        pass
    for d in data.get("blockdevices", []):
        if d.get("type") != "disk":
            continue
        name = d["name"]
        if name.startswith("/dev/loop") or name.startswith("/dev/sr"):
            continue
        if name == live_dev:
            continue
        size = int(d.get("size", 0))
        if size < 4 * 1024**3:
            continue
        model = (d.get("model") or "(unknown)").strip()
        out.append((name, size, model))
    return out

def list_nics():
    """List physical NICs."""
    out = []
    for d in sorted(os.listdir("/sys/class/net")):
        if d in ("lo",) or d.startswith(("docker", "veth", "virbr", "tap", "wg", "br-", "tun")):
            continue
        if os.path.exists(f"/sys/class/net/{d}/device"):
            out.append((d,))
    return out

def fmt_size(n):
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024:
            return f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} PB"


# ---------------------------------------------------------------------
# 3-step install flow
# ---------------------------------------------------------------------
def pick_disk():
    transition()
    header("Install · Step 1 of 3")
    disks = list_disks()
    if not disks:
        print(f"  {RED}No installable disks found.{RESET}")
        pause()
        return None
    section("Available disks",
            [f"{BOLD}{i}){RESET} {p}  {fmt_size(s):>10}  {m}"
             for i, (p, s, m) in enumerate(disks, 1)])
    print()
    pick = prompt_int("Pick disk", 1, len(disks))
    return disks[pick - 1][0]


def pick_hostname():
    transition()
    header("Install · Step 2 of 3")
    section("Hostname",
            ["The hostname identifies this firewall on your network.",
             f"Default: {BOLD}wml{RESET}"])
    print()
    return prompt_hostname(default="wml")


def pick_password():
    transition()
    header("Install · Step 3 of 3")
    section("Root password",
            ["This password works for:",
             f"  • Console login as {BOLD}root{RESET}",
             f"  • SSH login as {BOLD}root{RESET}",
             f"  • Web UI at {BOLD}https://172.30.30.1/{RESET}",
             "",
             "Minimum 8 characters."])
    print()
    return prompt_password("Password", min_len=8)


def confirm_install(disk, hostname, password):
    transition()
    header("Install · Confirm")
    body = [
        f"{RED}WARNING:{RESET} the disk below will be COMPLETELY ERASED.",
        "",
        f"  {BOLD}Disk     :{RESET} {disk}",
        f"  {BOLD}Hostname :{RESET} {hostname}",
        f"  {BOLD}Password :{RESET} (set, {len(password)} chars)",
    ]
    section("Review", body)
    print()
    return prompt_yn("Proceed with install", default=False)


# ---------------------------------------------------------------------
# Install backend
# ---------------------------------------------------------------------
def do_install(disk, hostname, password):
    transition()
    header("Installing")
    print()
    print(f"  Installing WML OS to {BOLD}{disk}{RESET}")
    print(f"  Log file: {LOG_FILE}")
    print()

    steps = [
        ("Partitioning disk",      lambda: step_partition(disk)),
        ("Formatting partitions",  lambda: step_format(disk)),
        ("Mounting target",        lambda: step_mount(disk)),
        ("Extracting system image", step_extract),
        ("Configuring system",     lambda: step_config(hostname, password)),
        ("Installing bootloader",  lambda: step_grub(disk)),
        ("Finalizing",             step_finalize),
    ]
    for i, (label, fn) in enumerate(steps, 1):
        full_label = f"[{i}/{len(steps)}] {label}"
        try:
            with Spinner(full_label):
                fn()
            print(f"  {GREEN}✓{RESET} {full_label}")
        except Exception as e:
            print(f"  {RED}✗{RESET} {full_label}")
            print(f"  {RED}{e}{RESET}")
            print(f"\n  See {LOG_FILE} for details.")
            return False
    return True


def step_partition(disk):
    run(["wipefs", "-af", disk], check=False)
    # Wipe any old partitions too
    r = run(["lsblk", "-no", "PATH", disk], capture=True, check=False)
    for line in r.stdout.strip().split("\n"):
        if line and line != disk:
            run(["wipefs", "-af", line], check=False)
    run(["sgdisk", "-Z", disk], check=False)
    run(["sgdisk", "-a1",
         "-n1:2048:4095",   "-t1:EF02", "-c1:BIOS boot",
         "-n2:4096:+256M",  "-t2:EF00", "-c2:EFI System",
         "-n3:0:0",         "-t3:8300", "-c3:WML root",
         disk])
    run(["sync"])
    run(["partx", "-u", disk], check=False)
    run(["blockdev", "--rereadpt", disk], check=False)
    run(["udevadm", "settle", "--timeout=15"], check=False)
    time.sleep(3)
    efi_part, root_part = find_partitions(disk)
    if not (os.path.exists(efi_part) and os.path.exists(root_part)):
        raise RuntimeError(f"partition device nodes did not appear: {efi_part} / {root_part}")


def find_partitions(disk):
    """After sgdisk, find the EFI + root partition device nodes."""
    r = run(["lsblk", "-Jp", "-o", "PATH,PARTTYPE", disk], capture=True)
    data = json.loads(r.stdout)
    parts = []
    for d in data.get("blockdevices", []):
        if d.get("parttype"):
            parts.append(d)
        for child in d.get("children", []) or []:
            if child.get("parttype"):
                parts.append(child)
    efi = root = None
    for p in parts:
        pt = (p.get("parttype") or "").lower()
        if pt.startswith("c12a7328"):  # EFI System
            efi = p["path"]
        elif pt.startswith("21686148"):  # BIOS boot
            continue
        else:
            root = p["path"]
    return efi, root


def step_format(disk):
    efi, root = find_partitions(disk)
    run(["mkfs.vfat", "-F32", "-n", "EFI", efi])
    run(["mkfs.ext4", "-F", "-L", "wml-root", root])
    run(["udevadm", "settle", "--timeout=15"], check=False)


def step_mount(disk):
    efi, root = find_partitions(disk)
    os.makedirs(TARGET, exist_ok=True)
    run(["mount", root, TARGET])
    os.makedirs(f"{TARGET}/boot/efi", exist_ok=True)
    run(["mount", efi, f"{TARGET}/boot/efi"])


def step_extract():
    src = None
    for c in ("/usr/lib/live/mount/medium/live/filesystem.squashfs",
              "/run/live/medium/live/filesystem.squashfs"):
        if os.path.exists(c):
            src = c
            break
    if not src:
        raise RuntimeError("live squashfs not found")
    run(["unsquashfs", "-f", "-d", TARGET, src])
    run(["sync"])
    # Writability test
    test = f"{TARGET}/.wml-install-write-test"
    try:
        with open(test, "w") as f:
            f.write("ok")
        os.remove(test)
    except OSError as e:
        raise RuntimeError(f"target not writable: {e}")
    # Copy kernel + initrd to /boot for GRUB to find
    os.makedirs(f"{TARGET}/boot", exist_ok=True)
    for f in ("vmlinuz", "initrd.img"):
        for src_dir in ("/run/live/medium/live",
                        "/usr/lib/live/mount/medium/live"):
            p = f"{src_dir}/{f}"
            if os.path.exists(p):
                shutil.copy(p, f"{TARGET}/boot/{f}")
                break
    run(["sync"])


def step_config(hostname, password):
    # Hostname
    with open(f"{TARGET}/etc/hostname", "w") as f:
        f.write(f"{hostname}\n")
    # hosts
    try:
        with open(f"{TARGET}/etc/hosts") as f:
            lines = [l for l in f.readlines() if not l.startswith("127.0.1.1")]
    except Exception:
        lines = ["127.0.0.1\tlocalhost\n"]
    lines.append(f"127.0.1.1\t{hostname}\n")
    with open(f"{TARGET}/etc/hosts", "w") as f:
        f.writelines(lines)
    # fstab by label
    with open(f"{TARGET}/etc/fstab", "w") as f:
        f.write("# WML OS\n"
                "LABEL=wml-root  /         ext4  errors=remount-ro 0 1\n"
                "LABEL=EFI       /boot/efi vfat  umask=0077        0 2\n"
                "tmpfs           /tmp      tmpfs defaults,nosuid   0 0\n")
    # Root password: hash with openssl ONCE, then write the same hash to
    # BOTH the target /etc/shadow (for console + SSH login) AND
    # /etc/wml/admin.json (for web UI login). Single source of truth.
    # We don't use chroot chpasswd because subprocess stdin piping into
    # chroot has subtle bugs - direct file write is bulletproof.
    h = subprocess.check_output(
        ["openssl", "passwd", "-6", password], text=True).strip()
    days_since_epoch = int(time.time() / 86400)

    # Rewrite /etc/shadow with the new root hash. Format:
    #   root:$6$salt$hash:<days_since_epoch>:0:99999:7:::
    shadow_path = f"{TARGET}/etc/shadow"
    new_shadow = []
    root_seen = False
    try:
        with open(shadow_path) as f:
            for line in f:
                if line.startswith("root:"):
                    new_shadow.append(
                        f"root:{h}:{days_since_epoch}:0:99999:7:::\n")
                    root_seen = True
                else:
                    new_shadow.append(line)
    except Exception:
        pass
    if not root_seen:
        new_shadow.append(f"root:{h}:{days_since_epoch}:0:99999:7:::\n")
    with open(shadow_path, "w") as f:
        f.writelines(new_shadow)
    os.chmod(shadow_path, 0o640)
    log(f"Wrote root password hash to {shadow_path} (len={len(h)})")

    # admin.json for web UI single-credential
    os.makedirs(f"{TARGET}/etc/wml", exist_ok=True)
    with open(f"{TARGET}/etc/wml/admin.json", "w") as f:
        json.dump({
            "username": "root",
            "password_hash": h,
            "role": "superadmin",
            "must_change_password": False,
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
        }, f, indent=2)
    os.chmod(f"{TARGET}/etc/wml/admin.json", 0o600)
    # /etc/wml/installed marker
    open(f"{TARGET}/etc/wml/installed", "w").close()

    # SSH config: ensure PermitRootLogin + PasswordAuthentication so user
    # can SSH with the same root password they set during install.
    # Debian's sshd_config defaults to "PermitRootLogin prohibit-password"
    # which BANS password login as root - bad UX for an appliance.
    sshd_dropin_dir = f"{TARGET}/etc/ssh/sshd_config.d"
    os.makedirs(sshd_dropin_dir, exist_ok=True)
    with open(f"{sshd_dropin_dir}/10-wml.conf", "w") as f:
        f.write("# WML OS - allow root password login for appliance admin\n"
                "PermitRootLogin yes\n"
                "PasswordAuthentication yes\n"
                "PubkeyAuthentication yes\n")
    log("Wrote SSH config drop-in: root password login enabled")
    # Strip live-only artifacts
    for p in (f"{TARGET}/root/.bash_profile",
              f"{TARGET}/etc/systemd/system/getty@tty1.service.d/plymouth.conf"):
        if os.path.exists(p):
            os.remove(p)
    # tty1 autologin -> wml-console
    autologin_dir = f"{TARGET}/etc/systemd/system/getty@tty1.service.d"
    os.makedirs(autologin_dir, exist_ok=True)
    with open(f"{autologin_dir}/autologin.conf", "w") as f:
        f.write("[Service]\n"
                "ExecStart=\n"
                "ExecStart=-/sbin/agetty --autologin root --noclear %I $TERM\n"
                "Type=idle\n")
    # /root/.bash_profile exec wml-console
    with open(f"{TARGET}/root/.bash_profile", "w") as f:
        f.write('if [ -t 0 ]; then exec /usr/local/sbin/wml-console; fi\n')
    # Enable WML services on target
    # Minimal essential services - everything else stays disabled until the
    # admin opts in via the web UI. This is the OPNsense model: lean boot,
    # opt-in to heavy services.
    # Unmask + enable firewall services that were disabled in the live chroot
    # so they're inactive during install but active on the booted target.
    for svc in ("suricata", "clamav-daemon", "clamav-freshclam", "c-icap", "squid",
                "kea-dhcp4-server", "unbound", "nftables", "ssh", "apparmor"):
        run(["chroot", TARGET, "systemctl", "unmask", svc], check=False)
        run(["chroot", TARGET, "systemctl", "enable", svc], check=False)

    for svc in ("wml-firstboot", "wml-api", "wml-wan-monitor", "wml-carrier-watcher", "nginx", "chrony"):
        run(["chroot", TARGET, "systemctl", "enable", f"{svc}.service"],
            check=False)
    # NOTE: wml-network-ensure and wml-carrier-watcher are NOT enabled at
    # install. They caused cascade-failures with systemd-networkd on first
    # boot. Admin can enable via console option 11 (Reload all services)
    # or by running `systemctl enable wml-network-ensure wml-carrier-watcher`.
    # Filter timers also disabled - admin opts in via web UI.
    # Ensure nginx vhost
    target_vhost = f"{TARGET}/etc/nginx/sites-available/wml-admin"
    if os.path.exists(target_vhost):
        try:
            os.remove(f"{TARGET}/etc/nginx/sites-enabled/default")
        except FileNotFoundError:
            pass
        en = f"{TARGET}/etc/nginx/sites-enabled/wml-admin"
        if not os.path.exists(en):
            os.symlink("/etc/nginx/sites-available/wml-admin", en)
    # NOTE: previously masked services here; now we unmask+enable above so
    # wml-firstboot can start them. Firewall services gate on their own config
    # files (kea, squid, suricata) and fail-fast if not configured.


def step_grub(disk):
    boot = f"{TARGET}/boot"
    efi = f"{TARGET}/boot/efi"

    # CRITICAL: bind-mount /dev /proc /sys into the target BEFORE chrooting
    # to run grub-install.  Without /dev, grub-install can't enumerate disks
    # and fails silently — leaves disk without bootloader → UEFI shell on
    # next boot.  Same for /proc (needed for udev queries) and /sys.
    for d in ("proc", "sys", "dev", "dev/pts"):
        os.makedirs(f"{TARGET}/{d}", exist_ok=True)
        run(["mount", "--bind", f"/{d}", f"{TARGET}/{d}"], check=False)
    # On UEFI systems also bind efivars so efibootmgr can write NVRAM entry.
    target_efivars = f"{TARGET}/sys/firmware/efi/efivars"
    if os.path.isdir("/sys/firmware/efi/efivars"):
        os.makedirs(target_efivars, exist_ok=True)
        run(["mount", "--bind", "/sys/firmware/efi/efivars", target_efivars],
            check=False)

    bios_ok = False
    efi_ok = False

    # BIOS install (always attempt - hybrid disk works on legacy and UEFI)
    r = subprocess.run(["chroot", TARGET, "grub-install",
                        "--target=i386-pc", "--recheck", "--force", disk],
                       capture_output=True, text=True)
    bios_ok = (r.returncode == 0)
    log(f"grub-install BIOS rc={r.returncode}: stdout={r.stdout[:200]} stderr={r.stderr[:200]}")

    # UEFI install (only if firmware reports EFI mode)
    if os.path.isdir("/sys/firmware/efi"):
        # Named entry
        r = subprocess.run(["chroot", TARGET, "grub-install",
                            "--target=x86_64-efi",
                            "--efi-directory=/boot/efi",
                            "--bootloader-id=wml-os",
                            "--recheck"],
                           capture_output=True, text=True)
        efi_ok = (r.returncode == 0)
        log(f"grub-install EFI named rc={r.returncode}: stdout={r.stdout[:200]} stderr={r.stderr[:200]}")

        # Removable-media fallback (writes /EFI/BOOT/BOOTX64.EFI for boxes
        # without persistent NVRAM)
        r2 = subprocess.run(["chroot", TARGET, "grub-install",
                             "--target=x86_64-efi",
                             "--efi-directory=/boot/efi",
                             "--bootloader-id=wml-os",
                             "--removable", "--recheck"],
                            capture_output=True, text=True)
        log(f"grub-install EFI removable rc={r2.returncode}: stderr={r2.stderr[:200]}")
        if r2.returncode == 0:
            efi_ok = True

    # Unmount the bind mounts we set up (in reverse order)
    if os.path.ismount(target_efivars):
        run(["umount", target_efivars], check=False)
    for d in ("dev/pts", "dev", "sys", "proc"):
        run(["umount", "-lf", f"{TARGET}/{d}"], check=False)

    if not bios_ok and not efi_ok:
        raise RuntimeError("both BIOS and UEFI grub-install failed - see /tmp/wml-install.log")

    # Write grub.cfg
    cfg = f"""# WML OS GRUB config
set default=0
set timeout=3
insmod gzio
insmod part_gpt
insmod ext2
insmod fat
if loadfont /boot/grub/fonts/unicode.pf2 ; then
    insmod efi_gop
    insmod efi_uga
    insmod gfxterm
    set gfxmode=auto
    terminal_output gfxterm
fi
set color_normal=light-gray/black
set color_highlight=white/black

menuentry "WML OS {IMAGE_VERSION}" {{
    linux  /boot/vmlinuz  root=LABEL=wml-root ro quiet splash net.ifnames=0 biosdevname=0
    initrd /boot/initrd.img
}}

menuentry "WML OS (safe mode)" {{
    linux  /boot/vmlinuz  root=LABEL=wml-root ro nomodeset nosplash verbose
    initrd /boot/initrd.img
}}
"""
    with open(f"{boot}/grub/grub.cfg", "w") as f:
        f.write(cfg)


def step_finalize():
    # Copy the install log + a result summary to the TARGET system so it
    # survives reboot and can be inspected post-install.
    try:
        os.makedirs(f"{TARGET}/var/log/wml", exist_ok=True)
        if os.path.exists(LOG_FILE):
            import shutil
            shutil.copy(LOG_FILE, f"{TARGET}/var/log/wml/install.log")
        with open(f"{TARGET}/etc/wml/last-install-result.txt", "w") as f:
            f.write("Install completed successfully.\n"
                    "Detailed log: /var/log/wml/install.log\n")
    except Exception as e:
        log(f"WARN: couldn't persist install log: {e}")

    run(["sync"])
    for d in ("dev/pts", "dev", "proc", "sys", "run"):
        run(["umount", "-lf", f"{TARGET}/{d}"], check=False)
    run(["umount", "-lf", f"{TARGET}/boot/efi"], check=False)
    run(["umount", "-lf", TARGET], check=False)
    run(["sync"])


# ---------------------------------------------------------------------
# Final screen
# ---------------------------------------------------------------------
def install_complete(disk):
    transition()
    header("Install Complete")
    section("Success",
            [f"WML OS has been installed to {BOLD}{disk}{RESET}.",
             "",
             f"After reboot you can access the admin UI at:",
             f"  {BOLD}https://172.30.30.1/{RESET} (from a LAN-connected laptop)",
             "",
             f"Login: {BOLD}root{RESET} + the password you set."])
    print()
    print(f"  {YELLOW}>>> Remove the install media (USB) before rebooting <<<{RESET}")
    print()
    input(f"  {DIM}Press ENTER to reboot...{RESET}")
    run(["sync"])
    # Kernel-level reboot to avoid systemd "shutdown binary not found"
    try:
        with open("/proc/sys/kernel/sysrq", "w") as f:
            f.write("1\n")
        with open("/proc/sysrq-trigger", "w") as f:
            f.write("b\n")
    except Exception:
        subprocess.run(["reboot", "-f"])


# ---------------------------------------------------------------------
# Main menu
# ---------------------------------------------------------------------
INSTALLER_MENU = [
    (1, "Install WML OS to disk"),
    (2, "Shell (recovery / advanced)"),
    (3, "Reboot"),
    (4, "Power off"),
]


def main_menu():
    # Don't let the user Ctrl-C out of the installer.
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    signal.signal(signal.SIGTSTP, signal.SIG_IGN)

    while True:
        show_header()
        menu(INSTALLER_MENU)
        try:
            choice = prompt_int("Enter an option", 1, 4)
        except (EOFError, KeyboardInterrupt):
            continue

        if choice == 1:
            run_install_flow()
        elif choice == 2:
            transition()
            header("Recovery Shell")
            print(f"  {DIM}Type 'exit' to return to the installer menu.{RESET}")
            print()
            subprocess.run(["/bin/bash", "--norc", "--noprofile", "-i"])
        elif choice == 3:
            if prompt_yn("Reboot now", default=False):
                subprocess.run(["systemctl", "reboot", "--force", "--force"])
        elif choice == 4:
            if prompt_yn("Power off now", default=False):
                subprocess.run(["systemctl", "poweroff", "--force", "--force"])


def run_install_flow():
    try:
        disk = pick_disk()
        if not disk:
            return
        hostname = pick_hostname()
        password = pick_password()
        if not confirm_install(disk, hostname, password):
            print(f"  {DIM}Cancelled.{RESET}")
            pause()
            return
        if do_install(disk, hostname, password):
            install_complete(disk)
    except (EOFError, KeyboardInterrupt):
        print(f"\n  {DIM}Cancelled.  Returning to menu.{RESET}")
        time.sleep(0.5)
    except Exception as e:
        import traceback
        print(f"\n  {RED}Installer crashed:{RESET} {e}")
        traceback.print_exc()
        print(f"\n  Drop to recovery shell.")
        pause()
        subprocess.run(["/bin/bash", "--norc", "--noprofile", "-i"])


if __name__ == "__main__":
    if os.geteuid() != 0:
        print(f"{RED}Must run as root.{RESET}")
        sys.exit(1)
    main_menu()
