// Per-platform boot handling.
//
// The boot project builds each image with one of three boot layouts
// (install_base.sh); the installed system must reproduce the layout of
// the medium it was copied from. Detection reads the running medium
// itself, so no per-target configuration needs to ship in the image.

use std::path::Path;
use std::process::Command;

#[ derive( Clone, Copy, Debug, PartialEq, Eq ) ]
pub enum BootFlavor
{
	/// UEFI or BIOS: GRUB on an ESP (amd64, Qualcomm UEFI boards).
	Grub,
	/// u-boot in the raw area before the first partition, extlinux
	/// config generated by u-boot-menu (imx8: Librem 5, Liberux Nexx).
	UBootExtlinux,
	/// u-boot in the raw area, boot.scr + env.txt (Rockchip boards).
	UBootScript,
}

pub fn detect() -> BootFlavor
{
	if cfg!( target_arch = "x86_64" ) || cfg!( target_arch = "x86" )
		|| Path::new( "/sys/firmware/efi" ).exists()
	{
		return BootFlavor::Grub;
	}
	if Path::new( "/boot/extlinux/extlinux.conf" ).exists()
	{
		return BootFlavor::UBootExtlinux;
	}
	if Path::new( "/boot/boot.scr" ).exists()
	{
		return BootFlavor::UBootScript;
	}
	BootFlavor::Grub
}

/// u-boot lives between the GPT structures and the first partition:
/// flash.bin at 33 KiB on imx8, idbloader at sector 64 plus u-boot.itb
/// at sector 16384 on Rockchip. The images start the first partition at
/// 16 MiB, so copying `[ 32 KiB, 16 MiB )` carries whichever loader the
/// platform uses without touching either partition table.
pub const BOOT_AREA_START: u64 = 32 * 1024;
pub const BOOT_AREA_END:   u64 = 16 * 1024 * 1024;

pub fn copy_boot_area( source_disk: &Path, target_disk: &Path ) -> Result<(), String>
{
	use std::io::{ Read, Seek, SeekFrom, Write };

	let mut src = std::fs::File::open( source_disk )
		.map_err( |e| format!( "open {}: {e}", source_disk.display() ) )?;
	let mut dst = std::fs::OpenOptions::new()
		.write( true )
		.open( target_disk )
		.map_err( |e| format!( "open {}: {e}", target_disk.display() ) )?;
	src.seek( SeekFrom::Start( BOOT_AREA_START ) )
		.map_err( |e| format!( "seek {}: {e}", source_disk.display() ) )?;
	dst.seek( SeekFrom::Start( BOOT_AREA_START ) )
		.map_err( |e| format!( "seek {}: {e}", target_disk.display() ) )?;

	let mut remaining = BOOT_AREA_END - BOOT_AREA_START;
	let mut buf = vec![ 0u8; 1024 * 1024 ];
	while remaining > 0
	{
		let want = remaining.min( buf.len() as u64 ) as usize;
		let n = src.read( &mut buf[ ..want ] )
			.map_err( |e| format!( "read {}: {e}", source_disk.display() ) )?;
		if n == 0 { break; }
		dst.write_all( &buf[ ..n ] )
			.map_err( |e| format!( "write {}: {e}", target_disk.display() ) )?;
		remaining -= n as u64;
	}
	dst.sync_all().ok();
	Ok( () )
}

/// Point the copied system at its new root and regenerate
/// /boot/extlinux/extlinux.conf inside the target. `root_arg` is the
/// full kernel argument ("root=UUID=…" or "root=/dev/mapper/…").
pub fn configure_uboot_extlinux( target: &Path, root_arg: &str ) -> Result<(), String>
{
	let path = target.join( "etc/default/u-boot" );
	let current = std::fs::read_to_string( &path ).unwrap_or_default();
	let mut new = String::new();
	let mut replaced = false;
	for line in current.lines()
	{
		let t = line.trim_start();
		if t.starts_with( "U_BOOT_ROOT=" ) || t.starts_with( "#U_BOOT_ROOT=" )
		{
			new.push_str( &format!( "U_BOOT_ROOT=\"{root_arg}\"" ) );
			replaced = true;
		}
		else
		{
			new.push_str( line );
		}
		new.push( '\n' );
	}
	if !replaced
	{
		new.push_str( &format!( "U_BOOT_ROOT=\"{root_arg}\"\n" ) );
	}
	std::fs::write( &path, new )
		.map_err( |e| format!( "write {}: {e}", path.display() ) )?;

	crate::install::run_ok( Command::new( "chroot" ).arg( target ).arg( "u-boot-update" ) )
}

/// Rewrite the root device in /boot/env.txt and the boot.cmd default,
/// then rebuild boot.scr with the target's mkimage. `root_val` is the
/// value u-boot passes as root= ("UUID=…" or a mapper path).
pub fn configure_uboot_script( target: &Path, root_val: &str ) -> Result<(), String>
{
	let env_path = target.join( "boot/env.txt" );
	if let Ok( current ) = std::fs::read_to_string( &env_path )
	{
		let new: String = current.lines().map( |l|
			if l.trim_start().starts_with( "rootdev=" )
			{
				format!( "rootdev={root_val}\n" )
			}
			else
			{
				format!( "{l}\n" )
			} ).collect();
		std::fs::write( &env_path, new )
			.map_err( |e| format!( "write {}: {e}", env_path.display() ) )?;
	}

	let cmd_path = target.join( "boot/boot.cmd" );
	let current = std::fs::read_to_string( &cmd_path )
		.map_err( |e| format!( "read {}: {e}", cmd_path.display() ) )?;
	let new: String = current.lines().map( |l|
		if l.trim_start().starts_with( "setenv rootdev " )
		{
			format!( "setenv rootdev \"{root_val}\"\n" )
		}
		else
		{
			format!( "{l}\n" )
		} ).collect();
	std::fs::write( &cmd_path, new )
		.map_err( |e| format!( "write {}: {e}", cmd_path.display() ) )?;

	crate::install::run_ok( Command::new( "chroot" ).arg( target ).args( [
		"mkimage", "-C", "none", "-A", "arm64", "-T", "script",
		"-d", "/boot/boot.cmd", "/boot/boot.scr",
	] ) )
}

/// boot.scr loads /boot/uInitrd; after update-initramfs refreshes
/// initrd.img-* the u-boot-wrapped copies must be rebuilt to match.
pub fn regen_uinitrd( target: &Path ) -> Result<(), String>
{
	let boot = target.join( "boot" );
	let entries = std::fs::read_dir( &boot )
		.map_err( |e| format!( "read {}: {e}", boot.display() ) )?;
	for entry in entries.flatten()
	{
		let name = entry.file_name();
		let Some( name ) = name.to_str() else { continue };
		let Some( ver ) = name.strip_prefix( "initrd.img-" ) else { continue };
		crate::install::run_ok( Command::new( "chroot" ).arg( target ).args( [
			"mkimage", "-n", "uInitrd", "-A", "arm64", "-O", "linux",
			"-T", "ramdisk", "-C", "gzip",
			"-d", &format!( "/boot/initrd.img-{ver}" ),
			&format!( "/boot/uInitrd-{ver}" ),
		] ) )?;
	}
	Ok( () )
}
