use clap::Parser;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};

mod c_includes;
mod udfread;

use udfread::UDF_BLOCK_SIZE;

/// Replaces contents of files in a UDF image with new data of the same size
#[derive(Parser)]
struct Args {
    /// UDF image to update
    udf_image: PathBuf,

    /// Directory containing files with updated data
    input_dir: PathBuf,

    /// Paths of the files to update, relative to image root and input_dir
    #[clap(required = true)]
    paths: Vec<PathBuf>,
}

fn process_paths<'a>(
    mut udf_image: &File,
    udf_reader: &udfread::UdfRead<impl Read + Seek>,
    input_dir: &Path,
    paths: impl Iterator<Item = &'a Path>,
    dry_run: bool,
) -> bool {
    for path in paths {
        let utf8_path = match path.to_str() {
            Some(path) => path,
            None => {
                eprintln!(
                    "Relative path contains non-UTF-8 characters: {}",
                    path.to_string_lossy()
                );
                return false;
            }
        };

        let udf_file = match udf_reader.open_file(utf8_path) {
            Some(file) => file,
            None => {
                eprintln!("Failed to open {} from UDF image", utf8_path);
                return false;
            }
        };

        let udf_size = match udf_file.size() {
            Some(size) => size,
            None => {
                eprintln!("Failed to get UDF file size of {}", utf8_path);
                return false;
            }
        };

        let fs_path = input_dir.join(path);

        let mut fs_file = match File::open(&fs_path) {
            Ok(file) => file,
            Err(err) => {
                eprintln!("Failed to open {}: {}", fs_path.to_string_lossy(), err);
                return false;
            }
        };

        let fs_size = match fs_file.seek(SeekFrom::End(0)) {
            Ok(size) => size,
            Err(err) => {
                eprintln!(
                    "Failed to get file size of {}: {}",
                    fs_path.to_string_lossy(),
                    err
                );
                return false;
            }
        };

        if udf_size != fs_size {
            eprintln!(
                "Size mismatch for {} ({} vs {})",
                utf8_path, fs_size, udf_size
            );
            return false;
        }

        if udf_size == 0 {
            continue;
        }

        let blocks: u32 = ((udf_size - 1) / u64::from(UDF_BLOCK_SIZE) + 1)
            .try_into()
            .expect("UDF file too large!");

        let mut bytes_remaining = udf_size;

        for block in 0..blocks {
            let image_block = match udf_file.block_lba(block) {
                Some(image_block) => image_block,
                None => {
                    eprintln!("Unsupported: file inline or sparse: {}", utf8_path);
                    return false;
                }
            };

            let file_offset = u64::from(block) * u64::from(UDF_BLOCK_SIZE);
            let image_offset = u64::from(image_block) * u64::from(UDF_BLOCK_SIZE);

            let mut data = [0u8; UDF_BLOCK_SIZE as usize];
            let bytes = if bytes_remaining > u64::from(UDF_BLOCK_SIZE) {
                UDF_BLOCK_SIZE as usize
            } else {
                bytes_remaining as usize
            };

            if !dry_run {
                if let Err(err) = fs_file
                    .seek(SeekFrom::Start(file_offset))
                    .and_then(|_| fs_file.read_exact(&mut data[0..bytes]))
                {
                    eprintln!("Error reading from {}: {}", fs_path.to_string_lossy(), err);
                    return false;
                }

                if let Err(err) = udf_image
                    .seek(SeekFrom::Start(image_offset))
                    .and_then(|_| udf_image.write_all(&data[0..bytes]))
                {
                    eprintln!("Error writing to UDF image: {}", err);
                    return false;
                }
            }

            bytes_remaining -= bytes as u64;
        }
    }
    return true;
}

fn main() -> Result<(), &'static str> {
    let args = Args::parse();

    let file = File::options()
        .read(true)
        .write(true)
        .open(args.udf_image)
        .expect("Failed to open UDF file");

    let udf_reader =
        udfread::UdfRead::open_image_reader(&file).expect("Failed to create UDF reader");

    if !process_paths(
        &file,
        &udf_reader,
        &args.input_dir,
        (&args.paths).into_iter().map(AsRef::as_ref),
        true,
    ) {
        Err("Dry run failed. UDF image is unchanged.")?;
    }

    if !process_paths(
        &file,
        &udf_reader,
        &args.input_dir,
        (&args.paths).into_iter().map(AsRef::as_ref),
        false,
    ) {
        Err("Replace failed after successful dry run. Image likely in an inconsistant state.")?;
    }

    Ok(())
}