#![allow(non_upper_case_globals)]
#![allow(clippy::too_many_arguments)]

use crate::codegen;
use crate::errors::*;
use crate::utils;
use std::path::Path;
use std::process::Stdio;
use tokio::io::AsyncWriteExt;
use tokio::process::{Child, ChildStdin, Command};

const RUSTC_BINARY: &str = utils::compile_env!("SH4D0WUP_RUSTC_BINARY", "rustc");

const SIGCHLD: u64 = 17;

const __NR_write: u64 = 1;
const __NR_close: u64 = 3;
const __NR_pipe: u64 = 22;
const __NR_dup2: u64 = 33;
const __NR_fork: u64 = 57;
const __NR_execve: u64 = 59;
const __NR_exit: u64 = 60;
const __NR_wait4: u64 = 61;

pub async fn stream_bin_std(orig: &[u8], stdin: &mut ChildStdin) -> Result<()> {
    debug!("Passing through binary...");
    let mut buf = String::new();
    for chunk in orig.chunks(2048) {
        buf.clear();
        codegen::escape(chunk, &mut buf)?;
        stdin.write_all(b"if f.write_all(b\"").await?;
        stdin.write_all(buf.as_bytes()).await?;
        stdin.write_all(b"\").is_err() { exit(1) };\n").await?;
    }
    Ok(())
}

pub async fn stream_bin_nostd(orig: &[u8], stdin: &mut ChildStdin) -> Result<()> {
    debug!("Passing through binary...");
    let mut buf = String::new();
    for chunk in orig.chunks(2048) {
        buf.clear();
        codegen::escape(chunk, &mut buf)?;
        stdin.write_all(b"if write_all(f, b\"").await?;
        stdin.write_all(buf.as_bytes()).await?;
        stdin.write_all(b"\".as_ptr(), ").await?;
        stdin.write_all(chunk.len().to_string().as_bytes()).await?;
        stdin.write_all(b") == -1 { exit(1) }\n").await?;
    }
    Ok(())
}

pub struct Compiler {
    child: Child,
    pub stdin: ChildStdin,
}

impl Compiler {
    pub async fn spawn(out: &Path, target: Option<&str>) -> Result<Self> {
        let target = target.unwrap_or("x86_64-unknown-linux-musl");
        info!("Spawning Rust compiler...");
        let mut cmd = Command::new(RUSTC_BINARY);
        cmd.arg("-Copt-level=3")
            .arg("-Cpanic=abort")
            .arg("-Cstrip=symbols")
            .arg(format!("--target={target}"))
            .arg("-o")
            .arg(out)
            .arg("-")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped());
        debug!(
            "Setting up process: {:?} {:?}",
            cmd.as_std().get_program(),
            cmd.as_std().get_args()
        );
        let mut child = cmd
            .spawn()
            .with_context(|| anyhow!("Failed to spawn Rust compiler: {RUSTC_BINARY:?}"))?;

        let stdin = child.stdin.take().unwrap();
        let compiler = Compiler { child, stdin };

