#include <sys/ptrace.h>
#include <sys/wait.h>
#include <sys/uio.h>
#include <cstring>
#include <cerrno>
#include <elf.h>

#ifndef __aarch64__
#error "Unsupported Host Architecture"
#endif
#include <asm/ptrace.h>

#include "ptrace.h"
#include "log.h"
#include "trampoline.h"
#include "memory.h"

// Helper Functions
static void safe_ptrace(__ptrace_request request, pid_t pid, void *addr, void *data) {
    if (ptrace(request, pid, addr, data) == -1) {
        ERR("PTrace Error: %s", strerror(errno));
    }
}
static void safe_wait(pid_t guest_pid, int *status) {
    if (waitpid(guest_pid, status, 0) == -1) {
        ERR("Unable To Wait For Guest");
    }
    if (WIFEXITED(*status)) {
        // Process Exited
        exit(WEXITSTATUS(*status));
    }
}
static void ptrace_wait_syscall(pid_t guest_pid) {
    while (true) {
        safe_ptrace(PTRACE_SYSCALL, guest_pid, nullptr, nullptr);
        int status;
        safe_wait(guest_pid, &status);
        if (WIFSTOPPED(status) && (WSTOPSIG(status) & 0x80)) {
            return;
        }
    }
}

// Registers
struct arm32_pt_regs {
    uint32_t regs[18];
};
static bool check_syscall(pid_t guest_pid, arm32_pt_regs *out) {
    // Read Registers
    arm32_pt_regs regs = {};
    iovec io = {
        .iov_base = &regs,
        .iov_len = sizeof(regs)
    };
    safe_ptrace(PTRACE_GETREGSET, guest_pid, nullptr, &regs);
    if (io.iov_len != sizeof(regs)) {
        ERR("Guest Must Be 32-Bit");
    }
    // Check Syscall
    if (regs.regs[7] != TRAMPOLINE_SYSCALL) {
        // Not Trampoline
        return false;
    }
    // Export Registers
    *out = regs;
    // Return
    return true;
}
static void get_arguments(arm32_pt_regs *regs, uint32_t *arg1, uint32_t *arg2, uint32_t *arg3) {
    *arg1 = regs->regs[1];
    *arg2 = regs->regs[2];
    *arg3 = regs->regs[3];
}
static void set_syscall_return(pid_t guest_pid, arm32_pt_regs *regs) {
    regs->regs[0] = 0;
    iovec io = {
        .iov_base = regs,
        .iov_len = sizeof(arm32_pt_regs)
    };
    safe_ptrace(PTRACE_SETREGSET, guest_pid, (void *) NT_PRSTATUS, &io);
}

// Main Loop
void loop_ptrace(pid_t guest_pid) {
    while (true) {
        // Wait For Syscall
        ptrace_wait_syscall(guest_pid);

        // Handle Syscall Entrance
        arm32_pt_regs regs = {};
        bool is_trampoline = check_syscall(guest_pid, &regs);
        if (is_trampoline) {
            // Get Arguments
            uint32_t id;
            uint32_t length;
            uint32_t args_addr;
            get_arguments(&regs, &id, &length, &args_addr);
            static unsigned char args[MAX_TRAMPOLINE_ARGS_SIZE];
            memory_reader(args_addr, args, length);
            // Run Trampoline
            uint32_t ret = trampoline(id, args);
            memory_writer(args_addr, &ret, sizeof(uint32_t));
        }

        // Handle Syscall Exit
        ptrace_wait_syscall(guest_pid);

        // Set Syscall Return
        if (is_trampoline) {
            set_syscall_return(guest_pid, &regs);
        }
    }
}