// Integration test: aarch64 network PXE boot.
//
// The boot-ISO path already covered ARM (arm-iso-provision.test.ts). This covers the
// network path: DHCP option 93 handing an arm64 client an arm64 iPXE binary, dispatch
// serving an aarch64 kernel, and `provision debug` reaching a rescue shell -- which is
// what the DGX Sparks actually need and could not do.
//
// Two suites, because they cost very different amounts of time:
//
//   "ARM PXE rescue"   NBP handoff -> discovery -> rescue with SSH.       ~20-30 min
//   "ARM PXE install"  full discover -> install -> installed.             ~60-90 min
//
// The install suite only runs with ARM_PXE_FULL=1. No ARM machine in the lab is ever
// PXE-installed except the MS-R1, and an hour-plus test that runs by default is a test
// nobody runs.
//
// IMPORTANT: aarch64 has no KVM on an x86_64 host, so all of this is emulated and
// roughly 10x slower than native.
//
// A note for whoever debugs a failure here: if the VM panics with
//   VFS: Unable to mount root fs on unknown-block(0,0)
// that is very likely iPXE silently dropping the initrd because the build lacks
// EFI_LOAD_FILE2_PROTOCOL -- on arm64 the kernel EFI stub fetches the initrd over
// LoadFile2, and an iPXE without it accepts the `initrd` line and does nothing. It is
// NOT a reproduction of the DGX Spark kernel bug that motivated this work, despite
// being the identical message. assertIpxeSupportsLoadFile2() below checks the build up
// front so that failure names itself; to check by hand:
//   node -e 'const b=require("fs").readFileSync("/usr/share/ipxe/arm64-efi/snponly.efi");
//            console.log(b.indexOf(Buffer.from("c1c00640b3fc3e40996d4a6c8724e06d","hex")))'
// Fedora's ipxe-bootimgs-aarch64-20240119 has it at 0x3bbf0.
//
// Prerequisites:
//   - qemu-system-aarch64        (sudo dnf install qemu-system-aarch64)
//   - edk2-aarch64               (sudo dnf install edk2-aarch64)
//   - ipxe-bootimgs-aarch64      (sudo dnf install ipxe-bootimgs-aarch64)
//   - libvirtd, sudo, internet access
//
// Run: sudo ./scripts/test-provision.sh arm-pxe

import { describe, it, expect, beforeAll, afterAll } from "vitest";
import { readFileSync, existsSync, mkdirSync, rmSync, copyFileSync, writeFileSync } from "node:fs";
import { execSync } from "node:child_process";
import { join } from "node:path";
import { homedir, tmpdir } from "node:os";
import { log, waitForSsh } from "./helpers/libvirt.js";
import { ensurePxeNetwork, destroyPxeNetwork, deleteNftablesRejectRules, PXE_NETWORK_NAME, PXE_GATEWAY, PXE_SUBNET } from "./helpers/pxe-network.js";
import { createPxeVm, destroyPxeVm, getVmMac, rebootPxeVm, readSerialLog } from "./helpers/pxe-vm.js";
import { sshExec } from "./helpers/ssh.js";

const IPXE_ARM64 = "/usr/share/ipxe/arm64-efi/snponly.efi";
const AAVMF = "/usr/share/edk2/aarch64/QEMU_EFI.fd";

const VM_MEMORY = 4096;
const VM_VCPUS = 2;
const VM_DISK_GB = 250;
const SSH_USER = "lab";
const BASTION_IP = PXE_GATEWAY;
const DHCP_RANGE_START = `${PXE_SUBNET}.100`;
const DHCP_RANGE_END = `${PXE_SUBNET}.200`;
const SERIAL_PORT = 4555;

// Emulated aarch64 -- generous timeouts throughout.
const DISCOVERY_TIMEOUT_MS = 20 * 60_000;
const RESCUE_TIMEOUT_MS = 20 * 60_000;
const INSTALL_TIMEOUT_MS = 75 * 60_000;
const SSH_TIMEOUT_MS = 20 * 60_000;

