#include <dirent.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.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 | PTRACE_O_TRACECLONE;
static const int INTERRUPT_ON_SEIZE = 1;

static int detect_breakpoint(struct thread *th) {
    struct user_regs_struct regs;
    int check = 0;
    int restart = 0;

    ptrace(PTRACE_GETREGS, th->id, NULL, &regs);

    struct list *breaks = &th->proc->breakpoints;
    for (struct breakpoint *b = breaks->head; b != breaks->end; b = b->next) {
        /* should be if (b->active) ?
         * but there is a potential other bug that can happen
         * with the order of events in dbg_wait() */
        if (b->enabled) {
            if (!check) {
                if (regs.rip - 1 == b->address) {
                    regs.rip--;
                    ptrace(PTRACE_SETREGS, th->id, NULL, &regs);
                    check = 1;

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

                    if (b->tid != 0 && b->tid != th->id) {
                        restart = 1;
                    }
                }
            }

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

    return restart;
}

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

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

static void uninstall_breakpoints(struct process *proc) {
    struct list *breaks = &proc->breakpoints;
    for (struct breakpoint *b = breaks->tail; b != breaks->end; b = b->prev) {
        if (b->active) {
            ptrace(PTRACE_POKETEXT, proc->id, b->address, b->orig);
            b->active = 0;
        }
    }
}

static void clear_breakpoints(struct process *proc) {
    while (proc->breakpoints.head != proc->breakpoints.end) {
        struct breakpoint *b = proc->breakpoints.head;
        list_remove(b);
        free(b);
    }
}

static void clear_states(struct thread *th) {
    while (th->states.head != th->states.end) {
        struct state *s = th->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);
    }

    th->state = NULL;
    th->clearstates = 0;
}

static void capture_state(struct thread *th) {
    if (th->clearstates) {
        clear_states(th);
    }

    struct state *s = xmalloc(sizeof(*s));
    ptrace(PTRACE_GETREGS, th->id, NULL, &s->regs);
    ptrace(PTRACE_GETFPREGS, th->id, NULL, &s->fpregs);
    list_init(&s->maps);

    list_insert(th->states.end, s);
    th->state = s;

    char mapspath[32], entry[512];
    snprintf(mapspath, sizeof(mapspath), "/proc/%li/maps", (long)th->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(th->id, &loc, 1, &rem, 1, 0) < 0) {
                free(m->data);
                free(m);
                continue;
            }

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

        fclose(maps);
    }
}

static struct thread *new_thread(struct process *proc, pid_t id) {
    struct thread *th = xmalloc(sizeof(*th));
    th->proc = proc;
    list_init(&th->states);
    th->state = NULL;
    th->clearstates = 0;
    th->id = id;
    th->stopped = 0;
    th->signal = 0;
    th->cont = 0;
    th->status = "RUNNING";
    return th;
}

static void interrupt_all_threads(struct process *proc) {
    struct list *threads = &proc->threads;
    for (struct thread *th = threads->head; th != threads->end; th = th->next) {
        if (!th->stopped) {
            ptrace(PTRACE_INTERRUPT, th->id, NULL, NULL);
            while (!dbg_wait(th, 0)) {}
        }
    }
}

void add_breakpoint(struct process *proc, unsigned long address, unsigned long stack, pid_t tid, int enabled) {
    struct breakpoint *b = xmalloc(sizeof(*b));
    b->orig = 0;
    b->address = address;
    b->stack = stack;
    b->tid = tid;
    b->enabled = enabled;
    b->active = 0;
    list_insert(proc->breakpoints.end, b);
}

int is_breakpoint(struct process *proc, unsigned long address) {
    struct list *breaks = &proc->breakpoints;
    for (struct breakpoint *b = breaks->head; b != breaks->end; b = b->next) {
        if (b->address == address) {
            return 1;
        }
    }
    return 0;
}