        Ok(compiler)
    }

    pub async fn add_line(&mut self, line: &str) -> Result<()> {
        debug!("Sending to compiler: {:?}", line);
        self.stdin.write_all(line.as_bytes()).await?;
        Ok(())
    }

    pub async fn add_lines(&mut self, lines: &[&str]) -> Result<()> {
        for line in lines {
            self.add_line(line).await?;
        }
        Ok(())
    }

    pub fn done(self) -> PendingCompile {
        PendingCompile { child: self.child }
    }

    pub async fn syscall0_readonly(&mut self, ret: &str, nr: u64) -> Result<()> {
        self.add_lines(&[
            "unsafe {\n",
            &format!("let r0: {ret};\n"),
            "asm!(\"syscall\",\n",
            &format!("inlateout(\"rax\") {nr}{ret} => r0,\n"),
            "lateout(\"rcx\") _,\n",
            "lateout(\"r11\") _,\n",
            "options(nostack, preserves_flags)\n",
            ");\n",
            "r0\n",
            "}\n",
        ])
        .await?;
        Ok(())
    }

    pub async fn syscall1_readonly(&mut self, ret: &str, nr: u64, a0: &str) -> Result<()> {
        self.add_lines(&[
            "unsafe {\n",
            &format!("let r0: {ret};\n"),
            "asm!(\"syscall\",\n",
            &format!("inlateout(\"rax\") {nr}{ret} => r0,\n"),
            &format!("in(\"rdi\") {a0},\n"),
            "lateout(\"rcx\") _,\n",
            "lateout(\"r11\") _,\n",
            "options(nostack, preserves_flags)\n",
            ");\n",
            "r0\n",
            "}\n",
        ])
        .await?;
        Ok(())
    }

    pub async fn syscall2_readonly(
        &mut self,
        ret: &str,
        nr: u64,
        a0: &str,
        a1: &str,
    ) -> Result<()> {
        self.add_lines(&[
            "unsafe {\n",
            &format!("let r0: {ret};\n"),
            "asm!(\"syscall\",\n",
            &format!("inlateout(\"rax\") {nr}{ret} => r0,\n"),
            &format!("in(\"rdi\") {a0},\n"),
            &format!("in(\"rsi\") {a1},\n"),
            "lateout(\"rcx\") _,\n",
            "lateout(\"r11\") _,\n",
            "options(nostack, preserves_flags)\n",
            ");\n",
            "r0\n",
            "}\n",
        ])
        .await?;
        Ok(())
    }

    pub async fn syscall3_readonly(
        &mut self,
        ret: &str,
        nr: u64,
        a0: &str,
        a1: &str,
        a2: &str,
    ) -> Result<()> {
        self.add_lines(&[
            "unsafe {\n",
            &format!("let r0: {ret};\n"),
            "asm!(\"syscall\",\n",
            &format!("inlateout(\"rax\") {nr}{ret} => r0,\n"),
            &format!("in(\"rdi\") {a0},\n"),
            &format!("in(\"rsi\") {a1},\n"),
            &format!("in(\"rdx\") {a2},\n"),
            "lateout(\"rcx\") _,\n",
            "lateout(\"r11\") _,\n",
            "options(nostack, preserves_flags)\n",
            ");\n",
            "r0\n",
            "}\n",
        ])
        .await?;
        Ok(())
    }

    pub async fn syscall4_readonly(
        &mut self,
        ret: &str,
        nr: u64,
        a0: &str,
        a1: &str,
        a2: &str,
        a3: &str,
    ) -> Result<()> {
        self.add_lines(&[
            "unsafe {\n",
            &format!("let r0: {ret};\n"),
            "asm!(\"syscall\",\n",
            &format!("inlateout(\"rax\") {nr}{ret} => r0,\n"),
            &format!("in(\"rdi\") {a0},\n"),
            &format!("in(\"rsi\") {a1},\n"),
            &format!("in(\"rdx\") {a2},\n"),
            &format!("in(\"r10\") {a3},\n"),
            "lateout(\"rcx\") _,\n",
            "lateout(\"r11\") _,\n",
            "options(nostack, preserves_flags)\n",
            ");\n",
            "r0\n",
            "}\n",
        ])
        .await?;
        Ok(())
    }

    pub async fn syscall5_readonly(
        &mut self,
        ret: &str,
        nr: u64,
        a0: &str,
        a1: &str,
        a2: &str,
        a3: &str,
        a4: &str,
    ) -> Result<()> {
        self.add_lines(&[
            "unsafe {\n",
            &format!("let r0: {ret};\n"),
            "asm!(\"syscall\",\n",
            &format!("inlateout(\"rax\") {nr}{ret} => r0,\n"),
            &format!("in(\"rdi\") {a0},\n"),
            &format!("in(\"rsi\") {a1},\n"),
            &format!("in(\"rdx\") {a2},\n"),
            &format!("in(\"r10\") {a3},\n"),
            &format!("in(\"r8\") {a4},\n"),
            "lateout(\"rcx\") _,\n",
            "lateout(\"r11\") _,\n",
            "options(nostack, preserves_flags)\n",
            ");\n",
            "r0\n",
            "}\n",
        ])
        .await?;
        Ok(())
    }

    pub async fn generate_entrypoint(&mut self) -> Result<()> {
        self.add_lines(&[
            "global_asm!(",
            "\".global _start\",",
            "\"_start:\", \"mov rdi, rsp\", \"call main\"",
            ");\n",
        ])
        .await
    }

    pub async fn generate_panic_handler(&mut self) -> Result<()> {
        self.add_lines(&[
            "#[panic_handler]\n",
            "fn panic(__info: &core::panic::PanicInfo) -> ! {\n",
            "exit(1)\n",
            "}\n",
        ])
        .await
    }

    pub async fn generate_exit_fn(&mut self) -> Result<()> {
        self
        .add_lines(&[
            "fn exit(code: i32) -> ! {\n",
            &format!("unsafe {{ asm!(\"syscall\", in(\"rax\") {}, in(\"rdi\") code, options(noreturn)) }}\n", __NR_exit),
            "}\n",
        ])
        .await
    }

    pub async fn generate_execve_fn(&mut self) -> Result<()> {
        self.add_line(
            "fn execve(prog: *const u8, argv: *const *const u8, envp: *const *const u8) -> u64 {\n",
        )
        .await?;
        self.syscall3_readonly("u64", __NR_execve, "prog", "argv", "envp")
            .await?;
        self.add_line("}\n").await?;
        Ok(())
    }

    pub async fn generate_fork_fn(&mut self) -> Result<()> {
        self.add_line("fn fork() -> i32 {\n").await?;
        let zero = "ptr::null_mut::<u64>()";
        self.syscall5_readonly(
            "i32",
            __NR_fork,
            &SIGCHLD.to_string(),
            zero,
            zero,
            zero,
            zero,
        )
        .await?;
        self.add_line("}\n").await?;
        Ok(())
    }

    pub async fn generate_wait4_fn(&mut self) -> Result<()> {
        self.add_line("fn wait4(pid: i32, status: *const i32, options: i32, rusage: *const core::ffi::c_void) -> i32 {\n").await?;
        self.syscall4_readonly("i32", __NR_wait4, "pid", "status", "options", "rusage")
            .await?;
        self.add_line("}\n").await?;
        Ok(())
    }

    pub async fn generate_pipe_fn(&mut self) -> Result<()> {
        self.add_line("fn pipe(pipefd: *const i32) -> i32 {\n")
            .await?;
        self.syscall1_readonly("i32", __NR_pipe, "pipefd").await?;
        self.add_line("}\n").await?;
        Ok(())
    }

    pub async fn generate_close_fn(&mut self) -> Result<()> {
        self.add_line("fn close(fd: i32) -> i32 {\n").await?;
        self.syscall1_readonly("i32", __NR_close, "fd").await?;
        self.add_line("}\n").await?;
        Ok(())
    }

    pub async fn generate_dup2_fn(&mut self) -> Result<()> {
        self.add_line("fn dup2(oldfd: i32, newfd: i32) -> i32 {\n")
            .await?;
        self.syscall2_readonly("i32", __NR_dup2, "oldfd", "newfd")
            .await?;
        self.add_line("}\n").await?;
        Ok(())
    }

    pub async fn generate_write_fn(&mut self) -> Result<()> {
        self.add_line("fn write(fd: i32, buf: *const u8, count: usize) -> isize {\n")
            .await?;
        self.syscall3_readonly("isize", __NR_write, "fd", "buf", "count")
            .await?;
        self.add_line("}\n").await?;
        Ok(())
    }

    pub async fn generate_write_all_fn(&mut self) -> Result<()> {
        self.add_lines(&[
            "fn write_all(fd: i32, mut buf: *const u8, mut count: usize) -> i32 {\n",
            "while count > 0 {\n",
            "let n = write(fd, buf, count);\n",
            "if n < 0 { return -1 }\n",
            "buf = unsafe { buf.offset(n) };\n",
            "count -= n as usize;\n",
            "}\n",
            "0",
            "}\n",
        ])
        .await
    }

    pub async fn generate_wait_child_fn(&mut self) -> Result<()> {
        self.add_lines(&[
            "fn wait_child(pid: i32) -> i32 {\n",
            "let wstatus: i32 = 0;\n",
            "loop {\n",
            "if wait4(pid, &wstatus as *const i32, 0, ptr::null()) == -1 { break }\n",
            "if (wstatus & 0x7f) != 0 || ((((wstatus & 0x7f) + 1) >> 1) <= 0) { break }\n",
            "}\n",
            "0\n",
            "}\n",
        ])
        .await
    }
}

pub struct PendingCompile {
    pub child: Child,
}

impl PendingCompile {
    pub async fn wait(&mut self) -> Result<()> {
        let status = self.child.wait().await?;
        if !status.success() {
            bail!("Compile failed, compiler exited with {:?}", status);
        } else {
            Ok(())
        }
    }
}
