/*
 * Copyright (c) 2022, Umut İnan Erdoğan <umutinanerdogan@pm.me>
 * Copyright (c) 2022, able <abl3theabove@gmail.com>
 *
 * SPDX-License-Identifier: MPL-2.0
 */

use colored::*;
use std::{
    fs::{self, File},
    os::fd::AsRawFd,
    process::Command,
};
use udisks::{
    filesystem::{MountOptions, UnMountOptions},
    manager::LoopSetupOptions,
};

struct Options {
    pub subcommand: Subcommand,
    pub arguments: Vec<String>,
}

enum Subcommand {
    BuildImage,
    Doc,
    Help,
    Run,
    Empty,
    /// Run all tests for all architectures
    Test,
    Unknown(String),
}

impl Subcommand {
    fn from_str<S: AsRef<str>>(str: S) -> Subcommand {
        match str.as_ref() {
            "build-image" => Subcommand::BuildImage,
            "doc" => Subcommand::Doc,
            "help" => Subcommand::Help,
            "run" | "r" => Subcommand::Run,
            "test" | "t" => Subcommand::Test,
            "" => Subcommand::Empty,
            unknown => Subcommand::Unknown(unknown.to_string()),
        }
    }
}

enum MachineType {
    X86_64,
    RiscV64,
    AArch64,
    Unknown(String),
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let options = options();

