#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <stdbool.h>
#include <stdint.h>
#include <inttypes.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/ptrace.h>
#include <sys/personality.h>
#include <sys/user.h>

#include "linenoise.h"
#include "mujs.h"
#include "hash.h"
#include "pmparser.h"

#define LOG_ERR(fmt, ...) fprintf(stderr, "Error: " fmt, ##__VA_ARGS__)
#define LOG_INF(fmt, ...) fprintf(stdout, "Info:  " fmt, ##__VA_ARGS__)

#define unused(x) ((void)(x))

#define shift(argc, argv, msg, ...) (argc <= 0 ? (LOG_ERR(msg, ##__VA_ARGS__), exit(EXIT_FAILURE)) : (void)0, argc--, *argv++)

#define INIT_SCRIPT ".debugusrc.js"

// Registers

typedef enum {
    rax,
    rbx,
    rcx,
    rdx,
    rdi,
    rsi,
    rbp,
    rsp,
    r8,
    r9,
    r10,
    r11,
    r12,
    r13,
    r14,
    r15,
    rip,
    rflags,
    cs,
    orig_rax,
    fs_base,
    gs_base,
    fs,
    gs,
    ss,
    ds,
    es,
    MAX_REGISTERS,
} Register;

typedef struct {
    Register r;
    int dwarf_r;
    const char *name;
} RegisterDescriptor;

static RegisterDescriptor reg_descriptors[MAX_REGISTERS] = {
    #define make_rd(r, dr) { r, dr, #r }
    make_rd(orig_rax, -1),
    make_rd(rip, -1),
    make_rd(rax, 0),
    make_rd(rdx, 1),
    make_rd(rcx, 2),
    make_rd(rbx, 3),
    make_rd(rsi, 4),
    make_rd(rdi, 5),
    make_rd(rbp, 6),
    make_rd(rsp, 7),
    make_rd(r8, 8),
    make_rd(r9, 9),
    make_rd(r10, 10),
    make_rd(r11, 11),
    make_rd(r12, 12),
    make_rd(r13, 13),
    make_rd(r14, 14),
    make_rd(r15, 15),
    make_rd(rflags, 49),
    make_rd(es, 50),
    make_rd(cs, 51),
    make_rd(ss, 52),
    make_rd(ds, 53),
    make_rd(fs, 54),
    make_rd(gs, 55),
    make_rd(fs_base, 58),
    make_rd(gs_base, 59),
    #undef make_rd
};

uint64_t get_reg_value(pid_t pid, Register r)
{
    struct user_regs_struct rs;
    ptrace(PTRACE_GETREGS, pid, NULL, &rs);

    for (int i = 0; i < sizeof(reg_descriptors)/sizeof(reg_descriptors[0]); i++) {
        if (reg_descriptors[i].r == r) {
            return ((uint64_t*)&rs)[i];
        }
    }
}

void set_reg_value(pid_t pid, Register r, uint64_t v)
{
    struct user_regs_struct rs;
    ptrace(PTRACE_GETREGS, pid, NULL, &rs);

    for (int i = 0; i < sizeof(reg_descriptors)/sizeof(reg_descriptors[0]); i++) {
        if (reg_descriptors[i].r == r) {
            ((uint64_t*)&rs)[i] = v;
            ptrace(PTRACE_SETREGS, pid, NULL, &rs);
            return;
        }
    }
}

const char *get_reg_name(Register r)
{
    for (int i = 0; i < sizeof(reg_descriptors)/sizeof(reg_descriptors[0]); i++) {
        if (reg_descriptors[i].r == r) {
            return reg_descriptors[i].name;
        }
    }
}

Register get_reg_from_name(const char *name)
{
    for (int i = 0; i < MAX_REGISTERS; i++) {
        if (strcmp(name, reg_descriptors[i].name) == 0) {
            return reg_descriptors[i].r;
        }
    }
}

// Memory

uint64_t mem_read(pid_t pid, uint64_t addr)
{
    return ptrace(PTRACE_PEEKDATA, pid, addr, NULL);
}

void mem_write(pid_t pid, uint64_t addr, uint64_t v)
{
    ptrace(PTRACE_POKEDATA, pid, addr, v);
}

// How breakpoints work?
// We can enable/disable breakpoints by putting/removing an int 3 instruction
// into/from the executed program. int 3 will trigger a SIGTRAP, which we can
// then handle. Pretty cool, eh?
//
// To handle a SIGTRAP, we can use waitpid(2) to listen for child process' events.

#define MAX_BRKS 0xff

typedef struct {
    pid_t pid;
    pid_t proc_pid;
    uintptr_t addr;
    bool enabled;
    uint8_t data;
} Brk;

