//! Process handling and monitoring using ptrace

use nix::unistd::{ForkResult, Pid, close, dup2_stderr, dup2_stdout, execve, fork, getpid};
use nix::{
    libc::SIGSTOP,
    sys::{
        memfd, ptrace,
        wait::{WaitPidFlag, WaitStatus, waitpid},
    },
};
use std::{
    collections::BTreeMap,
    env::{self, set_current_dir},
    ffi::{CString, NulError},
    path::PathBuf,
};

use crate::{
    args::RunArgs,
    errors::MekaError,
    meka::Meka,
    utils::{ParsedCommand, Proc, ProgramData, parse_cmd, read_memfd},
};

/// Process Options
#[derive(Debug, Clone)]
pub struct ProcessOptions {
    /// Capture environment
    pub capture_environ: bool,
    /// Capture status
    pub capture_status: bool,
}

/// Process Options
#[derive(Debug, Clone)]
pub struct RunOptions {
    /// Process cwd
    pub cwd: Option<PathBuf>,
    /// Wrap in a shell
    pub wrap_shell: bool,
    /// Write on insert
    pub write_on_insert: bool,
}

impl Meka {
    /// Process a command and return a map of process data
    /// # Errors
    /// Returns a `MekaError` if the command cannot be processed
    pub fn process_command(
        &self,
        path_var: Vec<PathBuf>,
        all_args: Vec<String>,
        run_options: RunArgs,
    ) -> Result<BTreeMap<i32, ProgramData>, (MekaError, BTreeMap<i32, ProgramData>)> {
        let wrap_shell = run_options.wrap_shell;
        let process_options = ProcessOptions {
            capture_environ: run_options.capture_environ,
            capture_status: run_options.capture_status,
        };
        let ParsedCommand {
            envs: cmd_envs,
            executable: path_executable,
            args: args_cmd,
        } = parse_cmd(&all_args, &path_var, wrap_shell)
            .map_err(|e| {
                if e.contains("not found") && self.should_show_debug(1) {
                    let debug_infos = {
                        // not using env::join_paths()
                        let paths = path_var
                            .iter()
                            .map(|p| p.to_string_lossy())
                            .collect::<Vec<_>>()
                            .join(":");
                        format!(" (PATH is '{paths}')")
                    };
                    MekaError::new(format!("{e}{debug_infos}"))
                } else {
                    MekaError::new(e)
                }
            })
            .map_err(|e| (e, BTreeMap::new()))?;
        let path_binary_cstr = CString::new(
            path_executable
                .to_str()
                .ok_or(MekaError::new(format!(
                    "Conversion of {} failed",
                    path_executable.display()
                )))
                .map_err(|e| (e, BTreeMap::new()))?,
        )
        .map_err(|e| {
            (
                MekaError::new(format!(
                    "Cannot create a CString from {}: {e}",
                    path_executable.display()
                )),
                BTreeMap::new(),
            )
        })?;

        let display = format!(
            "[{}]",
            args_cmd
                .iter()
                .map(|s| format!("\"{}\"", s)) // add quotes around each string
                .collect::<Vec<_>>()
                .join(", ")
        );
        self.logger1(format!(
            "Executing '{}' with args: {}",
            path_executable.display(),
            display
        ));

        let all_args_cstr = args_cmd
            .iter()
            .map(|s| CString::new(s.as_str()))
            .collect::<Result<Vec<CString>, NulError>>()
            .map_err(|e| {
                (
                    MekaError::new(format!(
                        "Cannot create CString(s) for args {}: {e}",
                        display
                    )),
                    BTreeMap::new(),
                )
            })?;

        let mut process_table = BTreeMap::<i32, ProgramData>::new();

        let fd_out = memfd::memfd_create("out", memfd::MFdFlags::MFD_CLOEXEC).map_err(|e| {
            (
                MekaError::new(format!("memfd_create error: {e}")),
                BTreeMap::new(),
            )
        })?;

        let fd_err = memfd::memfd_create("err", memfd::MFdFlags::MFD_CLOEXEC).map_err(|e| {
            (
                MekaError::new(format!("memfd_create error: {e}")),
                BTreeMap::new(),
            )
        })?;
        if let Some(pwd) = &run_options.cwd {
            set_current_dir(pwd).map_err(|e| {
                (
                    MekaError::new(format!(
                        "Cannot set the current dir to {}: {e}",
                        pwd.display()
                    )),
                    BTreeMap::new(),
                )
            })?;
        }
        // other option was to use unsafe set_var
        let env_vars: Vec<String> = env::vars()
            .map(|(key, val)| format!("{key}={val}"))
            .collect();
        let cmd_envs: Vec<String> = cmd_envs.iter().map(|(k, v)| format!("{k}={v}")).collect();
        let env_vars: Vec<String> = env_vars.into_iter().chain(cmd_envs).collect();
        let all_envs_cstr = env_vars
            .iter()
            .map(|s| CString::new(s.as_str()))
            .collect::<Result<Vec<CString>, NulError>>()
            .map_err(|e| {
                (
                    MekaError::new(format!("Cannot create CString(s) for envs {e}")),
                    BTreeMap::new(),
                )
            })?;
        match unsafe { fork() } {
            Ok(ForkResult::Child) => {
                if let Err(e) = dup2_stdout(&fd_out) {
                    return Err((e.into(), process_table));
                }
                if let Err(e) = dup2_stderr(&fd_err) {
                    return Err((e.into(), process_table));
                }
                if let Err(e) = close(fd_out) {
                    return Err((e.into(), process_table));
                }
                if let Err(e) = close(fd_err) {
                    return Err((e.into(), process_table));
                }
                if let Err(e) = ptrace::traceme() {
                    self.logger1(format!("PTRACE_TRACEME failed - most likely because the process is already being traced: {e}"));
                }
                let _ = execve(&path_binary_cstr, &all_args_cstr, &all_envs_cstr);
                unreachable!();
            }
            Ok(ForkResult::Parent { child: child_pid }) => {
                let father_process_pid = getpid().as_raw();
                self.logger1(format!(
                    "Parent PID: {father_process_pid}, Child PID: {child_pid}"
                ));

                match waitpid(child_pid, None) {
                    Ok(WaitStatus::Stopped(pid, sig)) => {
                        self.logger2(format!(
                            "Child process {} stopped (signal={:?})",
                            pid.as_raw(),
                            sig
                        ));
                        // insert the parent process
                        let program = ProgramData::from_proc(
                            &Proc::new(pid.as_raw()),
                            process_options.clone(),
                        );
                        process_table.insert(pid.as_raw(), program);
                        if run_options.write_on_insert
                            && let Err(e) =
                                Self::write_meka_file(&process_table, &run_options.output_file)
                        {
                            return Err((e, process_table));
                        }
                    }
                    Ok(WaitStatus::Exited(pid, return_code)) => {
                        return Err((
                            MekaError::new(format!(
                                "Failed to start the exec: pid({}) returned with code={} (it seems that {} is not an executable)",
                                pid.as_raw(),
                                return_code,
                                path_executable.display()
                            )),
                            process_table,
                        ));
                    }
                    status => {
                        return Err((
                            MekaError::new(format!("Unexpected return from wait(): {:?}", status)),
                            process_table,
                        ));
                    }
                }
                if let Err(e) = ptrace::setoptions(
                    child_pid,
                    ptrace::Options::PTRACE_O_TRACEFORK
                        | ptrace::Options::PTRACE_O_TRACEVFORK
                        | ptrace::Options::PTRACE_O_TRACECLONE
                        | ptrace::Options::PTRACE_O_TRACEEXEC
                        | ptrace::Options::PTRACE_O_TRACEVFORKDONE
                        | ptrace::Options::PTRACE_O_TRACEEXIT
                        | ptrace::Options::PTRACE_O_EXITKILL,
                ) {
                    return Err((e.into(), process_table));
                }

                // continue
                if let Err(e) = ptrace::cont(child_pid, None) {
                    return Err((e.into(), process_table));
                }

                loop {
                    self.logger4("[MONITOR] Starting waitpid(-1)");
                    let status = waitpid(Pid::from_raw(-1), Some(WaitPidFlag::__WALL));
                    self.logger4("[MONITOR] waitpid(-1) finished");
                    match status {
                        Ok(WaitStatus::Exited(exited_pid, code)) => {
                            self.logger2(format!(
                                "[MONITOR] Process {} exited with code={}",
                                exited_pid.as_raw(),
                                code
                            ));
                            let pid_raw = exited_pid.as_raw();
                            if let Some(process) = process_table.get_mut(&pid_raw) {
                                process.set_finished();
                            }
                        }
                        Ok(WaitStatus::Signaled(signaled_pid, signal, is_core_dump)) => {
                            // Process is killed by a signal
                            self.logger2(format!(
                                "[MONITOR] Process {signaled_pid} signaled (signal={:?}, core_dump={})",
                                signal, is_core_dump
                            ));
                            if let Err(e) = ptrace::cont(signaled_pid, None) {
                                return Err((e.into(), process_table));
                            }
                        }
                        Ok(WaitStatus::StillAlive) => {
                            self.logger3("[MONITOR] Process still alive");
                            // Still running, continue monitoring
                            continue;
                        }
                        Ok(WaitStatus::Stopped(stopped_pid, signal)) => {
                            self.logger2(format!(
                                "[MONITOR] Process {stopped_pid} stopped (signal={:?})",
                                signal
                            ));
                            let proc = Proc::new(stopped_pid.as_raw());
                            if let Some(parent_pid) = proc.parent_pid() {
                                let mut program =
                                    ProgramData::from_proc(&proc, process_options.clone());
                                program.parent_pid = parent_pid;
                                if let Some(_proc) = process_table.get(&stopped_pid.as_raw()) {
                                    // most likely stopped by a signal SIGCHLD - process is already in table
                                } else {
                                    process_table.insert(stopped_pid.as_raw(), program);
                                    if run_options.write_on_insert
                                        && let Err(e) = Self::write_meka_file(
                                            &process_table,
                                            &run_options.output_file,
                                        )
                                    {
                                        return Err((e, process_table));
                                    }
                                }
                                if let Some(parent) = process_table.get_mut(&parent_pid) {
                                    parent.subprograms.insert(stopped_pid.as_raw());
                                }
                            };
                            if signal as i32 == SIGSTOP {
                                if let Err(e) = ptrace::cont(stopped_pid, None) {
                                    return Err((e.into(), process_table));
                                }
                            } else {
                                // forward others signals
                                if let Err(e) = ptrace::cont(stopped_pid, signal) {
                                    return Err((e.into(), process_table));
                                }
                            }
                        }
                        Ok(WaitStatus::PtraceEvent(ptrace_event_pid, signal, event_type)) => {
                            self.logger2(format!(
                                "[MONITOR] Process {} ptraceEvent {:?} {}",
                                ptrace_event_pid, signal, event_type
                            ));
                            match event_type {
                                0 => {
                                    // nothing to be done for now
                                }
                                x if x == ptrace::Event::PTRACE_EVENT_FORK as i32
                                    || x == ptrace::Event::PTRACE_EVENT_VFORK as i32
                                    || x == ptrace::Event::PTRACE_EVENT_CLONE as i32 =>
                                {
                                    self.logger3(format!(
                                        "[MONITOR] Process {} forked/cloned (signal={:?})",
                                        ptrace_event_pid.as_raw(),
                                        signal
                                    ));
                                    let proc = Proc::new(ptrace_event_pid.as_raw());
                                    if let Some(parent_pid) = proc.parent_pid() {
                                        if let Some(parent) = process_table.get_mut(&parent_pid) {
                                            // sub fork - we add the pid to parent
                                            parent.subprograms.insert(ptrace_event_pid.as_raw());
                                        } else {
                                            // It's a fork of child_pid - the parent is the father process
                                        }
                                    }
                                    if let Err(e) = ptrace::setoptions(
                                        ptrace_event_pid,
                                        ptrace::Options::PTRACE_O_TRACEFORK
                                            | ptrace::Options::PTRACE_O_TRACEVFORK
                                            | ptrace::Options::PTRACE_O_TRACECLONE
                                            | ptrace::Options::PTRACE_O_TRACEEXEC
                                            | ptrace::Options::PTRACE_O_TRACEVFORKDONE
                                            | ptrace::Options::PTRACE_O_TRACEEXIT
                                            | ptrace::Options::PTRACE_O_EXITKILL,
                                    ) {
                                        return Err((e.into(), process_table));
                                    }
                                }
                                x if x == ptrace::Event::PTRACE_EVENT_EXEC as i32 => {
                                    self.logger3(format!(
                                        "[MONITOR] Process {} exec (signal={:?})",
                                        ptrace_event_pid, signal
                                    ));
                                    let proc = Proc::new(ptrace_event_pid.as_raw());
                                    if let Some(program) =
                                        process_table.get_mut(&ptrace_event_pid.as_raw())
                                    {
                                        program.after_exec.push(proc.cmdline().unwrap_or_default());
                                    }
                                    // nothing to be done for now
                                }
                                x if x == ptrace::Event::PTRACE_EVENT_VFORK_DONE as i32
                                    || x == ptrace::Event::PTRACE_EVENT_EXIT as i32 =>
                                {
                                    self.logger3(format!(
                                        "[MONITOR] Process {ptrace_event_pid} done/exit (signal={:?})",
                                        signal
                                    ));
                                    if let Some(process) =
                                        process_table.get_mut(&ptrace_event_pid.as_raw())
                                    {
                                        process.set_finished();
                                    }
                                }
                                event_num => {
                                    // do nothing
                                    self.logger3(format!(
                                        "[MONITOR] Process {} unknown (signal={:?}) num={}",
                                        ptrace_event_pid, signal, event_num
                                    ));
                                }
                            }

                            // continue
                            if let Err(e) = ptrace::cont(ptrace_event_pid, None) {
                                return Err((e.into(), process_table));
                            }
                        }
                        Ok(WaitStatus::PtraceSyscall(ptrace_syscall_pid)) => {
                            self.logger2(format!(
                                "[MONITOR] Process {ptrace_syscall_pid} ptraceSyscall "
                            ));
                            // do nothing for now
                            if let Err(e) = ptrace::cont(ptrace_syscall_pid, None) {
                                return Err((e.into(), process_table));
                            }
                        }
                        Ok(WaitStatus::Continued(continued_pid)) => {
                            self.logger2(format!("[MONITOR] Process {continued_pid} continued "));
                            // do nothing for now
                            if let Err(e) = ptrace::cont(continued_pid, None) {
                                return Err((e.into(), process_table));
                            }
                        }
                        Err(e) => {
                            self.logger1(format!("[MONITOR] waitpid(-1) failed: {e}"));
                            break;
                        }
                    }
                }
                self.logger(
                    1,
                    format!("[MONITOR] Child process has exited. PID: {child_pid}"),
                );
                if let Some(data) = read_memfd(fd_out) {
                    self.logger(2, "-------- Child stdout --------");
                    self.logger(2, &data);
                    self.logger(2, "-------- End --------");
                }
                if let Some(data) = read_memfd(fd_err) {
                    self.logger(2, "-------- Child stderr --------");
                    self.logger(2, &data);
                    self.logger(2, "-------- End --------");
                }
                self.logger(1, "[MONITOR] Parent process has finished.");
                Ok(process_table)
            }
            Err(e) => Err((MekaError::new(format!("fork() failed: {e}")), process_table)),
        }
    }
}
