mod dev;

use {
    core::fmt::Write as _,
    derive_more::Display,
    dev::Package,
    error_stack::{bail, report, Context, Report, Result, ResultExt},
    fatfs::{FileSystem, FormatVolumeOptions, FsOptions, ReadWriteSeek},
    std::{
        fs::{self, File},
        io::{self, Write},
        path::Path,
        process::{exit, Command, Stdio},
    },
    toml::Value,
};

fn main() -> Result<(), Error> {
    let mut args = std::env::args();
    args.next();

    log::set_logger(&hblang::Logger).unwrap();
    log::set_max_level(log::LevelFilter::Error);

    match args.next().as_deref() {
        Some("build" | "b") => {
            let mut release = false;
            let mut debuginfo = false;
            let mut target = Target::X86_64;
            let mut tests = false;
            for arg in args {
                if arg == "-r" || arg == "--release" {
                    release = true;
                } else if arg == "-d" || arg == "--debuginfo" {
                    debuginfo = true;
                } else if arg == "rv64" || arg == "riscv64" || arg == "riscv64-virt" {
                    target = Target::Riscv64Virt;
                } else if arg == "arm64" || arg == "aarch64" || arg == "aarch64-virt" {
                    target = Target::Aarch64;
                } else if arg == "avx2" {
                    target = Target::X86_64Avx2;
                } else if arg == "--ktest" {
                    tests = true;
                } else {
                    return Err(report!(Error::InvalidSubCom));
                }
            }

            build(release, target, debuginfo, tests).change_context(Error::Build)
        }
        // Some("test" | "t") => {
        //     let mut release = false;
        //     let mut debuginfo = false;
        //     let mut target = Target::X86_64;
        //     for arg in args {
        //         if arg == "-r" || arg == "--release" {
        //             release = true;
        //         } else if arg == "-d" || arg == "--debuginfo" {
        //             debuginfo = true;
        //         } else if arg == "rv64" || arg == "riscv64" || arg == "riscv64-virt" {
        //             target = Target::Riscv64Virt;
        //         } else if arg == "arm64" || arg == "aarch64" || arg == "aarch64-virt" {
        //             target = Target::Aarch64;
        //         } else if arg == "avx2" {
        //             target = Target::X86_64Avx2;
        //         } else {
        //             return Err(report!(Error::InvalidSubCom));
        //         }
        //     }

        //     test(release, target, debuginfo).change_context(Error::Build)
        // }
        Some("run" | "r") => {
            let mut release = false;
            let mut debuginfo = false;
            let mut target = Target::X86_64;
            let mut tests = false;
            let mut do_accel = true;
            for arg in args {
                if arg == "-r" || arg == "--release" {
                    release = true;
                } else if arg == "-d" || arg == "--debuginfo" {
                    debuginfo = true;
                } else if arg == "rv64" || arg == "riscv64" || arg == "riscv64-virt" {
                    target = Target::Riscv64Virt;
                } else if arg == "arm64" || arg == "aarch64" || arg == "aarch64-virt" {
                    target = Target::Aarch64;
                } else if arg == "--noaccel" {
                    do_accel = false;
                } else if arg == "avx2" {
                    target = Target::X86_64Avx2;
                } else if arg == "--ktest" {
                    tests = true; 
                } else {
                    return Err(report!(Error::InvalidSubCom));
                }
            }

            build(release, target, debuginfo, tests)?;
            run(release, target, do_accel)
        }
        Some("help" | "h") => {
            println!(concat!(
                "AbleOS RepBuild\n",
                "Subcommands:\n",
                "  build (b): Build a bootable disk image\n",
                "   help (h): Print this message\n",
                "    run (r): Build and run AbleOS in QEMU\n\n",
                "Options for build and run:\n",
                "        -r / --release: build in release mode\n",
                "        -d / --debuginfo: build with debug info\n",
                "        --noaccel: run without acceleration (e.g, no kvm)\n",
                "        --ktest: Enables tests via ktest\n",
                "[ rv64 / riscv64 / riscv64-virt / aarch64 / arm64 / aarch64-virt / avx2 ]: sets target"
            ),);
            Ok(())
        }
        _ => Err(report!(Error::InvalidSubCom)),
    }
}

fn get_path_without_boot_prefix(val: &Value) -> Option<&str> {
    val.as_str()?.split("boot:///").last()
}