void brk_enable(Brk *brk)
{
    long data = ptrace(PTRACE_PEEKDATA, brk->pid, brk->addr, NULL);
    brk->data = (uint8_t)(data&0xff);
    uint64_t int3 = 0xCC; // int 3 instruction, SIGTRAP
    uint64_t data_with_int3 = ((data& ~0xff) | int3);
    ptrace(PTRACE_POKEDATA, brk->pid, brk->addr, data_with_int3);

    brk->enabled = true;
}

void brk_disable(Brk *brk)
{
    long data = ptrace(PTRACE_PEEKDATA, brk->pid, brk->addr, NULL);
    long restored_data = ((data& ~0xff) | brk->data);
    ptrace(PTRACE_POKEDATA, brk->pid, brk->addr, restored_data);

    brk->enabled = false;
}

typedef struct {
    const char *file;
    pid_t pid;
    js_State *js;
    HashTable brks;
    uintptr_t program_load_offset;
} Dbg;

#define getdbg() ((Dbg*)js_currentfunctiondata(js))

void dbg_js_cont(js_State *js)
{
    Dbg *dbg = getdbg();
    ptrace(PTRACE_CONT, dbg->pid, NULL, NULL);

    int status, options = 0;
    waitpid(dbg->pid, &status, options);

    js_pushundefined(js);
}

void dbg_js_mk_brk_addr(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    uintptr_t addr;
    sscanf(addr_str, "0x%"SCNxPTR, &addr);
    Brk brk = { .pid = dbg->pid, .addr = dbg->program_load_offset + addr };
    brk_enable(&brk);
    hashtable_set(&dbg->brks, addr_str, &brk, sizeof(brk));
done:
    js_pushundefined(js);
}

void dbg_js_rm_brk_addr(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    Brk *brk = (Brk *)hashtable_get(&dbg->brks, addr_str);
    if (brk == NULL) {
        LOG_ERR("No breakpoint at address: %s\n", addr_str);
        goto done;
    }
    brk_disable(brk);
    hashtable_delete(&dbg->brks, addr_str);
done:
    js_pushundefined(js);
}

void dbg_js_set_program_load_offset(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    uintptr_t addr;
    sscanf(addr_str, "0x%"SCNxPTR, &addr);
    dbg->program_load_offset = addr;

    js_pushundefined(js);
}

void dbg_js_get_program_load_offset(js_State *js)
{
    Dbg *dbg = getdbg();
    js_pushnumber(js, dbg->program_load_offset);
}

void dbg_js_list_brks(js_State *js)
{
    Dbg *dbg = getdbg();
    int c = 0;
    for (int i = 0; i < dbg->brks.capacity; i++) {
        if (dbg->brks.buckets[i] != NULL) {
            Brk *brk = (Brk*)dbg->brks.buckets[i]->value;
            LOG_INF("Breakpoint %d (%s) at 0x%"PRIxPTR"\n",
                    c+1, dbg->brks.buckets[i]->key, brk->addr);
            c++;
        }
    }

    js_pushundefined(js);
}

void dbg_js_log_inf(js_State *js)
{
    const char *str = js_tostring(js, 1);
    LOG_INF("%s\n", str);
    js_pushundefined(js);
}

void dbg_js_log_err(js_State *js)
{
    const char *str = js_tostring(js, 1);
    LOG_ERR("%s\n", str);
    js_pushundefined(js);
}

void dbg_js_load_script(js_State *js)
{
    void dbg_load_script(Dbg *dbg, const char *script_path);
    Dbg *dbg = getdbg();
    const char *path = js_tostring(js, 1);
    dbg_load_script(dbg, path);
    js_pushundefined(js);
}

void dbg_js_get_file(js_State *js)
{
    Dbg *dbg = getdbg();
    js_pushstring(js, dbg->file);
}

void dbg_js_get_pid(js_State *js)
{
    Dbg *dbg = getdbg();
    js_pushnumber(js, dbg->pid);
}

void dbg_js_get_reg(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *name = js_tostring(js, 1);
    Register r = get_reg_from_name(name);
    uint64_t v = get_reg_value(dbg->pid, r);
    char buf[20];
    snprintf(buf, sizeof(buf), "0x%016"PRIx64, v);
    js_pushstring(js, buf);
}

void dbg_js_set_reg(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *name = js_tostring(js, 1);
    Register r = get_reg_from_name(name);
    const char *value_str = js_tostring(js, 2);
    uint64_t value;
    sscanf(value_str, "0x%"SCNx64, &value);
    set_reg_value(dbg->pid, r, value);
    js_pushundefined(js);
}

