use std::path::PathBuf;
use std::{env, ffi::OsStr, io, io::Write, path::Path, process::Command};

// We have to build stage 0 completely separately because it won't be part of
// the final binary per se (it's what loads the final binary in the first place).
// That's why we have this not-at-all overengineered custom build system.

fn main() {
    let make = Make::from_env();

    let stage0_srcs = ["src/boot/stage0.s"];
    let stage0_objs = make.all(stage0_srcs);
    make.link(Some("config/stage0.ld"), stage0_objs, "stage0.elf");
    make.objcopy("stage0.elf", "stage0.bin");

    make.end();
}

struct Make {
    /// Base directory for all source files
    src_base: String,
    /// Base directory for all generated object files
    obj_base: String,
    /// Base directory for finished, pre-linked binaries
    root_dir: String,
    /// Command to use for assembling
    cmd_as: String,
    /// Command to use for linking
    cmd_ld: String,
}

impl Make {
    pub fn from_env() -> Self {
        let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
        let out_dir = env::var("OUT_DIR").unwrap();
        // XXX this will probably be removed in the final version, i'm only throwing
        //     the finished binaries in the target root directory because that's easier
        //     to access from gdb and various helper scripts to set up a disk image
        let root_dir = Path::new(&out_dir)
            .parent()
            .and_then(Path::parent)
            .and_then(Path::parent)
            .and_then(Path::to_str)
            .map(String::from)
            .unwrap();

        eprintln!("=== BussyBuild starting ===");

        Self {
            src_base: manifest_dir,
            obj_base: out_dir,
            root_dir,
            cmd_as: env::var("AS").unwrap_or_else(|_| String::from("as")),
            cmd_ld: env::var("LD")
                .or_else(|_| env::var("RUSTC_LINKER"))
                .unwrap_or_else(|_| "ld".to_owned()),
        }
    }

    /// Build the specified source file (relative to `CARGO_MANIFEST_DIR`)
    /// and return the absolute path to the object file.
    /// The program used for building depends on the file extension:
    ///
    /// | extension | program    |
    /// |:---------:|:-----------|
    /// |   `*.c`   | C compiler |
    /// |   `*.s`   | assembler  |
    pub fn make(&self, src_file: &str) -> String {
        assert!(Path::new(src_file).is_relative());

        let src_path = Path::new(&self.src_base).join(src_file);
        let obj_path = Path::new(&self.obj_base).join(format!("{src_file}.o"));

        println!("cargo:rerun-if-changed={}", src_path.to_str().unwrap());

        let obj_dir = obj_path.parent().expect("output file has no parent??");
        if !obj_dir.exists() {
            std::fs::create_dir_all(obj_dir).unwrap();
        }
        if !obj_dir.is_dir() {
            panic!("invalid output path");
        }

        let extension = src_path
            .extension()
            .expect("source file has no extension")
            .to_str()
            .expect("invalid (non-Unicode?) source file extension");

        match extension {
            "c" => unimplemented!(),
            "s" => self.assemble(src_path.as_os_str(), obj_path.as_os_str()),
            _ => panic!("unsupported source file extension \".{extension}\""),
        }

        String::from(obj_path.to_str().unwrap())
    }

    pub fn all<I, S>(&self, src_files: I) -> Vec<String>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        src_files
            .into_iter()
            .map(|f| self.make(f.as_ref()))
            .collect()
    }

    pub fn link<I, S>(&self, script: Option<&str>, obj_files: I, out: &str)
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        eprintln!("=> LD {out}");

        let mut cmd = Command::new(&self.cmd_ld);

        if let Some(script) = script {
            let script_path = Path::new(script);
            let script_path = if script_path.is_relative() {
                Path::new(&self.src_base).join(script_path)
            } else {
                script_path.to_path_buf()
            };
            let script = script_path.as_os_str().to_str().unwrap();
            cmd.arg(format!("-T{script}"));
            println!("cargo:rerun-if-changed={script}");
        }

        cmd.arg("-melf_i386")
            .arg("-static")
            .arg("-nostdlib")
            .arg("-o")
            .arg(Path::new(&self.root_dir).join(out));

        for file in obj_files.into_iter() {
            cmd.arg(file.as_ref());
        }

        eprintln!("  {cmd:?}");

        let output = cmd.output().expect("failed to execute linker");
        self.stdout_raw(output.stdout.as_ref());
        if !output.status.success() {
            self.stdout_raw(output.stderr.as_ref());
            self.cmd_failed(&self.cmd_ld, output.status.code());
        }
    }

    pub fn objcopy(&self, elf: &str, out: &str) {
        eprintln!("=> OBJCOPY {elf}");
        let elf = Path::new(&self.root_dir).join(elf);
        let bin = Path::new(&self.root_dir).join(out);
        let output = Command::new("objcopy")
            .arg("-O")
            .arg("binary")
            .arg(&elf)
            .arg(&bin)
            .output()
            .unwrap();
        self.stdout_raw(output.stdout.as_ref());
        if !output.status.success() {
            self.stdout_raw(output.stderr.as_ref());
            self.cmd_failed(&self.cmd_ld, output.status.code());
        }
    }

    pub fn end(self) {
        eprintln!("=== BussyBuild finished ===");
    }

    /// Assemble a source file to an object file.
    fn assemble(&self, src_file: &OsStr, obj_file: &OsStr) {
        eprintln!("=> AS {}", obj_file.to_str().unwrap());

        let mut cmd = Command::new(&self.cmd_as);
        cmd.arg("--32")
            // Without this option, the linker mysteriously appends exactly 40 bytes
            // of nonsense to the final output binary, even after stripping the ELF.
            // I have no explanation for this, other than the GNU people being sadists.
            .arg("-mx86-used-note=no")
            .arg("-g")
            .arg("-o")
            .arg(obj_file)
            .arg(src_file);

        eprintln!("  {cmd:?}");

        let output = cmd.output().expect("failed to execute assembler");
        let stderr = String::from_utf8(output.stderr).unwrap();
        for line in stderr.lines() {
            if !line.is_empty() {
                println!("cargo:warning={line}");
            }
        }

        if !output.status.success() {
            self.cmd_failed(&self.cmd_as, output.status.code());
        }
    }

    fn cmd_failed(&self, cmd: &str, code: Option<i32>) -> ! {
        let code = code
            .map(|code| format!("{code}"))
            .unwrap_or_else(|| "<unknown>".to_owned());
        panic!("{cmd} exited with status code {code}",)
    }

    fn stdout_raw(&self, buf: &[u8]) {
        io::stdout().write_all(buf).unwrap();
    }
}