fn get_fs() -> Result<FileSystem<impl ReadWriteSeek>, io::Error> {
    let filename = "sysdata/system_config.toml";

    let mut img = File::options()
        .read(true)
        .write(true)
        .create(true)
        .open(Path::new("target/disk.img"))?;

    img.set_len(1024 * 1024 * 64)?;

    fatfs::format_volume(&mut img, FormatVolumeOptions::new())?;

    let fs = FileSystem::new(img, FsOptions::new())?;

    // Read the contents of the file using a `match` block
    // to return the `data: Ok(c)` as a `String`
    // or handle any `errors: Err(_)`.
    let contents = match fs::read_to_string(filename) {
        // If successful return the files text as `contents`.
        // `c` is a local variable.
        Ok(c) => c,
        // Handle the `error` case.
        Err(_) => {
            // Write `msg` to `stderr`.
            eprintln!("Could not read file `{}`", filename);
            // Exit the program with exit code `1`.
            exit(1);
        }
    };
    use toml::Value;

    let mut limine_str = String::new();

    let mut data: Value = toml::from_str(&contents).unwrap();
    let boot_table = data.get_mut("boot");
    let limine_table = boot_table.unwrap().get_mut("limine").unwrap();
    let default_entry = limine_table.get("default_entry").unwrap();
    let timeout = limine_table.get("timeout").unwrap();
    let interface_resolution = limine_table.get("interface_resolution").unwrap();
    let verbose = limine_table.get("verbose").unwrap();
    let term_wallpaper = limine_table.get("term_wallpaper").unwrap();

    let vb_post = match verbose.as_bool().unwrap() {
        true => "yes",
        false => "no",
    };
    let term_background = limine_table
        .get("term_backdrop")
        .unwrap_or(&Value::Integer(0));
    let base = format!(
        "DEFAULT_ENTRY={}
TIMEOUT={}
VERBOSE={}
INTERFACE_RESOLUTION={}
# Terminal related settings
TERM_WALLPAPER={}
TERM_BACKDROP={}
",
        default_entry,
        timeout,
        vb_post,
        interface_resolution,
        term_wallpaper.as_str().unwrap(),
        term_background
    );

    // Copy the term_wallpaper to the image
    let term_wallpaper_path = get_path_without_boot_prefix(term_wallpaper).unwrap();
    copy_file_to_img(&format!("sysdata/{}", term_wallpaper_path), &fs);

    limine_str.push_str(&base);

    let boot_entries = limine_table.as_table_mut().unwrap();
    let mut real_boot_entries = boot_entries.clone();
    for (key, value) in boot_entries.into_iter() {
        if !value.is_table() {
            real_boot_entries.remove(key);
        }
    }

    for (name, mut value) in real_boot_entries {
        let comment = value.get("comment").unwrap();
        let protocol = value.get("protocol").unwrap();
        let resolution = value.get("resolution").unwrap();
        let kernel_path = value.get("kernel_path").unwrap();
        let kernel_cmdline = value.get("kernel_cmdline").unwrap();

        let text_entry = format!(
            "
:{}
    COMMENT={}
    PROTOCOL={}
    RESOLUTION={}
    KERNEL_PATH={}
    KERNEL_CMDLINE={}

",
            name,
            comment.as_str().unwrap(),
            protocol.as_str().unwrap(),
            resolution.as_str().unwrap(),
            kernel_path.as_str().unwrap(),
            kernel_cmdline,
        );

        limine_str.push_str(&text_entry);

        let modules = value.get_mut("modules").unwrap().as_table_mut().unwrap();
        // let mut real_modules = modules.clone();

        let mut errors = String::new();
        let mut out = Vec::new();

        modules
            .into_iter()
            .map(|(_, value)| -> Result<(), io::Error> {
                if value.is_table() {
                    let path = get_path_without_boot_prefix(
                        value.get("path").expect("You must have `path` as a value"),
                    )
                    .unwrap()
                    .split(".")
                    .next()
                    .unwrap();
                    let p = Package::load_from_file(
                        format!("sysdata/programs/{}/meta.toml", path).to_owned(),
                    );
                    match p.build(&mut out) {
                        Ok(()) => {}
                        Err(_) => {
                            writeln!(errors, "========= while compiling {} =========", path)
                                .unwrap();
                            errors.push_str(core::str::from_utf8(&out).expect("no"));
                            out.clear();
                        }
                    }
                }
                Ok(())
            })
            .for_each(drop);

        if !errors.is_empty() {
            let _ = writeln!(errors, "!!! STOPPING DUE TO PREVIOUS ERRORS !!!");
            std::eprint!("{errors}");
            continue;
        }

        modules.into_iter().for_each(|(_key, value)| {
            if value.is_table() {
                let path = value.get("path").expect("You must have `path` as a value");
                let default_value = Value::String("".into());
                let cmd_line = value.get("cmd_line").unwrap_or(&default_value);

                let a = format!(
                    "    MODULE_PATH={}
    MODULE_CMDLINE={}\n\n",
                    path.as_str().unwrap(),
                    cmd_line
                );
                limine_str.push_str(&a);
            }
        });

        // Copy modules into the test_programs directory
        modules.into_iter().for_each(|(_key, value)| {
            if value.is_table() {
                let path = get_path_without_boot_prefix(
                    value
                        .get("path")
                        .expect("You must have a `path` as a value"),
                )
                .unwrap();
                let fpath = format!("target/programs/{}", path);
                copy_file_to_img(&fpath, &fs);
            }
        });
    }

    let bootdir = fs.root_dir().create_dir("efi")?.create_dir("boot")?;

    let mut f = fs.root_dir().create_file("limine.cfg")?;
    let _ = f.write(limine_str.as_bytes())?;
    drop(f);

    io::copy(
        &mut File::open("limine/BOOTX64.EFI")
            .map_err(Report::from)
            .attach_printable("Copying Limine (x86_64): have you pulled the submodule?")?,
        &mut bootdir.create_file("bootx64.efi")?,
    )?;

    io::copy(
        &mut File::open("limine/BOOTAA64.EFI")
            .map_err(Report::from)
            .attach_printable("Copying Limine (ARM): have you pulled the submodule?")?,
        &mut bootdir.create_file("bootaa64.efi")?,
    )?;

    drop(bootdir);
    Ok(fs)
}