const RUN_FULL_INSTALL = process.env["ARM_PXE_FULL"] === "1";

function sleep(ms: number): Promise<void> {
  return new Promise((r) => setTimeout(r, ms));
}

function findSshKey(): { pubKey: string; keyPath: string } {
  const candidates: string[] = [];
  if (process.env["SSH_KEY_PATH"]) candidates.push(process.env["SSH_KEY_PATH"]);
  const homes = [homedir()];
  const sudoUser = process.env["SUDO_USER"];
  if (sudoUser) homes.push(join("/home", sudoUser));
  for (const home of homes) {
    for (const name of ["id_ed25519", "id_ecdsa", "id_rsa"]) {
      candidates.push(join(home, ".ssh", name));
    }
  }
  for (const keyPath of candidates) {
    if (existsSync(keyPath) && existsSync(`${keyPath}.pub`)) {
      return { pubKey: readFileSync(`${keyPath}.pub`, "utf-8").trim(), keyPath };
    }
  }
  throw new Error("No SSH key found — set SSH_KEY_PATH or ensure keys exist in ~/.ssh/");
}

async function pollApi<T>(
  url: string,
  check: (data: T) => boolean,
  timeoutMs: number,
  intervalMs = 10_000,
): Promise<T> {
  const start = Date.now();
  while (Date.now() - start < timeoutMs) {
    try {
      const res = await fetch(url);
      if (res.ok) {
        const data = (await res.json()) as T;
        if (check(data)) return data;
      }
    } catch { /* bastion not up yet, or a network hiccup */ }
    await sleep(intervalMs);
  }
  throw new Error(`Timeout after ${timeoutMs}ms polling ${url}`);
}

function requirePrerequisites(): void {
  if (!existsSync("/usr/bin/qemu-system-aarch64")) {
    throw new Error("qemu-system-aarch64 not installed. Run: sudo dnf install qemu-system-aarch64");
  }
  if (!existsSync(AAVMF)) {
    throw new Error(`AAVMF firmware not found at ${AAVMF}. Run: sudo dnf install edk2-aarch64`);
  }
  if (!existsSync(IPXE_ARM64)) {
    throw new Error(`arm64 iPXE not found at ${IPXE_ARM64}. Run: sudo dnf install ipxe-bootimgs-aarch64`);
  }
}

/**
 * Confirm the arm64 iPXE binary implements EFI_LOAD_FILE2_PROTOCOL.
 *
 * Without it the `initrd` line is accepted and silently ignored, and the kernel panics
 * with unknown-block(0,0). Checking here turns a confusing 30-minute boot failure into
 * an immediate, explanatory one.
 *
 * GUID 4006c0c1-fcb3-403e-996d-4a6c8724e06d, little-endian in the binary's GUID table.
 */
function assertIpxeSupportsLoadFile2(): void {
  const LOAD_FILE2_GUID = Buffer.from("c1c00640b3fc3e40996d4a6c8724e06d", "hex");
  const binary = readFileSync(IPXE_ARM64);
  if (binary.indexOf(LOAD_FILE2_GUID) < 0) {
    throw new Error(
      `${IPXE_ARM64} does not reference EFI_LOAD_FILE2_PROTOCOL. On arm64 the kernel ` +
      `EFI stub fetches the initrd over LoadFile2; without it iPXE drops the initrd ` +
      `silently and the kernel panics with "unknown-block(0,0)". Rebuild iPXE with ` +
      `LoadFile2, or chainload grubaa64.efi for aarch64 instead.`,
    );
  }
  log(`iPXE arm64 implements LoadFile2 — initrd will be delivered to the EFI stub`);
}

interface Harness {
  testDir: string;
  app: { close: () => Promise<void> };
  stopDnsmasq: () => void;
  vmMac: string;
  httpPort: number;
}

