#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 "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++)

// 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_print_file(js_State *js)
{
    Dbg *dbg = getdbg();
    LOG_INF("File: %s\n", dbg->file);
    
    js_pushundefined(js);
}

void dbg_js_print_dbg_pid(js_State *js)
{
    Dbg *dbg = getdbg();
    LOG_INF("Debugger PID: %d\n", dbg->pid);
    
    js_pushundefined(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, "%"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, "%"SCNxPTR, &addr);
    dbg->program_load_offset = addr;

    js_pushundefined(js);
}

void dbg_js_print_program_load_offset(js_State *js)
{
    Dbg *dbg = getdbg();
    LOG_INF("Program load offset: 0x%"PRIxPTR"\n", dbg->program_load_offset);
    
    js_pushundefined(js);
}

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_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(print_file, 0);
    make_js_func(print_dbg_pid, 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);
    make_js_func(print_program_load_offset, 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 = pmparser_next(&maps_iter);

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

    pmparser_free(&maps_iter);
}

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);
}

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;
}