fn copy_file_to_img(fpath: &str, fs: &FileSystem<File>) {
    let path = Path::new(fpath);
    // println!("{path:?}");
    io::copy(
        &mut File::open(path).expect(&format!("Could not open file {fpath}")),
        &mut fs
            .root_dir()
            .create_file(&path.file_name().unwrap().to_string_lossy())
            .expect("Unable to create file"),
    )
    .expect("Copy failed");
}

fn build(release: bool, target: Target, debuginfo: bool, tests: bool) -> Result<(), Error> {
    let fs = get_fs().change_context(Error::Io)?;
    let mut com = Command::new("cargo");
    com.current_dir("kernel");
    com.args(["b"]);
    if release {
        com.arg("-r");
    }
    if debuginfo {
        com.env("RUSTFLAGS", "-Cdebug-assertions=true");
    }

    if tests {
        com.args(["--features", "ktest"]);
    }

    if target == Target::Riscv64Virt {
        com.args(["--target", "targets/riscv64-virt-ableos.json"]);
    }
    if target == Target::Aarch64 {
        com.args(["--target", "targets/aarch64-virt-ableos.json"]);
    }
    if target == Target::X86_64Avx2 {
        com.args(["--target", "targets/x86_64_v3-ableos.json"]);
    }

    match com.status() {
        Ok(s) if s.code() != Some(0) => bail!(Error::Build),
        Err(e) => bail!(report!(e).change_context(Error::Build)),
        _ => (),
    }

    let mut path: String = "kernel".to_string();
    let kernel_dir = match target {
        Target::X86_64 => {
            path.push_str("_x86-64");
            "target/x86_64-ableos"
        }
        Target::X86_64Avx2 => {
            path.push_str("_x86-64");
            "target/x86_64_v3-ableos"
        }
        Target::Riscv64Virt => "target/riscv64-virt-ableos",
        Target::Aarch64 => {
            path.push_str("_aarch64");
            "target/aarch64-virt-ableos"
        }
    };

    (|| -> std::io::Result<_> {
        io::copy(
            &mut File::open(
                Path::new(kernel_dir)
                    .join(if release { "release" } else { "debug" })
                    .join("kernel"),
            )?,
            &mut fs.root_dir().create_file(&path)?,
        )
        .map(|_| ())
    })()
    .map_err(Report::from)
    .change_context(Error::Io)
}