    match options.subcommand {
        Subcommand::BuildImage => {
            let machine_text = options.arguments.get(0).cloned().unwrap_or_default();

            match machine(machine_text) {
                MachineType::X86_64 => {
                    // Cleanup
                    // NOTE: we are not unwrapping these, as we don't want this to fail if they
                    //       don't exist yet, probably not the best idea tho.
                    // FIXME: figure out a better way to ignore errors about these not existing
                    #[allow(unused_must_use)]
                    {
                        fs::remove_dir_all("./limine");
                        fs::remove_dir_all("./disk");
                        fs::remove_file("./target/disk.img");
                    }

                    // Build ableOS in release mode
                    Command::new("cargo")
                        .args(["build", "--release"])
                        .current_dir(fs::canonicalize("./kernel").unwrap())
                        .status()
                        .unwrap();

                    // Create disk directory
                    fs::create_dir("./disk").unwrap();

                    // Clone limine 4.x binaries
                    Command::new("git")
                        .arg("clone")
                        .arg("https://github.com/limine-bootloader/limine.git")
                        .arg("--branch=v4.x-branch-binary")
                        .arg("--depth=1")
                        .status()
                        .unwrap();

                    println!("{}", "Building limine".bold());
                    Command::new("make")
                        .args(["-C", "limine"])
                        .status()
                        .unwrap();

                    println!("{}", "Allocating new disk image".bold());
                    Command::new("fallocate")
                        .args(["-l", "256M", "./target/disk.img"])
                        .status()
                        .unwrap();

                    println!("{}", "Partitioning disk image".bold());
                    let dbus_conn = zbus::blocking::Connection::system()?;

                    // Setup loop device
                    let disk_img = File::options()
                        .read(true)
                        .write(true)
                        .open("./target/disk.img")?;
                    let loopdev = udisks::manager::UDisks2ManagerProxyBlocking::new(&dbus_conn)?
                        .loop_setup(
                            disk_img.as_raw_fd().into(),
                            LoopSetupOptions {
                                no_user_interaction: true,
                                offset: 0,
                                size: 0,
                                readonly: false,
                                no_part_scan: false,
                            },
                        )?;

                    // Create MBR
                    udisks::block::BlockProxyBlocking::builder(&dbus_conn)
                        .path(&loopdev)?
                        .build()?
                        .format("dos", Default::default())?;

                    // Create and format partition
                    let filesystem =
                        udisks::partition::PartitionTableProxyBlocking::builder(&dbus_conn)
                            .destination("org.freedesktop.UDisks2")?
                            .path(&loopdev)?
                            .build()?
                            .create_partition_and_format(
                                0,
                                0,
                                "",
                                "",
                                Default::default(),
                                "ext2",
                                [("take-ownership", true.into())].into_iter().collect(),
                            )?;

                    let fsproxy = udisks::filesystem::FilesystemProxyBlocking::builder(&dbus_conn)
                        .path(&filesystem)?
                        .build()?;

                    // Mount the filesystem
                    let mountpoint = fsproxy
                        .mount(MountOptions {
                            no_user_interaction: true,
                            fs_type: String::new(),
                            mount_options: String::new(),
                        })
                        .or_else(|_| {
                            Ok::<String, zbus::Error>(loop {
                                if let Some(m) = fsproxy.mount_points()?.get(0) {
                                    break m.to_string();
                                }
                            })
                        })?;

                    // copy ./base/* over to ./disk
                    Command::new("sh")
                        .arg("-c")
                        .arg(format!("cp -r ./base/* {mountpoint}"))
                        .status()?;

                    // copy ./limine/limine.sys over to ./disk/boot
                    Command::new("cp")
                        .args(["./limine/limine.sys", &format!("{mountpoint}/boot")])
                        .status()?;

                    // copy the kernel over to ./disk/boot/kernel
                    Command::new("cp")
                        .arg("./target/x86_64-ableos/release/kernel")
                        .arg(&format!("{mountpoint}/boot/kernel"))
                        .status()?;

                    // Unmount the filesystem (and the rest of things will follow)
                    fsproxy.unmount(UnMountOptions {
                        no_user_interaction: true,
                        force: false,
                    })?;

                    println!("{}", "Deploying limine".bold());
                    Command::new("./limine/limine-deploy")
                        .arg("./target/disk.img")
                        .status()
                        .unwrap();
                }
                MachineType::Unknown(unknown) => {
                    eprintln!(
                        "{}: unknown machine type `{}`",
                        "error".red().bold(),
                        unknown.bold(),
                    );
                    eprintln!("expected one of x86_64, riscv64 or aarch64");
                }
                _ => {
                    eprintln!(
                        "{}: build-image not implemented for this machine type",
                        "error".red().bold(),
                    );
                }
            }
        }
        Subcommand::Test => {
            Command::new("cargo")
                .args(["test", "--target=json_targets/x86_64-ableos.json"])
                .current_dir(fs::canonicalize("./ableos").unwrap())
                .status()
                .unwrap();

            // panic!("Test Infrastructure missing");
        }
        Subcommand::Doc => {
            let machine_text = options.arguments.get(0).cloned().unwrap_or_default();

            match machine(machine_text) {
                MachineType::X86_64 => {
                    Command::new("cargo")
                        .args(["doc", "--open"])
                        .current_dir(fs::canonicalize("./ableos").unwrap())
                        .status()
                        .unwrap();
                }
                MachineType::RiscV64 => {
                    Command::new("cargo")
                        .args(["doc", "--open", "--target=riscv64gc-unknown-none-elf"])
                        .current_dir(fs::canonicalize("./ableos").unwrap())
                        .status()
                        .unwrap();
                }
                MachineType::AArch64 => {
                    Command::new("cargo")
                        .args(["doc", "--open", "--target=json_targets/aarch64-ableos.json"])
                        .current_dir(fs::canonicalize("./ableos").unwrap())
                        .status()
                        .unwrap();
                }
                MachineType::Unknown(unknown) => {
                    eprintln!(
                        "{}: unknown machine type `{}`",
                        "error".red().bold(),
                        unknown.bold(),
                    );
                    eprintln!("expected one of x86_64, riscv64 or aarch64");
                }
            }
        }
        Subcommand::Help => help(),
        Subcommand::Run => {
            let machine_text = options.arguments.get(0).cloned().unwrap_or_default();
            let debug = options.arguments.get(1).cloned().unwrap_or_default();
            let debug = matches!(debug.as_str(), "--debug" | "--dbg" | "-d");

            match machine(machine_text) {
                MachineType::X86_64 if debug => {
                    // Build ableOS
                    Command::new("cargo")
                        .arg("build")
                        .current_dir(fs::canonicalize("./kernel").unwrap())
                        .status()
                        .unwrap();

                    // Setup loopback device for disk.img, with partitions
                    // FIXME: don't do ths if running without changes
                    // Setup loop device
                    let disk_img = File::options()
                        .read(true)
                        .write(true)
                        .open("./target/disk.img")?;
                    let dbus_conn = zbus::blocking::Connection::system()?;
                    let loopdev = udisks::manager::UDisks2ManagerProxyBlocking::new(&dbus_conn)?
                        .loop_setup(
                            disk_img.as_raw_fd().into(),
                            LoopSetupOptions {
                                no_user_interaction: true,
                                offset: 0,
                                size: 0,
                                readonly: false,
                                no_part_scan: false,
                            },
                        )?;

                    let parts = udisks::partition::PartitionTableProxyBlocking::builder(&dbus_conn)
                        .destination("org.freedesktop.UDisks2")?
                        .path(loopdev)?
                        .build()?
                        .partitions()?;

                    let fsobjpath = parts.get(0).ok_or("missing boot partition")?;
                    let mountpoint =
                        udisks::filesystem::FilesystemProxyBlocking::builder(&dbus_conn)
                            .path(fsobjpath)?
                            .build()?
                            .mount(MountOptions {
                                no_user_interaction: true,
                                fs_type: String::new(),
                                mount_options: String::new(),
                            })?;

                    // copy the kernel over to ./disk/boot/kernel
                    Command::new("cp")
                        .arg("./target/x86_64-ableos/debug/kernel")
                        .arg(format!("{mountpoint}/boot/kernel"))
                        .status()
                        .unwrap();

                    udisks::filesystem::FilesystemProxyBlocking::builder(&dbus_conn)
                        .path(fsobjpath)?
                        .build()?
                        .unmount(UnMountOptions {
                            no_user_interaction: true,
                            force: false,
                        })?;

                    // run qemu with "-S", "-gdb", "tcp:9000"
                    Command::new("qemu-system-x86_64")
                        .args(["-device", "piix4-ide,id=ide"])
                        .arg("-drive")
                        .arg("file=./target/disk.img,format=raw,if=none,id=disk")
                        .args(["-device", "ide-hd,drive=disk,bus=ide.0"])
                        // .arg("--nodefaults")
                        .args(["-cpu", "Broadwell-v3"])
                        .args(["-m", "4G"])
                        .args(["-serial", "stdio"])
                        .args(["-smp", "cores=2"])
                        // .args(["-soundhw", "pcspk"])
                        // .args(["-device", "VGA"])
                        // .args(["-device", "virtio-gpu-pci"])
                        .args(["-device", "vmware-svga"])
                        .args(["-device", "sb16"])
                        // .args(["-machine", "pcspk-audiodev=0"])
                        // .args(["-qmp", "unix:../qmp-sock,server,nowait"])
                        .args(["-S", "-gdb", "tcp:9000"])
                        .status()
                        .unwrap();
                }
                MachineType::X86_64 => {
                    // Build ableOS
                    Command::new("cargo")
                        .args(["build", "--release"])
                        .current_dir(fs::canonicalize("./kernel").unwrap())
                        .status()
                        .unwrap();

                    // Setup loopback device for disk.img, with partitions
                    // FIXME: don't do ths if running without changes
                    let disk_img = File::options()
                        .read(true)
                        .write(true)
                        .open("./target/disk.img")?;
                    let dbus_conn = zbus::blocking::Connection::system()?;
                    let loopdev = udisks::manager::UDisks2ManagerProxyBlocking::new(&dbus_conn)?
                        .loop_setup(
                            disk_img.as_raw_fd().into(),
                            LoopSetupOptions {
                                no_user_interaction: true,
                                offset: 0,
                                size: 0,
                                readonly: false,
                                no_part_scan: false,
                            },
                        )?;

                    let parts = udisks::partition::PartitionTableProxyBlocking::builder(&dbus_conn)
                        .destination("org.freedesktop.UDisks2")?
                        .path(loopdev)?
                        .build()?
                        .partitions()?;

                    let fsproxy = udisks::filesystem::FilesystemProxyBlocking::builder(&dbus_conn)
                        .path(&parts[0])?
                        .build()?;

                    // Mount the filesystem
                    let mountpoint = fsproxy
                        .mount(MountOptions {
                            no_user_interaction: true,
                            fs_type: String::new(),
                            mount_options: String::new(),
                        })
                        .or_else(|_| {
                            Ok::<String, zbus::Error>(loop {
                                if let Some(m) = fsproxy.mount_points()?.get(0) {
                                    break m.to_string();
                                }
                            })
                        })?;

                    // copy the kernel over to ./disk/boot/kernel
                    Command::new("cp")
                        .arg("./target/x86_64-ableos/release/kernel")
                        .arg(format!("{mountpoint}/boot/kernel"))
                        .status()
                        .unwrap();

                    fsproxy.unmount(UnMountOptions {
                        no_user_interaction: true,
                        force: false,
                    })?;

                    // run qemu
                    Command::new("qemu-system-x86_64")
                        .args(["-device", "piix4-ide,id=ide"])
                        .arg("-drive")
                        .arg("file=./target/disk.img,format=raw,if=none,id=disk")
                        .args(["-device", "ide-hd,drive=disk,bus=ide.0"])
                        // .arg("--nodefaults")
                        .args(["-cpu", "Broadwell-v3"])
                        .args(["-m", "4G"])
                        .args(["-serial", "stdio"])
                        .args(["-smp", "cores=2"])
                        // .args(["-soundhw", "pcspk"])
                        // .args(["-device", "VGA"])
                        // .args(["-device", "virtio-gpu-pci"])
                        .args(["-device", "vmware-svga"])
                        .args(["-device", "sb16"])
                        // .args(["-machine", "pcspk-audiodev=0"])
                        // .args(["-qmp", "unix:../qmp-sock,server,nowait"])
                        .status()
                        .unwrap();
                }
                MachineType::RiscV64 if debug => {
                    eprintln!(
                        "{}: debug is not implemented for riscv64",
                        "error".red().bold()
                    );
                }
                MachineType::RiscV64 => {
                    Command::new("cargo")
                        .args(["build", "--release", "--target=riscv64gc-unknown-none-elf"])
                        .current_dir(fs::canonicalize("./ableos").unwrap())
                        .status()
                        .unwrap();

                    Command::new("qemu-system-riscv64")
                        .args(["-machine", "virt"])
                        .args(["-cpu", "rv64"])
                        .args(["-smp", "8"])
                        .args(["-m", "128M"])
                        .arg("-bios")
                        .arg("src/arch/riscv/firmwear/opensbi-riscv64-generic-fw_jump.bin")
                        .arg("-kernel")
                        .arg("target/riscv64gc-unknown-none-elf/release/ableos")
                        .current_dir(fs::canonicalize("./ableos").unwrap())
                        .status()
                        .unwrap();
                }
                MachineType::AArch64 if debug => {
                    eprintln!(
                        "{}: debug is not implemented for aarch64",
                        "error".red().bold()
                    );
                }
                MachineType::AArch64 => {
                    Command::new("cargo")
                        .args([
                            "build",
                            "--release",
                            "--target=json_targets/aarch64-ableos.json",
                        ])
                        .current_dir(fs::canonicalize("./ableos").unwrap())
                        .status()
                        .unwrap();

                    Command::new("qemu-system-aarch64")
                        .args(["-machine", "virt"])
                        .args(["-m", "1024M"])
                        .args(["-cpu", "cortex-a53"])
                        .args(["-kernel", "target/aarch64-ableos/release/ableos"])
                        .args(["-device", "virtio-keyboard"])
                        .current_dir(fs::canonicalize("./ableos").unwrap())
                        .status()
                        .unwrap();
                }
                MachineType::Unknown(unknown) => {
                    eprintln!(
                        "{}: unknown machine type `{}`",
                        "error".red().bold(),
                        unknown.bold(),
                    );
                    eprintln!("expected one of x86_64, riscv64 or aarch64");
                }
            }
        }
        Subcommand::Empty => {
            eprintln!("{}: no subcommand passed", "error".red().bold());
            help();
        }
        Subcommand::Unknown(unknown) => {
            eprintln!(
                "{}: unknown subcommand `{}`",
                "error".red().bold(),
                unknown.bold()
            );
            help();
        }
    }

    Ok(())
}

fn options() -> Options {
    let subcommand = std::env::args().nth(1).unwrap_or_default();
    let arguments = std::env::args().skip(2).collect();

    Options {
        subcommand: Subcommand::from_str(subcommand),
        arguments,
    }
}

fn machine<S: AsRef<str>>(text: S) -> MachineType {
    match text.as_ref() {
        "x86" | "x86_64" => MachineType::X86_64,
        "riscv" | "riscv64" => MachineType::RiscV64,
        "arm" | "arm64" | "aarch64" => MachineType::AArch64,
        "" => {
            eprintln!(
                "{}: no machine type passed, defaulting to x86_64",
                "warning".yellow().bold()
            );
            MachineType::X86_64
        }
        unknown => MachineType::Unknown(unknown.to_string()),
    }
}

fn help() {
    todo!("`help`")
}