mkwim 0.1.0

// SPDX-License-Identifier: GPL-3.0-or-later
// Copyright (c) 2026 Jarkko Sakkinen

//! End-to-end pipeline: parse the ISO, prepare a block device or raw image,
//! copy the installation media, and split `install.wim` when it would exceed
//! FAT32's 4 GiB per-file limit.

use std::fs;
use std::io::{self, Write};
use std::os::unix::fs::MetadataExt;
use std::path::{Path, PathBuf};

use anyhow::{Context, Result, anyhow};

use crate::ffi::{Wim, WimLib};
use crate::{disk, iso, util};

/// Maximum size of a single file on a FAT32 volume: 2^32 - 1 bytes.
const FAT32_MAX_FILE: u64 = 4 * 1024 * 1024 * 1024 - 1;

/// Default part size for split WIM: 3584 MiB, comfortably below the 4 GiB
/// FAT32 per-file limit.
const SWM_PART_SIZE: u64 = 3584 * 1024 * 1024;

/// Run the creator with the parsed command-line arguments.
pub fn run(args: &crate::cli::Args) -> Result<()> {
    let output = args.output.as_path();
    let output_kind = disk::output_kind(output)?;
    if output_kind == disk::OutputKind::BlockDevice && !util::running_as_root() {
        return Err(anyhow!(
            "must run as root to write and mount a block-device output"
        ));
    }

    // libwim is only needed when splitting; initialise it eagerly so any
    // failure surfaces before we touch the output.
    let lib = WimLib::new()?;

    let iso_path = args.iso.as_path();
    if output_kind == disk::OutputKind::Image && same_file(iso_path, output)? {
        return Err(anyhow!(
            "output image {} must not be the input ISO {}",
            output.display(),
            iso_path.display()
        ));
    }
    let staging_path = match output_kind {
        disk::OutputKind::BlockDevice => iso_path,
        disk::OutputKind::Image => output,
    };
    let staging_dir = staging_path
        .parent()
        .filter(|path| !path.as_os_str().is_empty())
        .unwrap_or_else(|| Path::new("."));
    let (mut iso_file, root) = iso::open(iso_path)?;
    tracing::debug!("parsed ISO {}", iso_path.display());

    match output_kind {
        disk::OutputKind::BlockDevice => {
            run_device(&lib, &mut iso_file, &root, output, staging_dir)
        }
        disk::OutputKind::Image => run_image(&lib, &mut iso_file, &root, output, staging_dir),
    }
}

fn same_file(first: &Path, second: &Path) -> Result<bool> {
    let second = match fs::metadata(second) {
        Ok(metadata) => metadata,
        Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(false),
        Err(error) => {
            return Err(error).with_context(|| format!("stat output {}", second.display()));
        }
    };
    let first = fs::metadata(first).with_context(|| format!("stat ISO {}", first.display()))?;
    Ok(first.dev() == second.dev() && first.ino() == second.ino())
}

fn run_device(
    lib: &WimLib,
    iso_file: &mut fs::File,
    root: &isomage::TreeNode,
    output: &Path,
    staging_dir: &Path,
) -> Result<()> {
    let part = disk::partition_drive(output)?;
    tracing::debug!("created EFI System Partition {}", part.display());
    disk::format_fat32(&part)?;
    tracing::debug!("formatted FAT32 on {}", part.display());

    let guard = disk::mount_partition(&part)?;
    let dest = guard.point();
    tracing::debug!("mounted {} at {}", part.display(), dest.display());

    let copy_rep = crate::progress::Reporter::spinner("copying files");
    iso::copy_tree(iso_file, root, "", dest, &skip_install_wim, &copy_rep)?;
    copy_rep.finish();

    if let Some(node) = iso::find_install_wim(root) {
        handle_install_wim_on_device(lib, iso_file, node, dest, staging_dir)?;
    } else {
        tracing::warn!("no sources/install.wim found; nothing to split or copy");
    }

    util::sync_all();
    drop(guard);
    tracing::debug!("done: {} is a bootable Windows install USB", part.display());
    Ok(())
}