fn run(release: bool, target: Target, do_accel: bool) -> Result<(), Error> {
    let target_str = match target {
        Target::X86_64 | Target::X86_64Avx2 => "qemu-system-x86_64",
        Target::Riscv64Virt => "qemu-system-riscv64",
        Target::Aarch64 => "qemu-system-aarch64",
    };
    let (mut com, mut com2) = (Command::new(target_str), Command::new(target_str));
    let ovmf_path = fetch_ovmf(target);
    #[cfg(target_arch = "x86_64")]
    let accel = if do_accel {
        let supported = String::from_utf8(
            com2.args(["--accel", "help"])
                .stdout(Stdio::piped())
                .stderr(Stdio::piped())
                .spawn()
                .unwrap()
                .wait_with_output()
                .unwrap()
                .stdout,
        )
        .unwrap();

        let cpuid = raw_cpuid::CpuId::new();
        let vmx = cpuid.get_feature_info().unwrap().has_vmx();
        let svm = cpuid.get_svm_info().is_some();

        if supported.contains("kvm") && (vmx || svm) {
            "accel=kvm"
        } else if cpuid
            .get_processor_brand_string()
            .filter(|a| a.as_str() == "GenuineIntel")
            .is_some()
            && supported.contains("hax")
            && vmx
        {
            "accel=hax"
        } else if supported.contains("whpx") {
            "accel=whpx"
        } else {
            "accel=tcg"
        }
    } else {
        "accel=tcg"
    };
    #[cfg(not(target_arch = "x86_64"))]
    let accel = "accel=tcg";

    match target {
        Target::X86_64 | Target::X86_64Avx2 => {
            #[rustfmt::skip]
            com.args([
                "-bios", &ovmf_path.change_context(Error::OvmfFetch)?,
                "-drive", "file=target/disk.img,format=raw",
                "-device", "vmware-svga",
                // "-serial", "stdio",
                "-m", "2G",
                "-smp", "1",
																"-audiodev",
																"pa,id=speaker",
																"-machine",
																"pcspk-audiodev=speaker",
                "-parallel", "none",
                "-monitor", "none",
                "-machine", accel,
                "-cpu", "max",
                "-device", "isa-debug-exit,iobase=0xf4,iosize=0x04",
            ]);
        }
        Target::Riscv64Virt => {
            #[rustfmt::skip]
            com.args([
                "-M", "virt",
                "-m", "128M",
                "-serial", "stdio",
                "-kernel",
                if release {
                    "target/riscv64-virt-ableos/release/kernel"
                } else {
                    "target/riscv64-virt-ableos/debug/kernel"
                }
            ]);
        }
        Target::Aarch64 => {
            #[rustfmt::skip]
            com.args([
                "-M", "virt",
                "-cpu", "max",
                "-device", "ramfb",
                "-device", "qemu-xhci",
                "-device", "usb-kbd",
                "-m", "2G",
                "-bios", &ovmf_path.change_context(Error::OvmfFetch)?,
                "-drive", "file=target/disk.img,format=raw",
            ]);
        }
    }

    match com
        .status()
        .map_err(Report::from)
        .change_context(Error::ProcessSpawn)?
    {
        s if s.success() => Ok(()),
        s => Err(report!(Error::Qemu(s.code()))),
    }
}


fn fetch_ovmf(target: Target) -> Result<String, OvmfFetchError> {
    let (ovmf_url, ovmf_path) = match target {
        Target::X86_64 | Target::X86_64Avx2 => (
            "https://retrage.github.io/edk2-nightly/bin/RELEASEX64_OVMF.fd",
            "target/RELEASEX64_OVMF.fd",
        ),
        Target::Riscv64Virt => return Err(OvmfFetchError::Empty.into()),
        Target::Aarch64 => (
            "https://retrage.github.io/edk2-nightly/bin/RELEASEAARCH64_QEMU_EFI.fd",
            "target/RELEASEAARCH64_QEMU_EFI.fd",
        ),
    };

    let mut file = match std::fs::metadata(ovmf_path) {
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => std::fs::OpenOptions::new()
            .create(true)
            .write(true)
            .read(true)
            .open(ovmf_path)
            .map_err(Report::from)
            .change_context(OvmfFetchError::Io)?,
        Ok(_) => return Ok(ovmf_path.to_owned()),
        Err(e) => return Err(report!(e).change_context(OvmfFetchError::Io)),
    };
    let req = ureq::get(ovmf_url)
        .call()
        .map_err(Report::from)
        .change_context(OvmfFetchError::Fetch)?;

    std::io::copy(&mut req.into_reader(), &mut file)
        .map_err(Report::from)
        .change_context(OvmfFetchError::Io)?;

    Ok(ovmf_path.to_owned())
}

#[derive(Debug, Display)]
enum OvmfFetchError {
    #[display("Failed to fetch OVMF package")]
    Fetch,
    #[display("No OVMF package available")]
    Empty,
    #[display("IO Error")]
    Io,
}

impl Context for OvmfFetchError {}

#[derive(Clone, Copy, PartialEq, Eq)]
enum Target {
    X86_64,
    X86_64Avx2,
    Riscv64Virt,
    Aarch64,
}

#[allow(unused)]
#[derive(Debug, Display)]
enum Error {
    #[display("Failed to build the kernel")]
    Build,
    #[display("Missing or invalid subcommand (available: build, run)")]
    InvalidSubCom,
    #[display("IO Error")]
    Io,
    #[display("Failed to spawn a process")]
    ProcessSpawn,
    #[display("Failed to fetch UEFI firmware")]
    OvmfFetch,
    #[display("Failed to assemble Holey Bytes code")]
    Assembler,
    #[display("QEMU Error: {}", "fmt_qemu_err(*_0)")]
    Qemu(Option<i32>),
}

impl Context for Error {}

#[allow(dead_code)]
fn fmt_qemu_err(e: Option<i32>) -> impl Display {
    struct W(Option<i32>);
    impl Display for W {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            if let Some(c) = self.0 {
                c.fmt(f)
            } else {
                f.write_str("Interrupted by signal")
            }
        }
    }

    W(e)
}