/** Bring up an isolated network, a bastion with both arch payloads, and an arm64 VM. */
async function startHarness(vmName: string, httpPort: number, pubKey: string): Promise<Harness> {
  requirePrerequisites();
  assertIpxeSupportsLoadFile2();

  log("Setting up PXE test network...");
  ensurePxeNetwork();

  const testDir = join(tmpdir(), `lab-arm-pxe-test-${Date.now()}`);
  for (const sub of ["tftp", "http", "logs"]) {
    mkdirSync(join(testDir, sub), { recursive: true });
  }

  const { createApp } = await import("../../src/bastion/src/server.js");
  const { loadConfig } = await import("../../src/bastion/src/config.js");
  const { generateDnsmasqConf, startDnsmasq, stopDnsmasq } = await import("../../src/bastion/src/services/dnsmasq.js");
  const { generateDiscoverKickstart } = await import("../../src/bastion/src/services/kickstart-generator.js");
  const { renderBootIpxe, kernelPath, initrdPath } = await import("../../src/bastion/src/templates/boot.ipxe.js");
  // Relative, not "@lab/shared": these tests run from the repo root against sources,
  // where the workspace package alias is not resolvable.
  const { SUPPORTED_ARCHES, fedoraMirrorFor } = await import("../../src/shared/src/hardware/index.js");

  const config = loadConfig({
    bastionDir: testDir,
    httpPort,
    iface: "virbr-pxe",
    serverIp: BASTION_IP,
    network: `${PXE_SUBNET}.0`,
    gateway: BASTION_IP,
    dhcpMode: "full",
    dhcpRangeStart: DHCP_RANGE_START,
    dhcpRangeEnd: DHCP_RANGE_END,
    domain: "arm-pxe-test.local",
    sshKeys: [pubKey],
    adminUser: SSH_USER,
  });

  // iPXE binaries. The arm64 one is the whole point: dnsmasq hands it out on DHCP
  // option 93 = 11, and x86_64 is staged too so the config renders as it does in
  // production.
  log("Staging iPXE binaries...");
  copyFileSync(IPXE_ARM64, join(config.tftpDir, "ipxe-arm64.efi"));
  const ipxeX86 = "/usr/share/ipxe/ipxe-snponly-x86_64.efi";
  if (existsSync(ipxeX86)) copyFileSync(ipxeX86, join(config.tftpDir, "ipxe.efi"));

  // Fedora kernel + initrd for both architectures, cached across runs.
  const cacheDir = "/var/lib/libvirt/images/lab-pxe-cache";
  execSync(`mkdir -p "${cacheDir}"`, { stdio: "pipe" });

  for (const arch of SUPPORTED_ARCHES) {
    const mirror = fedoraMirrorFor(config.fedoraVersion, arch);
    const kernelCache = join(cacheDir, `vmlinuz-${arch}`);
    const initrdCache = join(cacheDir, `initrd-${arch}.img`);

    if (!existsSync(kernelCache)) {
      log(`Downloading Fedora ${config.fedoraVersion} ${arch} kernel...`);
      execSync(`curl -# -L -f -o "${kernelCache}" "${mirror}/images/pxeboot/vmlinuz"`, { stdio: "inherit", timeout: 600_000 });
    }
    if (!existsSync(initrdCache)) {
      log(`Downloading Fedora ${config.fedoraVersion} ${arch} initrd...`);
      execSync(`curl -# -L -f -o "${initrdCache}" "${mirror}/images/pxeboot/initrd.img"`, { stdio: "inherit", timeout: 600_000 });
    }

    // Staged under the exact names the iPXE templates will ask for.
    copyFileSync(kernelCache, join(config.httpDir, kernelPath(arch)));
    copyFileSync(initrdCache, join(config.httpDir, initrdPath(arch)));
    log(`Staged ${arch}: ${kernelPath(arch)} + ${initrdPath(arch)}`);
  }

  writeFileSync(join(config.httpDir, "discover.ks"), generateDiscoverKickstart(config));
  writeFileSync(
    join(config.httpDir, "boot.ipxe"),
    renderBootIpxe({ serverIp: config.serverIp, httpPort: config.httpPort }),
  );
  generateDnsmasqConf(config);

  const { app, syslog } = createApp(config);
  await app.listen({ port: config.httpPort, host: "0.0.0.0" });
  syslog.start();
  log(`Bastion HTTP listening on :${config.httpPort}`);

  log("Starting dnsmasq (full DHCP)...");
  startDnsmasq(config).catch((err) => {
    log(`dnsmasq failed: ${err instanceof Error ? err.message : String(err)}`);
  });
  await sleep(1500);

  log("Creating aarch64 PXE VM (emulated — this is slow)...");
  createPxeVm({
    name: vmName,
    memory: VM_MEMORY,
    vcpus: VM_VCPUS,
    diskSize: VM_DISK_GB,
    network: PXE_NETWORK_NAME,
    arch: "aarch64",
  });

  const vmMac = getVmMac(vmName);
  if (!vmMac) throw new Error("Could not determine VM MAC address");
  log(`ARM VM MAC: ${vmMac}`);

  return { testDir, app, stopDnsmasq, vmMac, httpPort: config.httpPort };
}

