#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 <errno.h>
#include <libelf.h>
#include <gelf.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"
#include "da.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
// Taken from da goat https://source.winehq.org/source/dlls/dbghelp/cpu_x86_64.c

typedef enum {
    rax,
    rbx,
    rcx,
    rdx,
    rdi,
    rsi,
    rbp,
    rsp,
    r8,
    r9,
    r10,
    r11,
    r12,
    r13,
    r14,
    r15,
    rip,
    eflags,
    cs,
    fs,
    gs,
    ss,
    ds,
    es,
    tr,
    ldtr,
    mxcsr,
    ctrl,
    stat,
    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(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(rip, 16),
    make_rd(eflags, 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(tr, 62),
    make_rd(ldtr, 63),
    make_rd(mxcsr, 64),
    make_rd(ctrl, 65),
    make_rd(stat, 66),
    #undef make_rd
};

// 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 *name;
    uintptr_t addr;
} Symbol;

typedef struct {
    Symbol *items;
    size_t count, capacity;
} Symbols;

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

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

// Memory

uint64_t dbg_mem_read(Dbg *dbg, uint64_t addr)
{
    return ptrace(PTRACE_PEEKDATA, dbg->pid, addr, NULL);
}

void dbg_mem_write(Dbg *dbg, uint64_t addr, uint64_t v)
{
    ptrace(PTRACE_POKEDATA, dbg->pid, addr, v);
}

uint64_t dbg_get_reg_value(Dbg *dbg, Register r)
{
    struct user_regs_struct rs;
    ptrace(PTRACE_GETREGS, dbg->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 dbg_set_reg_value(Dbg *dbg, Register r, uint64_t v)
{
    struct user_regs_struct rs;
    ptrace(PTRACE_GETREGS, dbg->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, dbg->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;
        }
    }
}

uint64_t dbg_get_rip(Dbg *dbg)
{
    return dbg_get_reg_value(dbg, rip);
}

void dbg_set_rip(Dbg *dbg, uint64_t v)
{
    dbg_set_reg_value(dbg, rip, v);
}

void dbg_step_brk(Dbg *dbg)
{
    uint64_t loc = dbg_get_rip(dbg) - 1;
    char key[20];
    snprintf(key, sizeof(key), "0x%"PRIxPTR, (uintptr_t)loc);
    Brk *brk = hashtable_get(&dbg->brks, key);
    if ((brk != NULL  && brk->enabled)) {
        uint64_t prev_instr = loc;
        dbg_set_rip(dbg, prev_instr);
        brk_disable(brk);
        ptrace(PTRACE_SINGLESTEP, brk->pid, NULL, NULL);
        dbg_wait(dbg);
        brk_enable(brk);
    }
}

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

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

    js_pushundefined(js);
}

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

void dbg_js_rmbaddr(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_splo(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_plo(js_State *js)
{
    Dbg *dbg = getdbg();
    js_pushnumber(js, dbg->program_load_offset);
}

void dbg_js_lsbrk(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" (%s)\n",
                    c+1, dbg->brks.buckets[i]->key, brk->addr,
                    brk->enabled ? "Enabled" : "Disabled");
            c++;
        }
    }

    js_pushundefined(js);
}

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

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

void dbg_js_ldscr(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_gf(js_State *js)
{
    Dbg *dbg = getdbg();
    js_pushstring(js, dbg->file);
}

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

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

void dbg_js_sr(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);
    dbg_set_reg_value(dbg, r, value);
    js_pushundefined(js);
}

void dbg_js_mrd(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 = dbg_mem_read(dbg, addr);
    char buf[20];
    snprintf(buf, sizeof(buf), "0x%"PRIx64, v);
    js_pushstring(js, buf);
}

void dbg_js_mwr(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);
    dbg_mem_write(dbg, addr, value);
    js_pushundefined(js);
}

void dbg_js_lsm(js_State *js)
{
    Dbg *dbg = getdbg();
    for (int i = 0; i < dbg->symbols.count; i++) {
        Symbol *sym = &dbg->symbols.items[i];
        LOG_INF("Symbol %s at %"PRIxPTR"\n", sym->name, sym->addr);
    }
    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(baddr, 1 /*addr*/);
    make_js_func(rmbaddr, 1 /*addr*/);
    make_js_func(lsbrk, 0);
    make_js_func(splo, 1 /*offset*/);
    make_js_func(lif, 1 /*string*/);
    make_js_func(ler, 1 /*string*/);
    make_js_func(ldscr, 1 /*path*/);
    make_js_func(gf, 0);
    make_js_func(pid, 0);
    make_js_func(plo, 0);
    make_js_func(gr, 1 /*reg name*/);
    make_js_func(sr, 2 /* reg name, value*/);
    make_js_func(mrd, 1 /*addr*/);
    make_js_func(mwr, 2 /*addr, value*/);
    make_js_func(lsm, 0);

    #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 = NULL;
    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);

    free(script_buf);
    fclose(script);
}

void dbg_load_symbols(Dbg *dbg)
{
    FILE *bin = fopen(dbg->file, "rb");
    if (bin == NULL) {
        LOG_ERR("Cannot load symbols from file\n");
        return;
    }

    Elf *elf;
    Elf_Scn *scn = NULL;
    GElf_Shdr shdr;
    Elf_Data *data;
    GElf_Sym sym;

    elf_version(EV_CURRENT);

    elf = elf_begin(fileno(bin), ELF_C_READ, NULL);
    if (elf == NULL) {
        LOG_ERR("Could not parse elf\n");
        fclose(bin);
        return;
    }

    while ((scn = elf_nextscn(elf, scn)) != NULL) {
        gelf_getshdr(scn, &shdr);
        if (shdr.sh_type == SHT_SYMTAB) {
            break;
        }
    }

    data = elf_getdata(scn, NULL);
    int count = shdr.sh_size / shdr.sh_entsize;

    for (int i = 0; i < count; i++) {
        gelf_getsym(data, i, &sym);
        if (sym.st_info != 2 /*func*/) {
            char *name = elf_strptr(elf, shdr.sh_link, sym.st_name);
            size_t len = strlen(name);
            if (len == 0 || sym.st_value == 0x0) {
                continue;
            }

            char *n = malloc(len+1);
            strcpy(n, name);
            n[len] = '\0';

            Symbol symbol = { .name = name, .addr = dbg->program_load_offset + (uintptr_t)sym.st_value };
            da_append(&dbg->symbols, symbol);
        }
    }

    elf_end(elf);
    fclose(bin);
}

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_symbols(dbg);
    dbg_load_script(dbg, INIT_SCRIPT);
}

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

void dbg_loop(Dbg *dbg)
{
    dbg_wait(dbg);
    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;
}