use std::{
    ffi::OsStr,
    fs,
    path::{Path, PathBuf},
    process::{Command, Output, Stdio},
};

use eyre::{WrapErr, eyre};

use crate::Context;

const NASM_COMMON_FLAGS: &[&str] = &["-werror"];

pub fn build(ctx: &Context) -> Result<(), eyre::Error> {
    let stage_1_src_dir = ctx.workspace.join("stage_1");
    let stage_2_src_dir = ctx.workspace.join("stage_2");
    let stage_3_src_dir = ctx.workspace.join("stage_3");

    let build_dir = ctx.workspace.join("build");
    let stage_1_build_dir = build_dir.join("stage_1");
    let stage_2_build_dir = build_dir.join("stage_2");

    mkdir_if_missing(&build_dir)?;
    mkdir_if_missing(&stage_1_build_dir)?;
    mkdir_if_missing(&stage_2_build_dir)?;

    build_stage_1(ctx, &stage_1_build_dir, &stage_1_src_dir)?;
    build_stage_2(ctx, &stage_2_build_dir, &stage_2_src_dir)?;
    build_stage_3(ctx, &stage_3_src_dir)?;

    Ok(())
}

fn build_stage_1(ctx: &Context, build_dir: &Path, src_dir: &Path) -> Result<(), eyre::Error> {
    let src_paths = ls_with_extension(src_dir, "s")?;

    let [src_path] = &*src_paths else {
        return Err(eyre!(
            "expected exactly one stage 1 source file, found {}",
            src_paths.len()
        ));
    };

    let out_path = build_dir.join("s1.bin");
    
    let include_dir = ctx.workspace.join("include");

    Command::new("nasm")
        .args(["-f", "bin"])
        .args(NASM_COMMON_FLAGS)
        .arg("-I")
        .arg(&include_dir)
        .arg("-o")
        .arg(&out_path)
        .arg(src_path)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .run_ok()?;

    Ok(())
}

fn build_stage_2(ctx: &Context, build_dir: &Path, src_dir: &Path) -> Result<(), eyre::Error> {
    let include_dir = ctx.workspace.join("include");
    let src_paths = ls_with_extension(src_dir, "s")?;
    let mut obj_paths = Vec::new();

    for src_path in src_paths {
        let Some(file_stem) = src_path.file_stem() else {
            continue;
        };
        
        let obj_path = build_dir.join(file_stem).with_extension("o");
        
        Command::new("nasm")
            .args(["-f", "elf"])
            .args(NASM_COMMON_FLAGS)
            .arg("-I")
            .arg(&include_dir)
            .arg("-o")
            .arg(&obj_path)
            .arg(src_path)
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .run_ok()?;

        obj_paths.push(obj_path);
    }

    let linker_file_path = src_dir.join("link.ld");
    let bin_path = build_dir.join("s2.bin");
    
    Command::new("ld")
        .args(["-m", "elf_i386"])
        .arg("-T")
        .arg(&linker_file_path)
        .arg("-o")
        .arg(bin_path)
        .args(&obj_paths)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .run_ok()?;
    
    Ok(())
}

fn build_stage_3(ctx: &Context, stage_3_src_dir: &Path) -> Result<(), eyre::Error> {
    Command::new(ctx.cargo)
        .arg("build")
        .arg("--release")
        .current_dir(stage_3_src_dir)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .run_ok()?;

    Ok(())
}

fn mkdir_if_missing(path: &Path) -> Result<(), eyre::Error> {
    let exists = fs::exists(path)
        .wrap_err_with(|| format!("failed to check if {} exists", path.to_string_lossy()))?;

    if !exists {
        fs::create_dir(path)
            .wrap_err_with(|| format!("failed to create {}", path.to_string_lossy()))?;
    }

    Ok(())
}

fn ls_with_extension(path: &Path, ext: &str) -> Result<Vec<PathBuf>, eyre::Error> {
    let read_dir = fs::read_dir(path)
        .wrap_err_with(|| format!("failed to read {}", path.to_string_lossy()))?;

    let mut paths = Vec::new();

    for dir in read_dir {
        let dir = dir.wrap_err_with(|| format!("failed to read {}", path.to_string_lossy()))?;

        if dir.path().extension() == Some(OsStr::new(ext)) {
            paths.push(dir.path());
        }
    }

    Ok(paths)
}

trait CommandExt {
    fn run_ok(&mut self) -> Result<Output, eyre::Error>;
}

impl CommandExt for Command {
    fn run_ok(&mut self) -> Result<Output, eyre::Error> {
        let output = self
            .output()
            .wrap_err_with(|| format!("failed to run {:?}", self))?;

        if !output.status.success() {
            return Err(eyre!("command failed: {:?}", self));
        }

        Ok(output)
    }
}