use libc::forkpty;
use nix::sys::wait::wait;
use nix::{
    errno::Errno,
    sys::{ptrace::Options, signal::Signal},
};
use std::{
    env::args,
    os::unix::process::CommandExt,
    process::{self, Command},
    thread,
    time::Duration,
};

use log::{debug, error, info, trace, warn};
use nix::{
    sys::{
        ptrace,
        wait::{self, WaitStatus},
    },
    unistd::{ForkResult, Pid, fork},
};

// pub type Error = Result<(), std::fmt::Error>;

#[derive(Debug)]
enum HavocError {
    Wait,
    Register(Errno),
    Ptrace(Errno),
}

fn main() -> Result<(), std::fmt::Error> {
    pretty_env_logger::formatted_builder()
        .filter_level(log::LevelFilter::Trace)
        .init();

    let args = &(args().collect::<Vec<String>>())[1..];

    match unsafe { fork() } {
        Ok(ForkResult::Child) => run_child(args),
        Ok(ForkResult::Parent { child }) => run_parent(child),
        Err(e) => panic!("Could not fork main process: {}", e),
    };

    Ok(())
}

fn run_child(args: &[String]) {
    info!("havoc child process executing as {}", process::id());
    info!("Running: {:?}", args);
    ptrace::traceme().expect("OS could not be bothered to trace me");

    // loop {
    //     println!("test");
    //     thread::sleep(Duration::from_secs(1));
    // }
    // unsafe { exit(0) };
    let e = Command::new(args.join(" ")).exec();
}

fn run_parent(pid: Pid) {
    // wait for our child process to be ready
    let ws = wait().expect("Parent failed waiting for child");
    // info!("Child process ready with signal: {ws:?}, will ask it to continue untill syscall");

    setup_tracing(pid).expect(format!("Parent failed to set up tracing for {pid:?}").as_str());
    // trace_syscall(pid, None).expect(format!("Parent failed tracing for {pid:?}").as_str());

    loop {
        match ptrace::syscall(pid, None).map_err(HavocError::Ptrace) {
            Ok(_) => { /* nop */ }
            Err(e) => {
                error!("error: {:?}", e);
                break;
            }
        }
    }
    println!("pid: {}", pid);
}
fn wait_for_signal() -> Result<(), HavocError> {
    match wait() {
        Ok(WaitStatus::Stopped(pid_t, sig_num)) => handle_child_stopped(sig_num, pid_t),

        Ok(WaitStatus::Exited(pid, exit_status)) => {
            debug!("Child with pid: {} exited with status {}", pid, exit_status);
            Ok(())
        }

        Ok(status) => {
            warn!("Received unhandled wait status: {:?}", status);
            Ok(())
        }

        Err(err) => {
            error!("An error occurred: {:?}", err);
            Err(HavocError::Wait)
        }
    }
}
fn handle_sigtrap(pid_t: Pid) -> Result<(), HavocError> {
    let regs = ptrace::getregs(pid_t).map_err(HavocError::Register)?;
    if regs.orig_rax == libc::SYS_msync as u64 {
        if regs.rax == -libc::ENOSYS as u64 {
            info!("Entry of syscall in {pid_t} : {}", regs.orig_rax);
        } else {
            info!("Exit of syscall in {pid_t} : {}", regs.orig_rax);
            // handle_msync(msync_counter, regs, pid_t)?;
        }
    }
    Ok(())
}
fn handle_child_stopped(sig_num: Signal, pid_t: Pid) -> Result<(), HavocError> {
    match sig_num {
        Signal::SIGTRAP => {
            handle_sigtrap(pid_t)?;
            trace_syscall(pid_t, None)
        }
        Signal::SIGSTOP => trace_syscall(pid_t, Some(Signal::SIGSTOP)),
        // ... some corner cases like SIGWINCH
        _ => trace_syscall(pid_t, Some(sig_num)),
    }
}
fn trace_syscall(pid: Pid, sig_num: Option<Signal>) -> Result<(), HavocError> {
    ptrace::syscall(pid, None).map_err(HavocError::Ptrace)?;
    trace!("{} -> {:?}", pid, sig_num);
    Ok(())
}

/// Setup the preace options to also trace fork, clone and vfork
fn setup_tracing(pid: Pid) -> Result<(), HavocError> {
    // ptrace::setoptions(
    //     pid,
    //     Options::PTRACE_O_TRACEFORK
    //         .union(Options::PTRACE_O_TRACECLONE)
    //         .union(Options::PTRACE_O_TRACEVFORK),
    // )
    ptrace::attach(pid).map_err(HavocError::Ptrace)
}

// fn next_syscall(child_pid: libc::pid_t) -> Result<i64, OsError> {
//     match WaitPID::from_process(child_pid)? {
//         WaitPID::SysCall { .. } => (),
//         x => {
//             return Err(OsError::new(
//                 ErrorKind::Other,
//                 format!("Expected syscall event, got {:?} instead", x),
//             ));
//         }
//     }
//     let number = get_syscall_number(child_pid)?;
//     p_trace_syscall(child_pid, None)?;
//     Ok(number)
// }