async function stopHarness(vmName: string, harness: Harness | undefined): Promise<void> {
  log("Cleaning up...");
  if (harness) {
    await harness.app.close().catch(() => {});
    harness.stopDnsmasq();
  }
  destroyPxeVm(vmName);
  destroyPxeNetwork();
  if (harness) rmSync(harness.testDir, { recursive: true, force: true });
}

/** Read the DHCP lease the bastion handed a MAC. Rescue mode reports no IP itself. */
function leaseIpFor(testDir: string, mac: string): string | null {
  const leaseFile = join(testDir, "dnsmasq.leases");
  if (!existsSync(leaseFile)) return null;
  for (const line of readFileSync(leaseFile, "utf-8").split("\n")) {
    // <expiry> <mac> <ip> <hostname> <clientid>
    const parts = line.trim().split(/\s+/);
    if (parts.length >= 3 && parts[1]?.toLowerCase() === mac.toLowerCase()) {
      return parts[2] ?? null;
    }
  }
  return null;
}

async function waitForLease(testDir: string, mac: string, timeoutMs: number): Promise<string> {
  const start = Date.now();
  while (Date.now() - start < timeoutMs) {
    const ip = leaseIpFor(testDir, mac);
    if (ip !== null) return ip;
    await sleep(5000);
  }
  throw new Error(`No DHCP lease for ${mac} within ${timeoutMs}ms`);
}

// ---------------------------------------------------------------------------
// Rescue path -- what the DGX Sparks need.
// ---------------------------------------------------------------------------

