chdlady-disc 0.2.2

//! Hard disk CHD image creation pipeline.

use super::common::DEFAULT_HD_CODECS;
use crate::error::DiscError;
use crate::hd::{guess_chs, HardDiskGeo, HD_TEMPLATES};
use crate::types::{HARD_DISK_IDENT_METADATA_TAG, HARD_DISK_METADATA_TAG};
use chdlady_core::{
    ChdFile, ChdWriteConfig, ChdWriter, CreateResult, MetadataItem, ProgressStatus,
    CHD_MDFLAGS_CHECKSUM,
};
use std::fs::File;
use std::io::{self, BufReader, BufWriter, Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};

use std::sync::atomic::AtomicBool;
use std::sync::Arc;

/// Options for creating a hard disk CHD image.
#[derive(Debug, Clone, Default)]
pub struct CreateHdOptions {
    /// Starting byte offset in the input file.
    pub start_byte: Option<u64>,
    /// Starting hunk index in the input file.
    pub start_hunk: Option<u64>,
    /// Total bytes to write.
    pub length_bytes: Option<u64>,
    /// Total hunks to write.
    pub length_hunks: Option<u64>,
    /// Hunk size in bytes.
    pub hunk_size: Option<u32>,
    /// Sector size in bytes (defaults to 512).
    pub sector_size: Option<u32>,
    /// Explicit physical CHS geometry: (cylinders, heads, sectors).
    pub chs: Option<(u32, u32, u32)>,
    /// Explicit capacity in bytes for blank drive.
    pub size: Option<u64>,
    /// Predefined template index (0 to 16).
    pub template: Option<usize>,
    /// Path to an ATA/IDE identify binary payload file.
    pub ident_path: Option<PathBuf>,
    /// Up to 4 FourCC compression codecs.
    pub codecs: Option<[u32; 4]>,
    /// Optional path to parent CHD for differential container.
    pub parent_path: Option<PathBuf>,
    /// Optional limit on worker threads for parallel compression.
    pub num_processors: Option<usize>,
    /// Optional shared atomic flag for cooperative cancellation.
    pub cancel_token: Option<Arc<AtomicBool>>,
}

