// First-boot fixups, run as root by `eydos-installer firstboot`
// (ExecStartPre of the daemon unit, ordered before loginmanager so
// everything here happens before any session or autologin decision).
//
// Also hosts the staging side the daemon calls from the Configure
// flow: the deferred account rename (usermod cannot rename a logged-in
// user) and the root re-encryption request the initramfs executes on
// the next boot (a mounted ext4 cannot be shrunk for the LUKS header).

use std::path::{ Path, PathBuf };
use std::process::Command;

pub const ACCOUNT_STAGED:   &str = "/var/lib/eydos-installer/account-staged";
pub const REENCRYPT_KEY:    &str = "/var/lib/eydos-installer/reencrypt-key";
pub const REENCRYPT_DEVICE: &str = "/var/lib/eydos-installer/reencrypt-device";

pub fn run()
{
	reencrypt_followup();
	apply_staged_account();
	direct_flash_fixup();
}

/// The image build always seeds `install-pending`, but an image
/// written straight to its final disk (eMMC via USB flashing, a PC
/// SSD) has nothing left to install: grow the root partition to the
/// disk and hand the first boot to the Configure flow instead.
fn direct_flash_fixup()
{
	if !Path::new( crate::install::INSTALL_PENDING ).exists() { return; }
	let Some( disk ) = crate::install::running_system_disk() else { return };
	if !disk_is_fixed( &disk ) { return; }

	eprintln!( "firstboot: root lives on fixed storage {}, switching to the Configure flow", disk.display() );
	if let Err( e ) = grow_root()
	{
		eprintln!( "firstboot: grow root: {e}" );
	}
	// Written in this order so a crash between the two still leaves a
	// pending marker; the wizard prefers Configure when both exist.
	let _ = std::fs::write( crate::install::CONFIGURE_PENDING, "" );
	let _ = std::fs::remove_file( crate::install::INSTALL_PENDING );
}

/// Removable media and anything on the USB bus count as install
/// media; eMMC, NVMe, SATA and VirtIO as final storage.
fn disk_is_fixed( disk: &Path ) -> bool
{
	let Some( name ) = disk.file_name().and_then( |s| s.to_str() ) else { return false };
	let removable = std::fs::read_to_string( format!( "/sys/block/{name}/removable" ) )
		.map( |s| s.trim() == "1" )
		.unwrap_or( false );
	if removable { return false; }
	crate::disks::classify_bus( name ) != crate::disks::BusKind::Usb
}

/// Grow the root partition to the end of its disk (only when it is
/// the last one) and the filesystem with it. The partition table is
/// edited directly; the kernel picks the new size up through
/// `partx -u`, which resizes mounted partitions in place.
fn grow_root() -> Result<(), String>
{
	let part = root_partition().ok_or( "cannot resolve the root partition" )?;
	let ( disk, num ) = partition_parent( &part ).ok_or( "cannot resolve the root disk" )?;

	let mut file = std::fs::File::open( &disk )
		.map_err( |e| format!( "open {}: {e}", disk.display() ) )?;
	let gpt = gptman::GPT::find_from( &mut file )
		.map_err( |e| format!( "no gpt on {}: {e:?}", disk.display() ) )?;
	let sector = gpt.sector_size;
	let entry = gpt.iter()
		.find( |( i, p )| *i == num && p.is_used() )
		.map( |( _, p )| p.clone() )
		.ok_or_else( || format!( "p{num}: not in GPT on {}", disk.display() ) )?;
	if gpt.iter().any( |( i, p )| i != num && p.is_used() && p.starting_lba > entry.starting_lba )
	{
		return Ok( () );
	}
	let cur_sectors = entry.ending_lba - entry.starting_lba + 1;
	let max_sectors = gpt.header.last_usable_lba + 1 - entry.starting_lba;
	if max_sectors * sector <= cur_sectors * sector + 64 * 1024 * 1024
	{
		return Ok( () );
	}
	drop( file );

	eprintln!( "firstboot: growing {} to fill {}", part.display(), disk.display() );
	crate::disks::resize_partition( &disk, num, max_sectors * sector )?;
	crate::install::run_ok( Command::new( "partx" ).arg( "-u" ).arg( &disk ) )?;
	crate::install::run_ok( Command::new( "resize2fs" ).arg( &part ) )
}

fn root_partition() -> Option<PathBuf>
{
	let mounts = std::fs::read_to_string( "/proc/self/mounts" ).ok()?;
	let source = mounts.lines().find_map( |l|
	{
		let mut it = l.split_whitespace();
		let src = it.next()?;
		let mnt = it.next()?;
		if mnt == "/" { Some( src.to_string() ) } else { None }
	} )?;
	if !source.starts_with( "/dev/" ) { return None; }
	Some( PathBuf::from( source ) )
}

fn partition_parent( part: &Path ) -> Option<( PathBuf, u32 )>
{
	let name = part.file_name()?.to_str()?;
	let sys = std::fs::canonicalize( format!( "/sys/class/block/{name}" ) ).ok()?;
	let num: u32 = std::fs::read_to_string( sys.join( "partition" ) )
		.ok()?
		.trim()
		.parse()
		.ok()?;
	let parent = sys.parent()?.file_name()?.to_str()?.to_string();
	Some( ( PathBuf::from( "/dev" ).join( parent ), num ) )
}

