package install

import (
	_ "embed"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
)

//go:embed assets/spinifex-boot-tool
var bootToolScript []byte

//go:embed assets/grub-install-wrapper
var grubInstallWrapper []byte

const (
	bootToolPath     = "/usr/sbin/spinifex-boot-tool"
	grubInstallPath  = "/usr/sbin/grub-install"
	grubInstallReal  = "/usr/sbin/grub-install.real"
	bootToolUUIDList = "/etc/kernel/spinifex-boot-uuids"
)

// installBootToolZFS sets up per-disk bootloaders on a ZFS root system.
//
// Nothing in the boot path is allowed to read the pool: GRUB's ZFS support
// trails OpenZFS, so a pool feature flag enabled years from now could make the
// machine unbootable. Each disk instead carries a self-contained ESP with its
// own GRUB, kernels and initramfs.
func installBootToolZFS(cfg DiskConfig) error {
	if err := writeTargetFile(bootToolPath, bootToolScript, 0o755); err != nil {
		return err
	}
	if err := writeGrubDefaults(cfg); err != nil {
		return err
	}

	if err := bindChrootMounts(); err != nil {
		return err
	}
	defer unbindChrootMounts()

	// Diverted before the first init: init runs a refresh, and refresh calls
	// grub-install. With the wrapper already in place that call has to reach
	// the real binary, which is what the diversion arranges.
	if err := divertGrubInstall(); err != nil {
		return err
	}

	// A partially-updated set of ESPs is an install that boots today and fails
	// after the next kernel upgrade, so every disk must succeed here.
	env := []string{"SPINIFEX_BOOT_TOOL_STRICT=1"}
	for _, d := range cfg.Disks {
		esp := d.PartitionPath(espPartNum)
		if err := runEnv(env, "chroot", mountRoot, bootToolPath, "init", esp); err != nil {
			return fmt.Errorf("spinifex-boot-tool init %s: %w", esp, err)
		}
	}

	// The one moment firmware boot variables are written. Refreshes run on
	// every kernel upgrade and deliberately leave NVRAM alone: repeatedly
	// rewriting boot entries has bricked more than one vendor's boards. Losing
	// the entry is survivable anyway, because each ESP also carries the
	// removable-path loader that firmware falls back to.
	if err := runEnv(append(env, "SPINIFEX_BOOT_TOOL_NVRAM=1"),
		"chroot", mountRoot, bootToolPath, "refresh"); err != nil {
		return fmt.Errorf("spinifex-boot-tool refresh: %w", err)
	}
	return verifyBootConfig(cfg)
}

// divertGrubInstall moves the real grub-install aside and puts the wrapper in
// its place. dpkg-divert makes the move survive package upgrades, which is the
// point: grub-pc's postinst runs grub-install and would otherwise quietly
// replace a working boot layout with one that reads the pool.
func divertGrubInstall() error {
	if err := run("chroot", mountRoot, "dpkg-divert", "--divert", grubInstallReal,
		"--rename", "--add", grubInstallPath); err != nil {
		return fmt.Errorf("divert grub-install: %w", err)
	}
	return writeTargetFile(grubInstallPath, grubInstallWrapper, 0o755)
}

// verifyBootConfig asserts that the config the machine will actually boot from
// names the ZFS root. grub-mkconfig degrades quietly when grub-probe cannot
// read the pool, emitting a config that gets as far as the GRUB menu and then
// panics with "unable to mount root fs" — which looks like a hardware fault.
func verifyBootConfig(cfg DiskConfig) error {
	uuids, err := os.ReadFile(filepath.Join(mountRoot, strings.TrimPrefix(bootToolUUIDList, "/")))
	if err != nil {
		return fmt.Errorf("read ESP list: %w", err)
	}
	if len(strings.Fields(string(uuids))) != len(cfg.Disks) {
		return fmt.Errorf("expected %d registered ESPs, found %d — some disks have no bootloader",
			len(cfg.Disks), len(strings.Fields(string(uuids))))
	}

	// One ESP is enough to prove grub-mkconfig produced the right thing; they
	// are all generated by the same code path in the same run.
	mp := filepath.Join(mountRoot, "var/tmp/verify-esp")
	if err := os.MkdirAll(mp, 0o755); err != nil {
		return err
	}
	defer func() { _ = os.Remove(mp) }()

	esp := cfg.Disks[0].PartitionPath(espPartNum)
	if err := run("mount", esp, mp); err != nil {
		return fmt.Errorf("mount %s for verification: %w", esp, err)
	}
	defer func() { _ = run("umount", mp) }()

	grubCfg, err := os.ReadFile(filepath.Join(mp, "grub/grub.cfg"))
	if err != nil {
		return fmt.Errorf("read grub.cfg from %s: %w", esp, err)
	}
	want := "root=ZFS=" + ZFSRootDataset
	if !strings.Contains(string(grubCfg), want) {
		return fmt.Errorf("grub.cfg on %s does not contain %q — the installed system would not find its root filesystem", esp, want)
	}
	if !strings.Contains(string(grubCfg), "vmlinuz-") {
		return fmt.Errorf("grub.cfg on %s references no kernel", esp)
	}
	return nil
}