void dbg_js_mem_read(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    uintptr_t addr;
    sscanf(addr_str, "0x%"SCNxPTR, &addr);
    uint64_t v = mem_read(dbg->pid, addr);
    char buf[20];
    snprintf(buf, sizeof(buf), "0x%"PRIx64, v);
    js_pushstring(js, buf);
}

void dbg_js_mem_write(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    uintptr_t addr;
    sscanf(addr_str, "0x%"SCNxPTR, &addr);
    const char *value_str = js_tostring(js, 2);
    uint64_t value;
    sscanf(value_str, "0x%"SCNx64, &value);
    mem_write(dbg->pid, addr, value);
    js_pushundefined(js);
}

void dbg_init_js(Dbg *dbg)
{
    dbg->js   = js_newstate(NULL, NULL, JS_STRICT);

    #define make_js_func(name, n) \
    do { \
        js_newcfunctionx(dbg->js, &dbg_js_##name, #name, n, dbg, NULL); \
        js_setglobal(dbg->js, #name); \
    } while(0)

    make_js_func(cont, 0);
    make_js_func(mk_brk_addr, 1 /*addr*/);
    make_js_func(rm_brk_addr, 1 /*addr*/);
    make_js_func(list_brks, 0);
    make_js_func(set_program_load_offset, 1 /*offset*/);
    make_js_func(log_inf, 1 /*string*/);
    make_js_func(log_err, 1 /*string*/);
    make_js_func(load_script, 1 /*path*/);
    make_js_func(get_file, 0);
    make_js_func(get_pid, 0);
    make_js_func(get_program_load_offset, 0);
    make_js_func(get_reg, 1 /*reg name*/);
    make_js_func(set_reg, 2 /* reg name, value*/);
    make_js_func(mem_read, 1 /*addr*/);
    make_js_func(mem_write, 2 /*addr, value*/);

    #undef make_js_func
}

void dbg_init_load_offset(Dbg *dbg)
{
    procmaps_iterator maps_iter = {0};
    procmaps_error_t parser_err = PROCMAPS_SUCCESS;

    parser_err = pmparser_parse(dbg->pid, &maps_iter);
    if (parser_err) {
        LOG_ERR("Failed to parse /proc/%d/maps (%d)\n", dbg->pid, (int)parser_err);
        return;
    }

    // We only need the first one
    procmaps_struct *mem_region = pmparser_next(&maps_iter);

    dbg->program_load_offset = (uintptr_t)mem_region->addr_start;

    pmparser_free(&maps_iter);
}

void dbg_load_script(Dbg *dbg, const char *script_path)
{
    FILE *script = fopen(script_path, "r");
    if (script == NULL) {
        return;
    }

    fseek(script, 0L, SEEK_END);
    long size = ftell(script);
    rewind(script);
    char *script_buf = malloc(size+1);
    fread(script_buf, size, 1, script);
    script_buf[size] = '\0';

    js_dostring(dbg->js, script_buf);

    fclose(script);
}

void dbg_init(Dbg *dbg, const char *file, pid_t pid)
{
    memset(dbg, 0, sizeof(*dbg));
    dbg->file = file;
    dbg->pid  = pid;
    dbg_init_js(dbg);
    dbg_init_load_offset(dbg);
    hashtable_init(&dbg->brks, MAX_BRKS);

    dbg_load_script(dbg, INIT_SCRIPT);
}

void dbg_deinit(Dbg *dbg)
{
    js_freestate(dbg->js);
    hashtable_deinit(&dbg->brks);
}

void dbg_loop(Dbg *dbg)
{
    int status, options = 0;
    waitpid(dbg->pid, &status, options);

    char *line = NULL;
    while ((line = linenoise("debugus # ")) != NULL) {
        js_dostring(dbg->js, line);
        linenoiseHistoryAdd(line);
        linenoiseFree(line);
    }
}

void start_debugging(const char *input_file_path)
{
    pid_t pid = fork();
    if (pid == 0) {
        personality(ADDR_NO_RANDOMIZE);
        ptrace(PTRACE_TRACEME, 0, NULL, NULL);
        execl(input_file_path, input_file_path, NULL);
    } else if (pid >= 1) {
        Dbg dbg;
        dbg_init(&dbg, input_file_path, pid);
        dbg_loop(&dbg);
        dbg_deinit(&dbg);
    }
}

int main(int argc, char ** argv)
{
    const char *program_name = shift(argc, argv, "");
    unused(program_name);

    const char *input_file_path = shift(argc, argv, "No input file provided\n");
    start_debugging(input_file_path);

    return 0;
}