describe("ARM PXE rescue", () => {
  const VM_NAME = "lab-arm-pxe-rescue";
  const HTTP_PORT = 8096;
  let harness: Harness | undefined;
  let sshKeyPath: string;
  let rescueIp: string;

  beforeAll(async () => {
    const { pubKey, keyPath } = findSshKey();
    sshKeyPath = keyPath;
    harness = await startHarness(VM_NAME, HTTP_PORT, pubKey);
    const { testDir, vmMac } = harness;

    // 1. The VM PXE boots. This exercises the whole chain: DHCP option 93 -> arm64
    //    iPXE -> /boot.ipxe -> /dispatch?arch=arm64 -> aarch64 kernel + initrd.
    //    Reaching discovery at all proves the initrd was delivered over LoadFile2 --
    //    without it the kernel panics before it can run the kickstart.
    log("Waiting for aarch64 VM to PXE boot and report discovery...");
    log("(emulated aarch64 — be patient)");
    type MachinesResponse = { discovered: Record<string, { arch: string }> };
    const discovered = await pollApi<MachinesResponse>(
      `http://${BASTION_IP}:${HTTP_PORT}/api/machines`,
      (data) => vmMac in data.discovered,
      DISCOVERY_TIMEOUT_MS,
    ).catch(async (err) => {
      log("Discovery timed out. Serial console:");
      try {
        log(await readSerialLog(SERIAL_PORT, { lastLines: 100, timeoutMs: 15_000 }));
      } catch { /* console unavailable */ }
      throw err;
    });
    log(`aarch64 VM discovered as ${discovered.discovered[vmMac]?.arch}`);

    // 2. Queue rescue, exactly as `labctl provision debug` does.
    log("Queueing debug/rescue mode...");
    const res = await fetch(`http://${BASTION_IP}:${HTTP_PORT}/api/debug`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ mac: vmMac }),
    });
    expect(res.ok).toBe(true);

    // 3. Reboot into rescue. Rescue mode runs no %pre/%post, so the machine never
    //    reports its own IP — read it from the DHCP lease instead.
    await sleep(20_000);
    rebootPxeVm(VM_NAME);
    await sleep(5_000);
    deleteNftablesRejectRules();

    log("Waiting for rescue environment DHCP lease...");
    rescueIp = await waitForLease(testDir, vmMac, RESCUE_TIMEOUT_MS);
    log(`Rescue IP: ${rescueIp}`);

    log("Waiting for SSH into the rescue shell (started by inst.sshd)...");
    await waitForSsh(rescueIp, "root", SSH_TIMEOUT_MS, sshKeyPath).catch(async (err) => {
      log("Rescue SSH timed out. Serial console:");
      try {
        log(await readSerialLog(SERIAL_PORT, { lastLines: 100, timeoutMs: 15_000 }));
      } catch { /* console unavailable */ }
      throw err;
    });
    log("ARM PXE rescue reached.");
  }, DISCOVERY_TIMEOUT_MS + RESCUE_TIMEOUT_MS + SSH_TIMEOUT_MS + 300_000);

  afterAll(async () => { await stopHarness(VM_NAME, harness); });

  it("discovered the machine as aarch64", async () => {
    const res = await fetch(`http://${BASTION_IP}:${HTTP_PORT}/api/machines`);
    const data = (await res.json()) as { discovered: Record<string, { arch: string }> };
    expect(data.discovered[harness!.vmMac]?.arch).toBe("aarch64");
  });

  it("rescue shell is reachable over SSH and is aarch64", () => {
    const result = sshExec(rescueIp, "root", "uname -m", { keyPath: sshKeyPath, timeout: 60_000 });
    expect(result.exitCode).toBe(0);
    expect(result.stdout.trim()).toBe("aarch64");
  });

  it("booted an initramfs — the LoadFile2 path worked", () => {
    // If iPXE had dropped the initrd the kernel would never have reached userspace at
    // all, but assert it explicitly so a regression names itself.
    const result = sshExec(rescueIp, "root", "cat /proc/cmdline; ls /run/install", {
      keyPath: sshKeyPath, timeout: 60_000,
    });
    expect(result.exitCode).toBe(0);
    expect(result.stdout).toContain("inst.rescue");
  });

  it("rescue kernel came from the bastion over HTTP", () => {
    const result = sshExec(rescueIp, "root", "cat /proc/cmdline", { keyPath: sshKeyPath, timeout: 60_000 });
    expect(result.stdout).toContain(`${BASTION_IP}:${HTTP_PORT}`);
    // arm64 gets serial console arguments, never nomodeset.
    expect(result.stdout).toContain("console=ttyAMA0");
    expect(result.stdout).not.toContain("nomodeset");
  });

  it("has LVM tools available for inspecting an installed system", () => {
    const result = sshExec(rescueIp, "root", "command -v vgchange && command -v lsblk", {
      keyPath: sshKeyPath, timeout: 60_000,
    });
    expect(result.exitCode).toBe(0);
  });
});

