use std::env;
use std::fs::File;
use std::io::{Read, Write};
use std::path::Path;

fn main() {
    let args: Vec<String> = env::args().collect();

    let iso_mount_dir = Path::new(args.get(1)
        .expect("Arguments: rsfus (path to XP ISO mount) [bootloader disk number] [installer disk number]"))
        .canonicalize().expect("Can't canonicalize path to XP ISO mount.");

    if !iso_mount_dir.exists() || !iso_mount_dir.is_dir() {
        panic!("XP ISO mount is not a directory.");
    }

    let bootloader_disk_number = if let Some(i) = args.get(2) {
        i.parse::<u8>().expect("Not a valid number for disk number!")
    } else { 1 };
    let installer_disk_number = if let Some(i) = args.get(3) {
        i.parse::<u8>().expect("Not a valid number for disk number!")
    } else { 1 };
    let setup_source_line = format!("SetupSourceDevice = \"\\device\\harddisk{installer_disk_number}\\partition1\"");

    let amd64_basedir = iso_mount_dir.join("amd64");
    let i386_basedir = iso_mount_dir.join("i386");

    let x64_edition = amd64_basedir.is_dir();

    {
        let mut ntdetect = File::open(amd64_basedir.join("ntdetect.com"))
            .or(File::open(i386_basedir.join("ntdetect.com")))
            .expect("Can't open file ntdetect.com.");
        let mut ntdetect_copy = File::create(iso_mount_dir.join("ntdetect.com"))
            .expect("Can't create file ntdetect.com.");

        let mut ntdetect_contents = Vec::new();
        println!("Reading {ntdetect:?}");
        ntdetect.read_to_end(&mut ntdetect_contents).expect("Can't read ntdetect.com.");
        println!("Writing {ntdetect_copy:?}");
        ntdetect_copy.write_all(&ntdetect_contents).expect("Can't write to ntdetect.com.");
    }

    {
        let mut txtsetup = File::open(amd64_basedir.join("txtsetup.sif"))
            .or(File::open(i386_basedir.join("txtsetup.sif")))
            .expect("Can't open file txtsetup.sif");
        let mut txtsetup_copy = File::create(iso_mount_dir.join("txtsetup.sif"))
            .expect("Can't create file txtsetup.sif");

        let mut txtsetup_contents = String::new();
        println!("Reading {txtsetup:?}");
        txtsetup.read_to_string(&mut txtsetup_contents).expect("Can't read from txtsetup.sif.");
        let mut txtsetup_lines: Vec<&str> = Vec::new();
        for line in txtsetup_contents.lines() {
            txtsetup_lines.push(line);
            if line == "[SetupData]" {
                println!("TXTSETP: Found [SetupData], writing \"{setup_source_line}\"");
                txtsetup_lines.push(setup_source_line.as_str());
            }
        }
        println!("Writing {txtsetup_copy:?}");
        txtsetup_copy.write_all(txtsetup_lines.join("\r\n").as_bytes())
            .expect("Couldn't write to txtsetup.sif.");
    }

    {
        let mut bootmgr = File::open(amd64_basedir.join("setupldr.bin"))
            .or(File::open(i386_basedir.join("setupldr.bin")))
            .expect("Can't open file setupldr.bin");
        let mut bootmgr_copy = File::create(iso_mount_dir.join("BOOTMGR"))
            .expect("Can't create file BOOTMGR");

        let mut bootmgr_contents = Vec::new();
        println!("Reading {bootmgr:?}");
        bootmgr.read_to_end(&mut bootmgr_contents).expect("Can't read from setupldr.bin");

        if bootmgr_contents.len() > 0x2061
            && bootmgr_contents[0x2060] == 0x74
            && bootmgr_contents[0x2061] == 0x03 {
            bootmgr_contents[0x2060] = 0xeb;
            bootmgr_contents[0x2061] = 0x1a;
            println!("BOOTMGR: 0x00002060  0x74 0x03  -->  0xeb 0x1a  (Disable Win2k3 CRC check)");
        }

        let rdisk0 = b"rdisk(0)".as_slice();
        let win_nt_bt = b"$win_nt$.~bt".as_slice();
        let win_nt_bt_caps = b"$WIN_NT$.~BT".as_slice();
        let new_dir_bytes = b"i386";
        let new_dir_bytes_x64 = b"amd64";

        for index in 0..(bootmgr_contents.len() - win_nt_bt.len()) {
            if &bootmgr_contents[index..index + rdisk0.len()] == rdisk0 {
                println!("BOOTMGR: {index:#010x}  Patching rdisk(0)  -->  rdisk({bootloader_disk_number})");
                bootmgr_contents[index + 6] = 0x30 + bootloader_disk_number;
            }

            if &bootmgr_contents[index..index + win_nt_bt.len()] == win_nt_bt
                || &bootmgr_contents[index..index + win_nt_bt.len()] == win_nt_bt_caps {
                if x64_edition {
                    println!("BOOTMGR: {index:#010x}  Patching {:?}  -->  \"amd64\"", bootmgr_contents[index..index + win_nt_bt.len()].to_vec());
                    for i in 0..new_dir_bytes_x64.len() {
                        bootmgr_contents[index + i] = new_dir_bytes_x64[i];
                    }
                    bootmgr_contents[index + new_dir_bytes_x64.len()] = 0;
                    bootmgr_contents[index + new_dir_bytes_x64.len()] = bootmgr_contents[index + win_nt_bt.len()];
                    bootmgr_contents[index + new_dir_bytes_x64.len() + 1] = 0;
                } else {
                    println!("BOOTMGR: {index:#010x}  Patching {:?}  -->  \"i386\"", bootmgr_contents[index..index + win_nt_bt.len()].to_vec());
                    for i in 0..new_dir_bytes.len() {
                        bootmgr_contents[index + i] = new_dir_bytes[i];
                    }
                    bootmgr_contents[index + new_dir_bytes.len()] = 0;
                    bootmgr_contents[index + new_dir_bytes.len()] = bootmgr_contents[index + win_nt_bt.len()];
                    bootmgr_contents[index + new_dir_bytes.len() + 1] = 0;
                }
            }
        }

        println!("Writing {bootmgr_copy:?}");
        bootmgr_copy.write_all(bootmgr_contents.as_slice()).expect("Couldn't write to BOOTMGR");
    }

    println!("Done!  :D");
}