use std::{
    collections::VecDeque,
    fs::File,
    io::{Cursor, Read, Seek, Write},
    process::{Command, Stdio},
};

use fatfs::{
    FileSystem, FormatVolumeOptions, FsOptions, IoBase, IoError, Read as FatRead, Seek as FatSeek,
    Write as FatWrite,
};
use mbrs::{Mbr, PartInfo, PartType};

const RELEASE: bool = true;
const VOLUME_LABEL: [u8; 11] = *b"EFIGIFEDEMO";

#[derive(Copy, Clone)]
enum Cpu {
    X86_64,
    AArch64,
}
impl std::fmt::Display for Cpu {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            Cpu::X86_64 => "x86_64",
            Cpu::AArch64 => "aarch64",
        })
    }
}

fn main() {
    let mut args: VecDeque<_> = std::env::args().skip(1).collect();

    match args.pop_front() {
        Some(cmd) => {
            let cpu = match args
                .pop_front()
                .unwrap_or(std::env::consts::ARCH.to_string())
                .to_lowercase()
                .as_str()
            {
                "aarch64" | "arm64" | "aa64" => Cpu::AArch64,
                "x86_64" | "x64" => Cpu::X86_64,
                x => unimplemented!("target cpu {:?}", x),
            };

            match cmd.as_str() {
                "build" => do_build(cpu, args),
                "run" => do_run(cpu, args),
                x => eprintln!("unknown command {:?}", x),
            }
        }
        None => eprintln!("usage: cargo xtask [build|run] [x86_64|aarch64]"),
    }
}

struct MyStdIoWrapper<T: Read + Write + Seek>(T);
#[derive(Debug)]
struct MyStdIoError(std::io::Error);
impl IoError for MyStdIoError {
    fn is_interrupted(&self) -> bool {
        self.0.kind() == std::io::ErrorKind::Interrupted
    }

    fn new_unexpected_eof_error() -> Self {
        Self(std::io::Error::new(
            std::io::ErrorKind::UnexpectedEof,
            "failed to fill whole buffer",
        ))
    }

    fn new_write_zero_error() -> Self {
        Self(std::io::Error::new(
            std::io::ErrorKind::WriteZero,
            "failed to write whole buffer",
        ))
    }
}
impl<T: Read + Write + Seek> IoBase for MyStdIoWrapper<T> {
    type Error = MyStdIoError;
}
impl<T: Read + Write + Seek> FatSeek for MyStdIoWrapper<T> {
    fn seek(&mut self, pos: fatfs::SeekFrom) -> Result<u64, Self::Error> {
        let pos = match pos {
            fatfs::SeekFrom::Start(x) => std::io::SeekFrom::Start(x),
            fatfs::SeekFrom::End(x) => std::io::SeekFrom::End(x),
            fatfs::SeekFrom::Current(x) => std::io::SeekFrom::Current(x),
        };
        Seek::seek(&mut self.0, pos).map_err(MyStdIoError)
    }
}
impl<T: Read + Write + Seek> FatRead for MyStdIoWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {
        Read::read(&mut self.0, buf).map_err(MyStdIoError)
    }
}
impl<T: Read + Write + Seek> FatWrite for MyStdIoWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
        Write::write(&mut self.0, buf).map_err(MyStdIoError)
    }

    fn flush(&mut self) -> Result<(), Self::Error> {
        Write::flush(&mut self.0).map_err(MyStdIoError)
    }
}

