#include <stdio.h>
#include <stdlib.h>
#include <sys/ptrace.h>
#include <sys/uio.h>
#include <sys/wait.h>

#include <capstone/capstone.h>

/* #include "config.h" */
#include "debugger.h"
#include "helpers.h"

static const int PTRACE_OPTIONS = PTRACE_O_TRACESYSGOOD;

static void capture_state(struct tracee *dbg) {
    struct state *s = xmalloc(sizeof(*s));
    ptrace(PTRACE_GETREGS, dbg->id, NULL, &s->regs);
    ptrace(PTRACE_GETFPREGS, dbg->id, NULL, &s->fpregs);
    list_init(&s->maps);
    list_insert(dbg->states.end, s);
    dbg->state = s;

    char mapspath[128], entry[512];
    snprintf(mapspath, sizeof(mapspath), "/proc/%li/maps", (long)dbg->id);

    FILE *maps = fopen(mapspath, "r");
    if (maps) {
        while (fgets(entry, sizeof(entry), maps)) {
            struct map *m = xmalloc(sizeof(*m));
            sscanf(entry, "%lx-%lx ", &m->start, &m->end);
            size_t size = m->end - m->start;
            m->data = xmalloc(size);

            struct iovec loc = { m->data, size };
            struct iovec rem = { (void *)m->start, size };
            if (process_vm_readv(dbg->id, &loc, 1, &rem, 1, 0) < 0) {
                free(m->data);
                free(m);
                continue;
            }

            list_insert(s->maps.end, m);
        }

        fclose(maps);
    }
}

static void clear_states(struct tracee *dbg) {
    while (dbg->states.head != dbg->states.end) {
        struct state *s = dbg->states.head;
        while (s->maps.head != s->maps.end) {
            struct map *m = s->maps.head;
            list_remove(m);
            free(m->data);
            free(m);
        }
        list_remove(s);
        free(s);
    }
    dbg->state = NULL;
}

static void install_breakpoints(struct tracee *dbg) {
    struct list *breaks = &dbg->breaks;
    for (struct breakpoint *b = breaks->head; b != breaks->end; b = b->next) {
        if (b->enabled) {
            unsigned long data;
            data = ptrace(PTRACE_PEEKTEXT, dbg->id, b->address, NULL);
            b->orig = data;

            data = (data & ~0xff) | BREAKPOINT_INSN;
            ptrace(PTRACE_POKETEXT, dbg->id, b->address, data);
            b->active = 1;
        }
    }
}

static int uninstall_breakpoints(struct tracee *dbg, int stop) {
    struct user_regs_struct regs;
    int ch = 0;
    int reenter = 0;
    if (stop) {
        ptrace(PTRACE_GETREGS, dbg->id, NULL, &regs);
    }

    struct list *breaks = &dbg->breaks;
    for (struct breakpoint *b = breaks->head; b != breaks->end; b = b->next) {
        if (b->active) {
            ptrace(PTRACE_POKETEXT, dbg->id, b->address, b->orig);
            b->active = 0;

            /* detect stop at breakpoint */
            if (stop && !ch) {
                if (regs.rip - 1 == b->address) {
                    regs.rip--;
                    ptrace(PTRACE_SETREGS, dbg->id, NULL, &regs);
                    ch = 1;

                    if (b->stack != 0 && b->stack != regs.rsp) {
                        reenter = 1;
                    }
                }
            }

            if (b->enabled < 0 && !reenter) {
                struct breakpoint *del = b;
                b = b->prev;
                list_remove(del);
                free(del);
            }
        }
    }

    return reenter;
}

int dbg_process(struct tracee *dbg, pid_t pid) {
    if (ptrace(PTRACE_ATTACH, pid, NULL, NULL) < 0) {
        return -1;
    }

    list_init(&dbg->breaks);
    list_init(&dbg->states);
    dbg->state = NULL;
    dbg->id = pid;
    dbg->child = 0;
    dbg->stopped = 0;
    dbg->status = 0;
    dbg->signal = 0;
    dbg->cont = 0;
    dbg->buff = NULL;
    dbg->buffsize = 0;

    while (!dbg_wait(dbg)) {}
    ptrace(PTRACE_SETOPTIONS, pid, NULL, PTRACE_OPTIONS);
    return 0;
}