int dbg_process(struct process *proc, pid_t pid, int child) {
    proc->id = pid;
    proc->child = child;
    list_init(&proc->breakpoints);
    list_init(&proc->threads);

    char taskpath[32];
    snprintf(taskpath, sizeof(taskpath), "/proc/%li/task", (long)pid);

    DIR *taskdir = opendir(taskpath);
    if (!taskdir) {
        return -1;
    }

    struct dirent *task;
    while ((task = readdir(taskdir))) {
        pid_t id = strtoul(task->d_name, NULL, 0);

        if (id != 0) {
            if (ptrace(PTRACE_SEIZE, id, NULL, PTRACE_OPTIONS) < 0) {
                closedir(taskdir);
                return -1;
            }

            struct thread *th = new_thread(proc, id);
            list_insert(proc->threads.end, th);
        }
    }

    if (INTERRUPT_ON_SEIZE) {
        interrupt_all_threads(proc);
    }

    closedir(taskdir);
    return 0;
}

int dbg_detach(struct process *proc) {
    interrupt_all_threads(proc);
    uninstall_breakpoints(proc);
    clear_breakpoints(proc);

    while (proc->threads.head != proc->threads.end) {
        struct thread *th = proc->threads.head;
        if (th->id >= 0) {
            ptrace(PTRACE_DETACH, th->id, NULL, th->signal);
        }
        clear_states(th);
        list_remove(th);
        free(th);
    }

    return 0;
}

int dbg_wait(struct thread *th, int dostops) {
    if (th->id < 0) {
        return -1;
    }

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

    if (WIFEXITED(status)) {
        th->id = -1;
        th->stopped = 1;
        th->signal = WEXITSTATUS(status);
        th->cont = 0;
        th->status = "EXITED";
        return 1;
    }

    if (WIFSIGNALED(status)) {
        th->id = -2;
        th->stopped = 1;
        th->signal = WTERMSIG(status);
        th->cont = 0;
        th->status = "TERMINATED";
        return 1;
    }

    unsigned long eventmsg;
    struct thread *newth;
    ptrace(PTRACE_GETEVENTMSG, th->id, NULL, &eventmsg);

    if (WIFSTOPPED(status)) {
        switch (status >> 8) {
            case SIGTRAP | (PTRACE_EVENT_CLONE << 8):
                newth = new_thread(th->proc, eventmsg);
                list_insert(th->proc->threads.end, newth);
                while (!dbg_wait(newth, 0)) {}

                th->stopped = 1;
                th->signal = 0;
                th->cont = 0;
                th->status = "CLONE EVENT";

                if (dostops) {
                    interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                }

                capture_state(th);
                return 1;

            case SIGTRAP | (PTRACE_EVENT_EXIT << 8):
                th->stopped = 1;
                th->signal = 0; //eventmsg;
                th->cont = 0;
                th->status = "EXIT EVENT";

                if (dostops) {
                    interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                }

                capture_state(th);
                return 1;

            case SIGTRAP | (PTRACE_EVENT_STOP << 8):
                th->stopped = 1;
                th->signal = 0;
                th->cont = 0;
                th->status = "STOP EVENT";

                if (dostops) {
                    interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                }

                capture_state(th);
                return 1;

            case SIGTRAP | 0x80:
                th->stopped = 1;
                th->signal = 0;
                th->cont = 0;
                th->status = "SYSCALL EVENT";

                if (dostops) {
                    interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                }

                capture_state(th);
                return 1;

            case SIGTRAP:
                th->stopped = 1;
                th->signal = 0;

                if (th->cont != 0) {
                    install_breakpoints(th);
                    ptrace(th->cont, th->id, NULL, NULL);
                    th->cont = 0;
                    th->stopped = 0;
                    th->status = "RUNNING";
                    return 0;
                }

                if (dostops) {
                    interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                }

                if (detect_breakpoint(th)) {
                    dbg_cont(th, PTRACE_CONT);
                    return 0;
                }

                th->status = "STEP/BREAKPOINT";
                capture_state(th);
                return 1;

            default:
                th->stopped = 1;
                th->signal = WSTOPSIG(status);
                th->cont = 0;
                th->status = "SIGNAL DELIVERY";

                if (dostops) {
                    interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                }

                capture_state(th);
                return 1;
        }
    }

    return -1;
}