/// Apply the account the Configure flow staged on the previous boot:
/// nobody is logged in yet, so the factory user can be renamed.
fn apply_staged_account()
{
	let Ok( raw ) = std::fs::read_to_string( ACCOUNT_STAGED ) else { return };
	let mut lines = raw.lines();
	let ( Some( username ), Some( hostname ), Some( password ) ) =
		( lines.next(), lines.next(), lines.next() )
	else
	{
		let _ = std::fs::remove_file( ACCOUNT_STAGED );
		return;
	};

	eprintln!( "firstboot: applying staged account {username}" );
	let account = crate::install::AccountSetup
	{
		username: username.to_string(),
		hostname: hostname.to_string(),
		password: password.to_string(),
	};
	if let Err( e ) = crate::install::apply_account( Path::new( "/" ), &account )
	{
		eprintln!( "firstboot: apply account: {e}" );
	}
	let _ = std::fs::remove_file( ACCOUNT_STAGED );
	// The Configure flow left its marker in place so this unit would
	// run; the flow is only over once the account is applied.
	let _ = std::fs::write( crate::install::CONFIGURED, "" );
	let _ = std::fs::remove_file( crate::install::CONFIGURE_PENDING );
}

/// Once the initramfs has encrypted the root (its backing device is
/// now the opened mapper), bind the TPM when there is one, retire the
/// staged key and rebuild the initramfs — from here on the standard
/// crypttab path unlocks at boot.
fn reencrypt_followup()
{
	if !Path::new( REENCRYPT_KEY ).exists() { return; }
	let on_mapper = root_partition().map_or( false, |p| p.to_string_lossy().contains( "/mapper/" ) );
	if !on_mapper { return; }

	eprintln!( "firstboot: re-encryption finished, retiring the staged key" );
	let partuuid = std::fs::read_to_string( REENCRYPT_DEVICE ).unwrap_or_default();
	let dev = PathBuf::from( format!( "/dev/disk/by-partuuid/{}", partuuid.trim() ) );
	if crate::install::has_tpm2() && dev.exists()
	{
		if let Err( e ) = crate::install::run_ok(
			Command::new( "clevis" )
				.args( [ "luks", "bind", "-f", "-k", REENCRYPT_KEY, "-d" ] )
				.arg( &dev )
				.args( [ "tpm2", "{\"pcr_ids\":\"7\"}" ] ) )
		{
			eprintln!( "firstboot: TPM bind failed, boot will ask for the password: {e}" );
		}
	}
	let _ = std::fs::remove_file( REENCRYPT_KEY );
	let _ = std::fs::remove_file( REENCRYPT_DEVICE );
	if dev.exists() && !crate::install::clevis_is_bound( &dev )
	{
		crate::install::drop_clevis_initramfs( Path::new( "/" ) );
	}
	if let Err( e ) = crate::install::run_ok(
		Command::new( "update-initramfs" ).arg( "-u" ) )
	{
		eprintln!( "firstboot: update-initramfs: {e}" );
	}
	if crate::platform::detect() == crate::platform::BootFlavor::UBootScript
	{
		if let Err( e ) = crate::platform::regen_uinitrd( Path::new( "/" ) )
		{
			eprintln!( "firstboot: uInitrd: {e}" );
		}
	}
}

// ── Staging (called by the daemon while configure-pending exists) ───────────

pub fn stage( username: &str, hostname: &str, password: &str, encrypt: bool ) -> Result<(), String>
{
	write_private( Path::new( ACCOUNT_STAGED ), &format!( "{username}\n{hostname}\n{password}\n" ) )?;
	if encrypt
	{
		stage_reencrypt( password )?;
	}
	Ok( () )
}

/// Prepare the next boot's initramfs to encrypt the root in place:
/// stage the key (the whole file is the passphrase — no trailing
/// newline), record the partition, write crypttab, and rebuild the
/// initramfs so the eydos-reencrypt hook embeds both.
fn stage_reencrypt( password: &str ) -> Result<(), String>
{
	let part = root_partition().ok_or( "cannot resolve the root partition" )?;
	if part.to_string_lossy().contains( "/mapper/" )
		|| crate::disks::blkid_fs_type( &part ).as_deref() == Some( "crypto_LUKS" )
	{
		return Err( "the root is already encrypted".to_string() );
	}
	let partuuid = blkid_value( &part, "PARTUUID" )
		.ok_or_else( || format!( "no PARTUUID for {}", part.display() ) )?;

	write_private( Path::new( REENCRYPT_KEY ), password )?;
	write_private( Path::new( REENCRYPT_DEVICE ), &partuuid )?;
	std::fs::write(
		"/etc/crypttab",
		format!( "eydos_crypt\tPARTUUID={partuuid}\tnone\t{}\n",
			crate::install::crypttab_options( Path::new( "/" ) ) ),
	).map_err( |e| format!( "write crypttab: {e}" ) )?;

	crate::install::run_ok( Command::new( "update-initramfs" ).arg( "-u" ) )?;
	if crate::platform::detect() == crate::platform::BootFlavor::UBootScript
	{
		crate::platform::regen_uinitrd( Path::new( "/" ) )?;
	}
	Ok( () )
}

fn write_private( path: &Path, content: &str ) -> Result<(), String>
{
	use std::os::unix::fs::PermissionsExt;
	std::fs::write( path, content ).map_err( |e| format!( "write {}: {e}", path.display() ) )?;
	std::fs::set_permissions( path, std::fs::Permissions::from_mode( 0o600 ) )
		.map_err( |e| format!( "chmod {}: {e}", path.display() ) )
}

fn blkid_value( path: &Path, tag: &str ) -> Option<String>
{
	let out = Command::new( "blkid" )
		.args( [ "-o", "value", "-s", tag ] )
		.arg( path )
		.output().ok()?;
	if !out.status.success() { return None; }
	let s = String::from_utf8_lossy( &out.stdout ).trim().to_string();
	if s.is_empty() { None } else { Some( s ) }
}