fn run_image(
    lib: &WimLib,
    iso_file: &mut fs::File,
    root: &isomage::TreeNode,
    output: &Path,
    staging_dir: &Path,
) -> Result<()> {
    let contents_size = iso::total_file_size(root)?;
    tracing::debug!(
        "creating image {} for {} of files",
        output.display(),
        util::human_size(contents_size)
    );
    let image = disk::create_image(output, contents_size)?;

    {
        let dest = image.root_dir();
        let copy_rep = crate::progress::Reporter::spinner("copying files");
        iso::copy_tree_to_fat(iso_file, root, "", &dest, &skip_install_wim, &copy_rep)?;
        copy_rep.finish();
    }

    if let Some(node) = iso::find_install_wim(root) {
        let dest = image.root_dir();
        handle_install_wim_in_image(lib, iso_file, node, &dest, staging_dir)?;
    } else {
        tracing::warn!("no sources/install.wim found; nothing to split or copy");
    }

    image.finish()?;
    util::sync_all();
    tracing::debug!(
        "done: {} is a bootable Windows installation image",
        output.display()
    );
    Ok(())
}

fn skip_install_wim(rel: &str) -> bool {
    rel.eq_ignore_ascii_case("sources/install.wim")
}

/// Place `sources/install.wim` on a mounted target, splitting it into
/// `install.swm` parts when it exceeds the FAT32 per-file limit.
fn handle_install_wim_on_device(
    lib: &WimLib,
    iso_file: &mut fs::File,
    node: &isomage::TreeNode,
    dest: &Path,
    staging_dir: &Path,
) -> Result<()> {
    let sources = dest.join("sources");
    fs::create_dir_all(&sources).with_context(|| format!("mkdir {}", sources.display()))?;

    if node.size > FAT32_MAX_FILE {
        tracing::debug!(
            "install.wim is {} (>4 GiB); splitting into .swm parts",
            util::human_size(node.size)
        );
        split_install_wim_to_device(lib, iso_file, node, &sources, staging_dir)
    } else {
        let output = sources.join("install.wim");
        tracing::debug!(
            "copying install.wim ({} bytes) to {}",
            node.size,
            output.display()
        );
        let mut writer =
            fs::File::create(&output).with_context(|| format!("create {}", output.display()))?;
        copy_install_wim(iso_file, node, &mut writer, "install.wim")?;
        writer.sync_all()?;
        Ok(())
    }
}

/// Place `sources/install.wim` directly in an image's FAT filesystem.
fn handle_install_wim_in_image(
    lib: &WimLib,
    iso_file: &mut fs::File,
    node: &isomage::TreeNode,
    dest: &disk::FatDir<'_>,
    staging_dir: &Path,
) -> Result<()> {
    if node.size > FAT32_MAX_FILE {
        tracing::debug!(
            "install.wim is {} (>4 GiB); splitting into .swm parts",
            util::human_size(node.size)
        );
        split_install_wim_to_image(lib, iso_file, node, dest, staging_dir)
    } else {
        tracing::debug!("copying install.wim ({} bytes) into image", node.size);
        let mut writer = disk::create_image_file(dest, "sources/install.wim")
            .context("create sources/install.wim in image")?;
        copy_install_wim(iso_file, node, &mut writer, "install.wim")?;
        writer
            .flush()
            .context("flush sources/install.wim in image")?;
        Ok(())
    }
}

fn copy_install_wim<W: Write>(
    iso_file: &mut fs::File,
    node: &isomage::TreeNode,
    writer: &mut W,
    message: &str,
) -> Result<()> {
    let rep = crate::progress::Reporter::bytes(message, node.size);
    let mut writer = crate::progress::ProgressWrite::new(writer, rep);
    iso::cat(iso_file, node, &mut writer)?;
    writer.finish();
    Ok(())
}