fn do_build(cpu: Cpu, _args: VecDeque<String>) {
    let mut cmd = Command::new(env!("CARGO"));
    cmd.stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .arg("build")
        .arg("--bin")
        .arg("efigife")
        .arg("--target")
        .arg(format!("{}-unknown-uefi", cpu));
    if RELEASE {
        cmd.arg("--release");
    }
    let build_success = cmd.output().unwrap().status.success();
    assert!(build_success);

    let efi_bin_path = format!(
        "{}/../target/{}-unknown-uefi/{}/efigife.efi",
        env!("CARGO_MANIFEST_DIR"),
        cpu,
        if RELEASE { "release" } else { "debug" }
    );
    let mut efi_bin_buf = Vec::new();
    let mut efi_bin_file = File::open(efi_bin_path).unwrap();
    let efi_bin_size = efi_bin_file.read_to_end(&mut efi_bin_buf);
    drop(efi_bin_file);

    let fat_disk_size = 64 * 1024; // todo!("get size from efi_bin_size");

    let mut fat_disk_buf = vec![0; fat_disk_size];
    let mut fat_disk = MyStdIoWrapper(Cursor::new(&mut fat_disk_buf));
    fatfs::format_volume(
        &mut fat_disk,
        FormatVolumeOptions::new()
            .fat_type(fatfs::FatType::Fat12)
            .volume_label(VOLUME_LABEL),
    )
    .unwrap();
    let fat_fs = FileSystem::new(fat_disk, FsOptions::new()).unwrap();
    let boot_efi_name = match cpu {
        Cpu::AArch64 => "BOOTAA64.EFI",
        Cpu::X86_64 => "BOOTX64.EFI",
    };

    let mut boot_efi = fat_fs
        .root_dir()
        .create_dir("EFI")
        .unwrap()
        .create_dir("BOOT")
        .unwrap()
        .create_file(boot_efi_name)
        .unwrap();

    FatWrite::write_all(&mut boot_efi, &efi_bin_buf).unwrap();
    drop(boot_efi);
    fat_fs.unmount().unwrap();

    // TODO: do the gif format juggling to add the comment in this xtask
    let src_gif_path = concat!(env!("CARGO_MANIFEST_DIR"), "/../efi-bin/src/bad-apple.gif");
    let mut src_gif_file = File::open(src_gif_path).unwrap();
    let mut buf = Vec::new();
    let src_gif_size = src_gif_file.read_to_end(&mut buf).unwrap();
    drop(src_gif_file);

    let fat_start_sector = (src_gif_size + 511) / 512;
    let fat_sector_count = (fat_disk_buf.len() + 511) / 512;
    buf.extend(std::iter::repeat(0).take(fat_start_sector * 512 - buf.len()));
    buf.extend(fat_disk_buf);

    let mut mbr = Mbr::default();
    mbr.partition_table.entries[0] = Some(
        PartInfo::try_from_lba(
            true,
            fat_start_sector as u32,
            fat_sector_count as u32,
            PartType::Efi,
        )
        .unwrap(),
    );
    let fake_mbr_buf = <[u8; 512]>::try_from(&mbr).unwrap();
    buf[0x1be..0x200].copy_from_slice(&fake_mbr_buf[0x1be..0x200]);

    let out_img_path = format!(
        "{}/../target/bad-apple.{}-efi.raw.gif",
        env!("CARGO_MANIFEST_DIR"),
        cpu
    );
    let mut out_img_file = File::create(out_img_path).unwrap();
    out_img_file.write_all(&buf).unwrap();
}

fn do_run(cpu: Cpu, args: VecDeque<String>) {
    do_build(cpu, args.clone());

    let mut cmd = Command::new(format!("qemu-system-{}", cpu));
    let kvm = cpu.to_string() == std::env::consts::ARCH;
    match cpu {
        Cpu::AArch64 => {
            if kvm {
                cmd.arg("-M")
                    .arg("virt,pflash0=code,accel=kvm,gic-version=3")
                    .arg("-cpu")
                    .arg("host");
            } else {
                cmd.arg("-M").arg("virt,pflash0=code,gic-version=3");
            }
            cmd.arg("-drive")
                .arg("if=none,id=code,format=raw,file=/usr/share/qemu/edk2-aarch64-code.fd,readonly=on")
                .arg("-device")
                .arg("qemu-xhci")
                .arg("-usb")
                .arg("-device")
                .arg("usb-kbd");
        }
        Cpu::X86_64 => {
            if kvm {
                cmd.arg("-M")
                    .arg("q35,smm=on,accel=kvm")
                    .arg("-cpu")
                    .arg("host");
            } else {
                cmd.arg("-M").arg("q35,smm=on");
            }
            cmd.arg("-global")
                .arg("ICH9-LPC.disable_s3=1")
                .arg("-drive")
                .arg("if=pflash,format=raw,unit=0,file=/usr/share/edk2-ovmf/OVMF_CODE.fd,readonly=on");
        }
    }
    cmd.arg("-m")
        .arg("512")
        .arg("-device")
        .arg("virtio-gpu-pci")
        .arg("-serial")
        .arg("stdio")
        .args(&args);
    let out_img_path = format!(
        "{}/../target/bad-apple.{}-efi.raw.gif",
        env!("CARGO_MANIFEST_DIR"),
        cpu
    );
    cmd.arg("-drive")
        .arg(format!("format=raw,file={}", out_img_path))
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .output()
        .unwrap();
}