// writeTargetFile writes a file into the install target.
func writeTargetFile(path string, content []byte, mode os.FileMode) error {
	full := filepath.Join(mountRoot, strings.TrimPrefix(path, "/"))
	if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
		return err
	}
	if err := os.WriteFile(full, content, mode); err != nil {
		return fmt.Errorf("write %s: %w", path, err)
	}
	// WriteFile does not apply the mode to a file that already exists, and this
	// path is re-run on a retried install.
	return os.Chmod(full, mode)
}

// writeGrubDefaults writes /etc/default/grub and the branding snippet.
//
// Both consoles are listed: tty0 stays primary for a local install while ttyS0
// mirrors kernel output for headless racks and CI. The last console= on the
// cmdline becomes the system console, so serial-only hardware still works.
func writeGrubDefaults(cfg DiskConfig) error {
	cmdline := "console=tty0 console=ttyS0,115200 systemd.show_status=1"
	if cfg.FS.IsZFS() {
		// Stated explicitly rather than relying on grub-mkconfig's probe. It
		// appears after the root= that 10_linux computes, so it wins if the
		// probe came back with something wrong.
		cmdline = fmt.Sprintf("root=ZFS=%s boot=zfs %s", ZFSRootDataset, cmdline)
	}
	grubDefault := fmt.Sprintf(`GRUB_DEFAULT=0
GRUB_TIMEOUT=5
GRUB_DISTRIBUTOR=Spinifex
GRUB_CMDLINE_LINUX_DEFAULT=""
GRUB_CMDLINE_LINUX="%s"
`, cmdline)
	if err := os.WriteFile(filepath.Join(mountRoot, "etc/default/grub"), []byte(grubDefault), 0o644); err != nil {
		return fmt.Errorf("write /etc/default/grub: %w", err)
	}

	// gfxterm must be activated before serial is appended — background_image
	// silently does nothing in text mode. $prefix is used rather than a literal
	// /boot/grub so the same snippet is correct whether GRUB's root is the ext4
	// filesystem or a per-disk ESP.
	grubTheme := `#!/bin/sh
exec tail -n +3 $0
insmod all_video
insmod font
if loadfont $prefix/fonts/unicode.pf2; then
  set gfxmode=auto
  insmod gfxterm
  terminal_output gfxterm
fi
insmod serial
if serial --unit=0 --speed=115200 --timeout=1; then
  terminal_input  --append serial
  terminal_output --append serial
fi

# --- Branding ---
insmod png
set theme=$prefix/theme.txt
export theme
`
	if err := os.WriteFile(filepath.Join(mountRoot, "etc/grub.d/06_spinifex"), []byte(grubTheme), 0o755); err != nil {
		return fmt.Errorf("write /etc/grub.d/06_spinifex: %w", err)
	}
	return nil
}

// refreshBootToolNote logs the command an operator needs after replacing a
// disk. Nothing in the installer can do this for them later.
func refreshBootToolNote(cfg DiskConfig) {
	if cfg.FS.IsZFS() {
		slog.Info("bootloader installed on every pool member",
			"disks", len(cfg.Disks),
			"after_disk_replacement", "spinifex-boot-tool format <esp> && spinifex-boot-tool init <esp>")
	}
}