use std::{env, ffi::{OsStr, OsString}, fmt, fs, io, path::{Path, PathBuf}, process::{self, Command}};

fn main() {
  let out_dir = env::var_os("OUT_DIR")
    .expect("OUT_DIR not set");
  let out_dir = PathBuf::from(out_dir);

  let search_path = SearchPath::load();
  
  let nasm_path = search_path.search("nasm")
    .next()
    .expect("failed to find nasm in PATH");
  let nasm = Nasm::new(nasm_path);

  build_asm(&nasm, &out_dir);
  
  emit_link_args();  

  rerun_if_changed("build.rs".as_ref());
}

fn emit_link_args() {
  let linker_script_path = "link.ld";
  println!("cargo::rustc-link-arg=-T{}", linker_script_path);
  rerun_if_changed(linker_script_path.as_ref());
}

fn rerun_if_changed(path: &Path) {
  let path_str = path.to_str()
    .expect("expected path to be valid utf8");
  println!("cargo::rerun-if-changed={}", path_str);
}

fn link_obj(path: &Path) {
  let path_str = path.to_str()
    .expect("expected path to be valid utf8");
  println!("cargo::rustc-link-arg={}", path_str);
}

fn build_asm(nasm: &Nasm, out_dir: &Path) {
  let asm_srcs = fs::read_dir("asm")
    .expect("failed to get asm sources");

  for asm_src in asm_srcs {
    let asm_src = asm_src.expect("failed to get asm source");
    
    let ty = asm_src.file_type().expect("failed to get file type");
    if !ty.is_file() {
      continue;
    }

    let src_path = asm_src.path();

    let is_asm = src_path
      .extension()
      .map(|ext| ext == "s")
      .unwrap_or(false);
    if !is_asm {
      continue;
    }

    let out_filename = src_path
      .file_name()
      .unwrap()
      .apply(PathBuf::from)
      .with_extension("o");
    let out_path = out_dir.join(out_filename);

    nasm.assemble(
      &out_path,
      &[&src_path],
      &["../include".as_ref()]
    ).expect("failed to assemble");

    link_obj(&out_path);
  }
}

struct Nasm {
  bin_path: PathBuf,
}

impl Nasm {
  fn new(bin_path: PathBuf) -> Self {
    Self { bin_path }
  }

  fn assemble(&self, output: &Path, sources: &[&Path], includes: &[&Path]) -> Result<(), CmdError>
  {
    for source in sources {
      rerun_if_changed(source);
    }

    let mut cmd = Command::new(&self.bin_path);
    cmd
      .arg("-Werror")
      .arg("-f")
      .arg("elf");

    for include in includes {
      let mut buf = OsString::new();
      buf.push("-I");
      buf.push(include);
      cmd.arg(buf);
    }

    cmd
      .arg("-o")
      .arg(output)
      .args(sources);

    run_cmd(&mut cmd)?;

    Ok(())
  }
}

struct SearchPath {
  paths: Vec<PathBuf>
}

impl SearchPath {
  fn load() -> Self {
    let path_var = env::var_os("PATH").unwrap_or_default();
    let paths = env::split_paths(&path_var).collect();
    Self { paths }
  }
  
  fn search<'a, 'b, 'c, T>(&'a self, bin: &'b T) -> impl Iterator<Item = PathBuf> + 'c
  where
    'a: 'c,
    'b: 'c,
    T: AsRef<OsStr> + ?Sized,
  {
    let bin = bin.as_ref();
    self.paths
      .iter()
      .filter_map(move |path| {
        let path = path.join(bin);
        fs::metadata(&path)
          .ok()
          .and_then(|meta| if meta.is_file() || meta.is_symlink() {
            Some(path)
          } else {
            None
          })
      })
  }
}

fn run_cmd(cmd: &mut Command) -> Result<process::Output, CmdError> {
  use fmt::Write;

  cmd
      .output()
      .map_err(CmdErrorKind::Io)
      .and_then(|out| if out.status.success() {
        Ok(out)
      } else {
        Err(CmdErrorKind::Status(out))
      })
      .map_err(|err| {
        let mut cmd_buf = String::new();
        write!(&mut cmd_buf, "{:?}", cmd).ok();
        CmdError::new(cmd_buf, err)
      })
}

#[derive(Debug)]
struct CmdError {
  cmd: String,
  kind: CmdErrorKind,
}

impl CmdError {
  fn new(cmd: String, kind: CmdErrorKind) -> Self {
    Self { cmd, kind }
  }
}

impl fmt::Display for CmdError {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "`{}` failed: ", self.cmd)?;
    match &self.kind {
      CmdErrorKind::Io(err) => err.fmt(f),
      CmdErrorKind::Status(out) => write!(f, "exited with status {}", out.status),
    }
  }
}

#[derive(Debug)]
enum CmdErrorKind {
  Io(io::Error),
  Status(process::Output),
}

trait Apply: Sized {
  fn apply<T, F>(self, f: F) -> T
  where
    F: FnOnce(Self) -> T;
}

impl<U> Apply for U
where
  U: Sized,
{
  fn apply<T, F>(self, f: F) -> T
  where
    F: FnOnce(Self) -> T
  {
    f(self)
  }
}