/// Extract `install.wim` to a temporary file to make it available to libwim.
fn extract_install_wim(
    iso_file: &mut fs::File,
    node: &isomage::TreeNode,
    staging_dir: &Path,
) -> Result<(tempfile::TempDir, PathBuf)> {
    let tmpdir = create_install_wim_tempdir(staging_dir)?;
    let wim_path = tmpdir.path().join("install.wim");
    tracing::debug!(
        "extracting install.wim ({} bytes) to {}",
        util::human_size(node.size),
        wim_path.display()
    );
    let mut writer =
        fs::File::create(&wim_path).with_context(|| format!("create {}", wim_path.display()))?;
    copy_install_wim(iso_file, node, &mut writer, "extract install.wim")?;
    writer.sync_all()?;
    Ok((tmpdir, wim_path))
}

fn create_install_wim_tempdir(staging_dir: &Path) -> Result<tempfile::TempDir> {
    match tempfile::tempdir_in(staging_dir) {
        Ok(tmpdir) => Ok(tmpdir),
        Err(error) => {
            tracing::warn!(
                "cannot create install.wim temp dir in {}: {error}; using system temp dir",
                staging_dir.display()
            );
            tempfile::tempdir().context("create temp dir for install.wim")
        }
    }
}

fn split_install_wim_to_device(
    lib: &WimLib,
    iso_file: &mut fs::File,
    node: &isomage::TreeNode,
    sources: &Path,
    staging_dir: &Path,
) -> Result<()> {
    let (_tmpdir, wim_path) = extract_install_wim(iso_file, node, staging_dir)?;
    split_wim(lib, &wim_path, &sources.join("install.swm"))
}

fn split_install_wim_to_image(
    lib: &WimLib,
    iso_file: &mut fs::File,
    node: &isomage::TreeNode,
    dest: &disk::FatDir<'_>,
    staging_dir: &Path,
) -> Result<()> {
    let (tmpdir, wim_path) = extract_install_wim(iso_file, node, staging_dir)?;
    let split_dir = tmpdir.path().join("split");
    fs::create_dir(&split_dir).with_context(|| format!("mkdir {}", split_dir.display()))?;
    split_wim(lib, &wim_path, &split_dir.join("install.swm"))?;
    copy_split_wims_to_image(&split_dir, dest)
}

fn split_wim(lib: &WimLib, wim_path: &Path, swm_path: &Path) -> Result<()> {
    tracing::debug!(
        "splitting into {} ({} per part)",
        swm_path.display(),
        util::human_size(SWM_PART_SIZE)
    );
    let wim = Wim::open(lib, wim_path)?;
    let split_rep = crate::progress::Reporter::spinner("splitting WIM");
    wim.split(swm_path, SWM_PART_SIZE)?;
    split_rep.finish();
    Ok(())
}

fn copy_split_wims_to_image(split_dir: &Path, dest: &disk::FatDir<'_>) -> Result<()> {
    let mut paths = fs::read_dir(split_dir)
        .with_context(|| format!("read {}", split_dir.display()))?
        .map(|entry| entry.map(|entry| entry.path()))
        .collect::<io::Result<Vec<_>>>()
        .with_context(|| format!("read {}", split_dir.display()))?;
    paths.sort();

    for path in paths {
        if !path.is_file() {
            continue;
        }
        let name = path
            .file_name()
            .and_then(|name| name.to_str())
            .ok_or_else(|| anyhow!("split WIM path {} is not valid UTF-8", path.display()))?;
        tracing::debug!("copying {name} into image");
        let mut writer = disk::create_image_file(dest, &format!("sources/{name}"))
            .with_context(|| format!("create sources/{name} in image"))?;
        copy_host_file(&path, &mut writer, name)?;
        writer
            .flush()
            .with_context(|| format!("flush sources/{name} in image"))?;
    }

    Ok(())
}

fn copy_host_file<W: Write>(source: &Path, writer: &mut W, message: &str) -> Result<()> {
    let mut input = fs::File::open(source).with_context(|| format!("open {}", source.display()))?;
    let size = input
        .metadata()
        .with_context(|| format!("stat {}", source.display()))?
        .len();
    let rep = crate::progress::Reporter::bytes(message, size);
    let mut writer = crate::progress::ProgressWrite::new(writer, rep);
    io::copy(&mut input, &mut writer).with_context(|| format!("copy {}", source.display()))?;
    writer.finish();
    Ok(())
}