// ---------------------------------------------------------------------------
// Full install -- opt-in, ~60-90 minutes emulated.
// ---------------------------------------------------------------------------

describe.runIf(RUN_FULL_INSTALL)("ARM PXE install", () => {
  const VM_NAME = "lab-arm-pxe-install";
  const HTTP_PORT = 8095;
  let harness: Harness | undefined;
  let sshKeyPath: string;
  let vmIp: string;

  beforeAll(async () => {
    const { pubKey, keyPath } = findSshKey();
    sshKeyPath = keyPath;
    harness = await startHarness(VM_NAME, HTTP_PORT, pubKey);
    const { vmMac } = harness;

    log("Waiting for aarch64 discovery...");
    await pollApi<{ discovered: Record<string, unknown> }>(
      `http://${BASTION_IP}:${HTTP_PORT}/api/machines`,
      (data) => vmMac in data.discovered,
      DISCOVERY_TIMEOUT_MS,
    );
    log("Discovered. Queueing install...");

    const installRes = await fetch(`http://${BASTION_IP}:${HTTP_PORT}/api/install`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ mac: vmMac, hostname: VM_NAME, disk: "", role: "vanilla" }),
    });
    expect(installRes.status).toBe(200);

    await sleep(30_000);
    rebootPxeVm(VM_NAME);

    log("Waiting for the emulated aarch64 install (60-90 min)...");
    type LogsResponse = { status: string; progress: string; ip?: string };
    const final = await pollApi<LogsResponse>(
      `http://${BASTION_IP}:${HTTP_PORT}/api/logs/${encodeURIComponent(vmMac)}`,
      (d) => d.status === "installed" || d.progress === "error",
      INSTALL_TIMEOUT_MS,
      30_000,
    );

    if (final.progress === "error") {
      const logs = await (await fetch(`http://${BASTION_IP}:${HTTP_PORT}/api/logs/${encodeURIComponent(vmMac)}`)).json();
      log(`ARM install FAILED: ${JSON.stringify(logs, null, 2)}`);
      throw new Error("ARM PXE install failed — see logs above");
    }

    vmIp = final.ip ?? "";
    log(`ARM install complete. IP: ${vmIp}`);

    await sleep(30_000);
    rebootPxeVm(VM_NAME);
    await sleep(5_000);
    deleteNftablesRejectRules();
    await waitForSsh(vmIp, SSH_USER, SSH_TIMEOUT_MS, sshKeyPath);
  }, DISCOVERY_TIMEOUT_MS + INSTALL_TIMEOUT_MS + SSH_TIMEOUT_MS + 600_000);

  afterAll(async () => { await stopHarness(VM_NAME, harness); });

  it("machine reached installed state", async () => {
    const res = await fetch(`http://${BASTION_IP}:${HTTP_PORT}/api/machines`);
    const data = (await res.json()) as { installed: Record<string, { hostname: string }> };
    expect(data.installed[harness!.vmMac]?.hostname).toBe(VM_NAME);
  });

  it("installed system is aarch64", () => {
    const result = sshExec(vmIp, SSH_USER, "uname -m", { keyPath: sshKeyPath, timeout: 60_000 });
    expect(result.stdout.trim()).toBe("aarch64");
  });

  it("SSH works with the admin user", () => {
    const result = sshExec(vmIp, SSH_USER, "whoami", { keyPath: sshKeyPath, timeout: 60_000 });
    expect(result.stdout.trim()).toBe(SSH_USER);
  });

  it("LVM layout is correct", () => {
    const result = sshExec(vmIp, SSH_USER, "sudo lvs labvg --noheadings -o lv_name", {
      keyPath: sshKeyPath, timeout: 60_000,
    });
    expect(result.exitCode).toBe(0);
    const lvs = result.stdout.trim().split("\n").map((l) => l.trim());
    for (const expected of ["root", "var", "varlog", "swap", "home", "srv"]) {
      expect(lvs).toContain(expected);
    }
  });
});