int dbg_new_process(struct tracee *dbg, char **argv, struct console *cons) {
    pid_t pid = fork();
    if (pid < 0) {
        return -1;
    }

    if (pid == 0) {
        console_configslave(cons);
        setpgid(0, 0);
        ptrace(PTRACE_TRACEME, 0, NULL, NULL);
        close_range(STDERR_FILENO+1, ~0U, CLOSE_RANGE_UNSHARE);
        execvp(argv[0], argv);
        exit(-1);
    }

    list_init(&dbg->breaks);
    list_init(&dbg->states);
    dbg->state = NULL;
    dbg->id = pid;
    dbg->child = 1;
    dbg->stopped = 0;
    dbg->status = 0;
    dbg->signal = 0;
    dbg->cont = 0;
    dbg->buff = NULL;
    dbg->buffsize = 0;

    while (!dbg_wait(dbg)) {}
    ptrace(PTRACE_SETOPTIONS, pid, NULL, PTRACE_O_EXITKILL | PTRACE_OPTIONS);
    return 0;
}

int dbg_wait(struct tracee *dbg) {
    if (dbg->stopped) {
        return 1;
    }

    if (waitpid(dbg->id, &dbg->status, WNOHANG | __WALL) <= 0) {
        return 0;
    }

    if (dbg->cont != 0) {
        install_breakpoints(dbg);
        ptrace(dbg->cont, dbg->id, NULL, NULL);
        dbg->cont = 0;
        return 0;
    }

    if (uninstall_breakpoints(dbg, 1)) {
        ptrace(PTRACE_SINGLESTEP, dbg->id, NULL, NULL);
        dbg->cont = PTRACE_CONT;
        return 0;
    }

    capture_state(dbg);
    dbg->stopped = 1;
    return 1;
}

/* note: move head < state < tail checks from main into these function
 * add a dbg_stepback function */

int dbg_stepin(struct tracee *dbg) {
    /* consider making this a no-op if we are not stopped on a call */
    ptrace(PTRACE_SINGLESTEP, dbg->id, NULL, NULL);
    dbg->stopped = 0;
    dbg->cont = 0;
    return 0;
}

int dbg_stepover(struct tracee *dbg) {
    csh handle;
    cs_open(CS_ARCH_X86, CS_MODE_64, &handle);

    uint64_t address = dbg->state->regs.rip;
    size_t size = 128;
    const uint8_t *code = deref(dbg, address, size);
    cs_insn *insn = cs_malloc(handle);

    cs_disasm_iter(handle, &code, &size, &address, insn);

    if (insn->id == X86_INS_CALL) {
        struct breakpoint *b = xmalloc(sizeof(*b));
        b->address = address;
        b->stack = dbg->state->regs.rsp;
        b->enabled = -1;
        b->active = 0;
        list_insert(dbg->breaks.end, b);
        dbg->cont = PTRACE_CONT;
        //dbg_cont(dbg, PTRACE_CONT);
    } else {
        dbg->cont = 0;
        //dbg_stepin(dbg);
    }

    ptrace(PTRACE_SINGLESTEP, dbg->id, NULL, NULL);
    dbg->stopped = 0;

    cs_free(insn, 1);
    cs_close(&handle);
    return 0;
}

/*
int dbg_stepout(struct tracee *dbg) {
    unsigned long bp = dbg->state->regs.rbp;
    unsigned long ra = *(unsigned long *)deref(dbg, bp+8, sizeof(unsigned long));

    struct breakpoint *b = xmalloc(sizeof(*b));
    b->address = ra;
    b->stack = 0;
    b->enabled = -1;
    b->active = 0;
    list_insert(dbg->breaks.end, b);

    ptrace(PTRACE_SINGLESTEP, dbg->id, NULL, NULL);
    //clear_states(dbg); // keep this?
    dbg->stopped = 0;
    dbg->cont = PTRACE_CONT;
    return 0;
}
*/

int dbg_cont(struct tracee *dbg, int cont) {
    ptrace(PTRACE_SINGLESTEP, dbg->id, NULL, NULL);
    clear_states(dbg);
    dbg->stopped = 0;
    dbg->cont = cont;
    return 0;
}

void *deref(struct tracee *dbg, unsigned long addr, size_t size) {
    (void)size; // todo

    struct list *maps = &dbg->state->maps;
    for (struct map *m = maps->head; m != maps->end; m = m->next) {
        if (m->start <= addr && addr < m->end) {
            return (unsigned char *)m->data + (addr - m->start);
        }
    }

    return NULL;
}