int dbg_cont(struct thread *th, int cont) {
    if (th->id < 0 || !th->stopped) {
        return -1;
    }

    struct list *threads = &th->proc->threads;
    for (struct thread *t = threads->head; t != threads->end; t = t->next) {
        ptrace(PTRACE_SINGLESTEP, t->id, NULL, t->signal);

        t->stopped = 0;
        t->signal = 0;
        t->cont = cont;
        t->status = "RUNNING";
        t->clearstates = 1;
    }

    /* are there timing concerns here?  Critically, we
     * are doing this before the event loop has a chance to
     * wait on any thread and actually start it in the cont */
    /* cant do it here, since no thread is stopped (ptrace
     * returns ESRCH). */
    //install_breakpoints(th->proc);
    return 0;
}

int dbg_stepin(struct thread *th) {
    if (th->id < 0 || !th->stopped) {
        return -1;
    }

    if (ptrace(PTRACE_SINGLESTEP, th->id, NULL, th->signal) < 0) {
        return -1;
    }

    th->stopped = 0;
    th->signal = 0;
    th->cont = 0;
    th->status = "RUNNING";
    th->clearstates = 0;
    return 0;
}

int dbg_intr(struct thread *th) {
    if (th->id < 0 || th->stopped) {
        return -1;
    }

    ptrace(PTRACE_INTERRUPT, th->id, NULL, NULL);
    return 0;
}

void *deref(struct thread *th, unsigned long addr, size_t size) {
    (void)size;
    if (!th->state) {
        return NULL;
    }

    struct list *maps = &th->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;
}















///*
//static const char *strstatus(int status) {
//    static char buff[128];
//
//    if (WIFEXITED(status)) {
//        snprintf(buff, sizeof(buff), "exited: %i", WEXITSTATUS(status));
//    } else if (WIFSIGNALED(status)) {
//        snprintf(buff, sizeof(buff), "terminated: %s", strsignal(WTERMSIG(status)));
//    } else if (WIFSTOPPED(status)) {
//        snprintf(buff, sizeof(buff), "stopped: %s", strsignal(WSTOPSIG(status)));
//    } else {
//        snprintf(buff, sizeof(buff), "UNKNOWN STATUS VALUE");
//    }
//
//    return buff;
//}
//*/
//
//static int all_cont(struct process *proc) {
//    /* install breakpoints */
//    struct list *breaks = &proc->breakpoints;
//    for (struct breakpoint *b = breaks->head; b != breaks->end; b = b->next) {
//        if (b->enabled) {
//            unsigned long data;
//            data = ptrace(PTRACE_PEEKTEXT, proc->id, b->address, NULL);
//            b->orig = data;
//
//            data = (data & ~0xff) | BREAKPOINT_INSN;
//            ptrace(PTRACE_POKETEXT, proc->id, b->address, data);
//            b->active = 1;
//        }
//    }
//
//    /* continue all threads */
//    struct list *threads = &proc->threads;
//    for (struct thread *th = threads->head; th != threads->end; th = th->next) {
//        ptrace(PTRACE_);
//    }
//    return -1;
//}
//
///* current problem to solve:
// * if a SINGLESTEP is interrupted (it was over a blocking syscall for example),
// * a CONT out of that interruption will stop AS IF the SINGLESTEP had completed.
// */
///*
// * new:
// * next step is to trigger all threads starting on continue.
// */








/* PTRACE_GETSIGINFO ! */
/* handle ptrace() -> ESRCH as unexpectedly dead tracee */










//
//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->gid = 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, NULL)) {}
//    ptrace(PTRACE_SETOPTIONS, pid, NULL, PTRACE_O_EXITKILL | PTRACE_OPTIONS);
//    return 0;
//}
//
///* 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;
//}
//
//void dbg_free(struct tracee *dbg) {
//    while (dbg->breaks.head != dbg->breaks.end) {
//        struct breakpoint *b = dbg->breaks.head;
//        list_remove(b);
//        free(b);
//    }
//    clear_states(dbg);
//    free(dbg->buff);
//}