//! Process handling and monitoring using ptrace

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

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

/// Process Options
#[derive(Debug, Clone)]
pub struct ProcessOptions {
    /// Capture environnement
    pub capture_environ: bool,
    /// Capture status
    pub capture_status: bool,
    /// Process cwd
    pub cwd: Option<PathBuf>,
    /// Wrap in a shell
    pub wrap_shell: 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>,
        process_options: ProcessOptions,
    ) -> Result<BTreeMap<i32, ProgramData>, MekaError> {
        let wrap_shell = process_options.wrap_shell;
        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)
            }
        })?;
        let path_binary_cstr = CString::new(path_executable.to_str().ok_or(MekaError::new(
            format!("Conversion of {} failed", path_executable.display()),
        ))?)
        .map_err(|_| {
            MekaError::new(format!(
                "Cannot create a CString from {}",
                path_executable.display()
            ))
        })?;

        let display = format!(
            "[{}]",
            args_cmd
                .iter()
                .map(|s| format!("\"{}\"", s)) // add quotes around each string
                .collect::<Vec<_>>()
                .join(", ")
        );
        self.logger(
            1,
            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(|_| {
                MekaError::new(format!("Cannot create CString(s) for args {}", display))
            })?;

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

        let fd_out = memfd::memfd_create("out", memfd::MFdFlags::MFD_CLOEXEC)?;

        let fd_err = memfd::memfd_create("err", memfd::MFdFlags::MFD_CLOEXEC)?;
        if let Some(pwd) = &process_options.cwd {
            set_current_dir(pwd).map_err(|_| {
                MekaError::new(format!("Cannot set the current dir to {}", pwd.display()))
            })?;
        }
        // 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(|_| MekaError::new("Cannot create CString(s) for envs"))?;
        match unsafe { fork() } {
            Ok(ForkResult::Child) => {
                dup2_stdout(&fd_out)?;
                dup2_stderr(&fd_err)?;
                close(fd_out)?;
                close(fd_err)?;
                ptrace::traceme()?;
                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.logger(
                    1,
                    format!(
                        "Parent PID: {}, Child PID: {}",
                        father_process_pid, child_pid
                    ),
                );

                match waitpid(child_pid, None)? {
                    WaitStatus::Stopped(pid, sig) => {
                        self.logger(
                            2,
                            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);
                    }
                    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()
                        )));
                    }
                    status => {
                        return Err(MekaError::new(format!(
                            "Unexpected return from wait(): {:?}",
                            status
                        )));
                    }
                }
                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,
                )?;

                // continue
                ptrace::cont(child_pid, None)?;

                loop {
                    let status = waitpid(Pid::from_raw(-1), Some(WaitPidFlag::__WALL));
                    match status {
                        Ok(WaitStatus::Exited(exited_pid, code)) => {
                            self.logger(
                                1,
                                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, _dump)) => {
                            self.logger(
                                1,
                                format!(
                                    "[MONITOR] Process {} signaled (signal={:?})",
                                    signaled_pid.as_raw(),
                                    signal
                                ),
                            );
                            ptrace::cont(signaled_pid, None)?;
                        }
                        Ok(WaitStatus::StillAlive) => {
                            // Still running, continue monitoring
                            continue;
                        }
                        Ok(WaitStatus::Stopped(stopped_pid, signal)) => {
                            self.logger(
                                1,
                                format!(
                                    "[MONITOR] Process {} stopped (signal={:?})",
                                    stopped_pid, 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 let Some(parent) = process_table.get_mut(&parent_pid) {
                                    parent.subprograms.insert(stopped_pid.as_raw());
                                }
                            };
                            ptrace::cont(stopped_pid, None)?;
                        }
                        Ok(WaitStatus::PtraceEvent(ptrace_event_pid, signal, event_type)) => {
                            match event_type {
                                0 => {
                                    // nothing to be done for now
                                    self.logger(1, format!("[MONITOR] Event {}", event_type));
                                }
                                x if x == ptrace::Event::PTRACE_EVENT_CLONE as i32
                                    || x == ptrace::Event::PTRACE_EVENT_FORK as i32
                                    || x == ptrace::Event::PTRACE_EVENT_VFORK as i32 =>
                                {
                                    self.logger(
                                        1,
                                        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
                                        }
                                    }
                                    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,
                                    )?;
                                }
                                x if x == ptrace::Event::PTRACE_EVENT_EXIT as i32
                                    || x == ptrace::Event::PTRACE_EVENT_VFORK_DONE as i32 =>
                                {
                                    if let Some(process) =
                                        process_table.get_mut(&ptrace_event_pid.as_raw())
                                    {
                                        process.set_finished();
                                    }
                                }
                                x if x == ptrace::Event::PTRACE_EVENT_EXEC as i32 => {
                                    self.logger(
                                        1,
                                        format!(
                                            "[MONITOR] Exec event detected in process {}",
                                            ptrace_event_pid
                                        ),
                                    );
                                    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
                                }
                                event_num => {
                                    // do nothing
                                    self.logger(
                                        1,
                                        format!(
                                            "[MONITOR] Event in process {} event_num={}",
                                            ptrace_event_pid, event_num
                                        ),
                                    );
                                }
                            }

                            // continue
                            ptrace::cont(ptrace_event_pid, None)?;
                        }
                        Ok(WaitStatus::PtraceSyscall(ptrace_syscall_pid)) => {
                            // do nothing for now
                            ptrace::cont(ptrace_syscall_pid, None)?;
                        }
                        Ok(WaitStatus::Continued(continued_pid)) => {
                            // do nothing for now
                            ptrace::cont(continued_pid, None)?;
                        }
                        Err(e) => {
                            self.logger(1, format!("[MONITOR] waitpid 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(_) => Err(MekaError::new("fork failed")),
        }
    }
}