/// Creates a CHD v5 container representing a hard disk image.
pub fn create_hd<F>(
    input_path: Option<&Path>,
    output_path: &Path,
    options: &CreateHdOptions,
    mut progress: F,
) -> Result<CreateResult, DiscError>
where
    F: FnMut(ProgressStatus),
{
    let mut sector_size = options.sector_size.unwrap_or(512);
    if sector_size == 0 {
        return Err(DiscError::InvalidMetadata(
            "sector size must be non-zero".into(),
        ));
    }
    let mut cylinders = 0u32;
    let mut heads = 0u32;
    let mut sectors = 0u32;

    if let Some(t_idx) = options.template {
        if t_idx >= HD_TEMPLATES.len() {
            return Err(DiscError::InvalidMetadata(format!(
                "invalid hard disk template index {}",
                t_idx
            )));
        }
        let t = &HD_TEMPLATES[t_idx];
        cylinders = t.cylinders;
        heads = t.heads;
        sectors = t.sectors;
        sector_size = t.sector_size;
    } else if let Some((c, h, s)) = options.chs {
        cylinders = c;
        heads = h;
        sectors = s;
    }

    let mut ident_data = Vec::new();
    if let Some(ref id_path) = options.ident_path {
        ident_data = std::fs::read(id_path)?;
        if ident_data.len() < 14 {
            return Err(DiscError::InvalidMetadata(
                "ident file too short (< 14 bytes)".into(),
            ));
        }
        if cylinders == 0 {
            let c = u16::from_le_bytes([ident_data[2], ident_data[3]]) as u32;
            let h = u16::from_le_bytes([ident_data[6], ident_data[7]]) as u32;
            let s = u16::from_le_bytes([ident_data[12], ident_data[13]]) as u32;
            if (c as u64 * h as u64 * s as u64) < 16_514_064 {
                cylinders = c;
                heads = h;
                sectors = s;
            }
        }
    }

    let hunk_size = options
        .hunk_size
        .unwrap_or_else(|| ((4096 / sector_size) * sector_size).max(sector_size));

    let (mut reader, logical_bytes): (Box<dyn Read>, u64) = if let Some(in_path) = input_path {
        let mut file = File::open(in_path)?;
        let total_file_size = file.metadata()?.len();
        let start = match (options.start_byte, options.start_hunk) {
            (Some(b), _) => b,
            (None, Some(h)) => h * hunk_size as u64,
            (None, None) => 0,
        };
        let length = match (options.length_bytes, options.length_hunks) {
            (Some(b), _) => b,
            (None, Some(h)) => h * hunk_size as u64,
            (None, None) => total_file_size.saturating_sub(start),
        };
        let end = start
            .checked_add(length)
            .ok_or_else(|| DiscError::InvalidMetadata("specified range overflows u64".into()))?;
        if end > total_file_size {
            return Err(DiscError::InvalidMetadata(
                "specified range exceeds input file size".into(),
            ));
        }
        if cylinders == 0 {
            let geo = guess_chs(length, sector_size)?;
            cylinders = geo.cylinders;
            heads = geo.heads;
            sectors = geo.sectors;
        }
        let totalsectors = cylinders as u64 * heads as u64 * sectors as u64;
        let expected_bytes = totalsectors * sector_size as u64;
        file.seek(SeekFrom::Start(start))?;
        let reader = BufReader::with_capacity(1024 * 1024, file.take(expected_bytes));
        (Box::new(reader), expected_bytes)
    } else {
        let capacity = if cylinders != 0 {
            (cylinders as u64 * heads as u64 * sectors as u64) * sector_size as u64
        } else if let Some(sz) = options.size {
            let geo = guess_chs(sz, sector_size)?;
            cylinders = geo.cylinders;
            heads = geo.heads;
            sectors = geo.sectors;
            (cylinders as u64 * heads as u64 * sectors as u64) * sector_size as u64
        } else {
            return Err(DiscError::InvalidMetadata(
                "length or CHS geometry must be specified for blank hard disk".into(),
            ));
        };
        (Box::new(io::repeat(0).take(capacity)), capacity)
    };

    let compressors = if input_path.is_none() {
        [0, 0, 0, 0]
    } else {
        options.codecs.unwrap_or(DEFAULT_HD_CODECS)
    };

    let (parent_map, parent_sha1) = if let Some(ref p) = options.parent_path {
        let pf = File::open(p)?;
        let mut parent_chd = ChdFile::open(pf)?;
        let p_sha1 = parent_chd.header().sha1;
        let map = ChdWriter::build_parent_map(&mut parent_chd, hunk_size, sector_size)?;
        (Some(map), p_sha1)
    } else {
        (None, [0u8; 20])
    };

    let config = ChdWriteConfig {
        logical_bytes,
        hunk_bytes: hunk_size,
        unit_bytes: sector_size,
        compressors,
        parent_sha1,
        num_processors: options.num_processors,
        cancel_token: options.cancel_token.clone(),
    };

    let mut metadata = Vec::new();
    let geo = HardDiskGeo::new(cylinders, heads, sectors, sector_size);
    let mut gdat_val = geo.to_metadata_string().into_bytes();
    gdat_val.push(0); // trailing null byte
    metadata.push(MetadataItem {
        tag: HARD_DISK_METADATA_TAG,
        flags: CHD_MDFLAGS_CHECKSUM,
        value: gdat_val,
    });

    if !ident_data.is_empty() {
        metadata.push(MetadataItem {
            tag: HARD_DISK_IDENT_METADATA_TAG,
            flags: CHD_MDFLAGS_CHECKSUM,
            value: ident_data,
        });
    }

    let output_file = File::create(output_path)?;
    let mut writer = BufWriter::new(output_file);

    let result = ChdWriter::write_chd_with_parent(
        &mut reader,
        &mut writer,
        &config,
        &metadata,
        parent_map.as_ref(),
        &mut progress,
    );
    match result {
        Ok(res) => Ok(res),
        Err(e) => {
            let _ = std::fs::remove_file(output_path);
            Err(DiscError::Chd(e))
        }
    }
}