#include <dirent.h>
#include <elf.h>
#include <signal.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 "debugger.h"
#include "helpers.h"

static const int PTRACE_OPTIONS =
    PTRACE_O_TRACECLONE |
    PTRACE_O_TRACEFORK  |
    PTRACE_O_TRACEEXEC  |
    PTRACE_O_TRACEEXIT  |
    PTRACE_O_TRACESYSGOOD;

static const int PTRACE_CHILD_OPTIONS = PTRACE_OPTIONS | PTRACE_O_EXITKILL;
//static const int STOP_ALL_ON_EVENT = 1;
static const useconds_t SCHEDULER_DELAY = 10000;

static struct process *new_process(pid_t id, int child) {
    struct process *proc = xmalloc(sizeof(*proc));
    proc->id = id;
    proc->child = child;
    list_init(&proc->breakpoints);
    list_init(&proc->threads);
    memset(proc->status, 0, sizeof(proc->status));
    return proc;
}

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->doing = 0;
    th->donext = 0;
    strcpy(th->status, "RUNNING");
    return th;
}

static struct thread *thread_by_id(struct process *proc, pid_t id) {
    struct list *threads = &proc->threads;
    for (struct thread *th = threads->head; th != threads->end; th = th->next) {
        if (th->id == id) {
            return th;
        }
    }
    return NULL;
}

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

static void install_breakpoints(struct thread *th) {
    struct archinfo archinfo;
    struct iovec regs = { &th->state->regs, th->state->regsize };
    architecture_info(&archinfo, &regs);

    struct list *breaks = &th->proc->breakpoints;
    for (struct breakpoint *b = breaks->head; b != breaks->end; b = b->next) {
        if (!b->installed) {
            unsigned long word;
            word = ptrace(PTRACE_PEEKTEXT, th->id, b->address, NULL);
            b->text = word;

            word = (word & ~archinfo.bp_mask) | archinfo.bp_insn;
            ptrace(PTRACE_POKETEXT, th->id, b->address, word);
            b->installed = 1;
            b->previously_installed = 1;
        }
    }
}

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

        if (b->previously_installed && b->enabled < 0) {
            struct thread *t = NULL;
            if (b->tid == 0 || ((t = thread_by_id(th->proc, b->tid)) && !t->doing)) {
                struct breakpoint *del = b;
                b = b->next;
                list_remove(del);
                free(del);
            }
        }
    }
}

static int detect_breakpoint(struct thread *th, int *restart) {
    int ret = 0;
    *restart = 0;

    /* Hack: Need to manually fetch registers here, since capture_state() has
     * not yet run for this stop.  It is not guaranteed that we even want to
     * preserve the current state. */
    user_regs_t regs;
    struct iovec ivregs = { &regs, sizeof(regs) };
    ptrace(PTRACE_GETREGSET, th->id, NT_PRSTATUS, &ivregs);

    struct archinfo archinfo;
    architecture_info(&archinfo, &ivregs);

    unsigned long breakpt_address = archinfo.progmctr - archinfo.bp_adjust;
    struct breakpoint *b = get_breakpoint(th->proc, breakpt_address);
    if (b && b->installed && th->doing != PTRACE_SINGLESTEP) {
        /* restore actual program counter to breakpoint address */
        if (ivregs.iov_len == sizeof(struct user_regs_32)) {
            regs.PROGMCTR_32 = breakpt_address;
        } else {
            regs.PROGMCTR_64 = breakpt_address;
        }
        ptrace(PTRACE_SETREGSET, th->id, NT_PRSTATUS, &ivregs);

        b->hits++; /* todo: consider whether this is firing too much */
        ret = b->user;

        if (b->stack != 0 && b->stack != archinfo.stackptr) {
            *restart = 1;
        } else if (b->tid != 0 && b->tid != th->id) {
            *restart = 1;
        } else if (!b->enabled) {
            *restart = 1;
        }
    }

    return ret;
}

static void free_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 process *proc) {
    struct list *threads = &proc->threads;
    for (struct thread *th = threads->head; th != threads->end; th = th->next) {
        if (!th->state) {
            if (th->clearstates) {
                free_states(th);
            }

            struct state *s = xmalloc(sizeof(*s));
            struct iovec regs = { &s->regs, sizeof(s->regs) };
            ptrace(PTRACE_GETREGSET, th->id, NT_PRSTATUS, &regs);
            s->regsize = regs.iov_len;
            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 char thread_state(struct thread *th) {
    char statpath[32], stat[512];
    snprintf(statpath, sizeof(statpath), "/proc/%li/stat", (long)th->id);

    FILE *f = fopen(statpath, "r");
    if (f) {
        fread(stat, 1, sizeof(stat), f);
        fclose(f);

        char *c = strchr(stat, ' ');
        c = strchr(c+1, ' ') + 1;
        return *c;
    }

    return 0;
}

static int wait_thread(struct thread *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);

            char state;
            do {
                state = thread_state(th);
            } while (state != 't' && state != 'D' && state != 'Z');

            wait_thread(th);

            //if (STOP_ALL_ON_EVENT) {
            //    th->cont = 0;
            //    th->shouldcont = 0;
            //}
        }
    }
}

static void continue_all_threads(struct process *proc) {
    struct list *threads = &proc->threads;
    for (struct thread *th = threads->head; th != threads->end; th = th->next) {
        if (th->id > 0) {
            char state;
            do {
                wait_thread(th);

                ptrace(PTRACE_CONT, th->id, NULL, th->signal);
                th->stopped = 0;
                th->signal = 0;
                th->donext = 0;
                th->doing = PTRACE_CONT;
                strcpy(th->status, "RUNNING");
                th->clearstates = 1;

                state = thread_state(th);
            } while (state == 't');
        }
    }
}

static void resume_threads(struct process *proc) {
    struct list *threads = &proc->threads;
    int once = 0;

    for (struct thread *th = threads->head; th != threads->end; th = th->next) {
        if (th->stopped && th->doing == PTRACE_SINGLESTEP) {
            ptrace(PTRACE_SINGLESTEP, th->id, NULL, th->signal);
            th->stopped = 0;
            th->signal = 0;
            strcpy(th->status, "RUNNING");
        }
    }

    for (struct thread *th = threads->head; th != threads->end; th = th->next) {
        if (th->stopped && th->doing && th->doing != PTRACE_SINGLESTEP) {
            if (!once) {
                usleep(SCHEDULER_DELAY);
                once = 1;
            }
            install_breakpoints(th);
            ptrace(th->doing, th->id, NULL, th->signal);
            th->stopped = 0;
            th->signal = 0;
            strcpy(th->status, "RUNNING");
        }
    }
}

static int wait_thread(struct thread *th) {
    if (th->id <= 0) {
        return -1;
    }

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

    if (WIFEXITED(status)) {
        th->id = 0;
        th->stopped = 1;
        th->signal = WEXITSTATUS(status);
        th->doing = 0;
        th->donext = 0;
        strcpy(th->status, "EXITED");
        return 1;
    } else if (WIFSIGNALED(status)) {
        th->id = 0;
        th->stopped = 1;
        th->signal = WTERMSIG(status);
        th->doing = 0;
        th->donext = 0;
        strcpy(th->status, "TERMINATED");
        return 1;
    } else if (WIFSTOPPED(status)) {
        struct process *eventproc;
        struct thread *eventth;
        unsigned long eventmsg;
        ptrace(PTRACE_GETEVENTMSG, th->id, NULL, &eventmsg);

        /* todo: process status messages */
        switch (status >> 8) {
            /* todo: other ptrace event stops */
            case SIGTRAP | (PTRACE_EVENT_CLONE << 8):
                eventth = new_thread(th->proc, eventmsg);
                list_insert(th->proc->threads.end, eventth);
                while (!wait_thread(eventth)) {}

                th->stopped = 1;
                th->signal = 0;
                th->doing = 0;
                th->donext = 0;
                strcpy(th->status, "CLONE");
                th->state = NULL;
                return 1;

            case SIGTRAP | (PTRACE_EVENT_FORK << 8):
                eventproc = new_process(eventmsg, th->proc->child);
                eventth = new_thread(eventproc, eventmsg);
                list_insert(eventproc->threads.end, eventth);
                list_insert(th->proc->next, eventproc);
                //while (!wait_thread(eventth)) {}
                dbg_sync(eventproc);

                th->stopped = 1;
                th->signal = 0;
                th->doing = 0;
                th->donext = 0;
                strcpy(th->status, "FORK");
                th->state = NULL;
                return 1;

            case SIGTRAP | (PTRACE_EVENT_EXEC << 8):
                eventth = thread_by_id(th->proc, eventmsg);
                if (eventth->id != th->id) {
                    eventth->id = 0;
                    eventth->stopped = 1;
                    eventth->signal = 0;
                    eventth->doing = 0;
                    eventth->donext = 0;
                    strcpy(eventth->status, "EXITED");
                }

                th->stopped = 1;
                th->signal = 0;
                th->doing = 0;
                th->donext = 0;
                strcpy(th->status, "EXECVE");
                th->state = NULL;
                return 1;

            case SIGTRAP | (PTRACE_EVENT_EXIT << 8):
                th->stopped = 1;
                th->signal = 0;
                th->doing = PTRACE_CONT;
                th->donext = 0;
                strcpy(th->status, "EXITING");
                th->state = NULL;
                return 1;

            case SIGTRAP | (PTRACE_EVENT_STOP << 8):
                th->stopped = 1;
                th->signal = 0;
                strcpy(th->status, "STOPPED");
                if (!th->doing) {
                    th->state = NULL;
                }
                return 1;

            case SIGTRAP | 0x80:
                th->stopped = 1;
                th->signal = 0;
                th->doing = 0;
                th->donext = 0;
                strcpy(th->status, "SYSCALL");
                th->state = NULL;
                return 1;

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

                int restart;
                int bp = detect_breakpoint(th, &restart);
                strcpy(th->status, (bp ? "BREAKPOINT" : "STEP"));

                if (restart) {
                    th->donext = th->doing;
                    th->doing = (th->doing ? PTRACE_SINGLESTEP : 0);
                } else {
                    th->doing = th->donext;
                    th->donext = 0;
                }

                if (!th->doing) {
                    th->state = NULL;
                }

                return 1;

            default:
                th->stopped = 1;
                th->signal = WSTOPSIG(status);
                th->doing = 0;
                th->donext = 0;
                strcpy(th->status, "SIGNAL DELIVERY");
                th->state = NULL;
                return 1;
        }
    }

    return -1;
}

struct breakpoint *add_breakpoint(struct process *proc, unsigned long address) {
    struct breakpoint *b = xmalloc(sizeof(*b));
    b->address = address;
    b->text = 0;
    b->installed = 0;
    b->previously_installed = 0;
    b->hits = 0;
    b->user = 1;
    b->stack = 0;
    b->tid = 0;
    b->enabled = 1;
    list_insert(proc->breakpoints.end, b);
    return b;
}

struct breakpoint *get_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 b;
        }
    }
    return NULL;
}

struct process *dbg_attach(pid_t pid, int child) {
    const int OPTIONS = (child ? PTRACE_CHILD_OPTIONS : PTRACE_OPTIONS);

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

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

    struct dirent *task;
    struct process *proc = new_process(pid, child);

    while ((task = readdir(taskdir))) {
        pid_t id = strict_strtoul(task->d_name, 0);

        if (id != 0) {
            if (ptrace(PTRACE_SEIZE, id, NULL, OPTIONS) < 0) {
                closedir(taskdir);
                dbg_detach(proc);
                return NULL;
            }
            struct thread *th = new_thread(proc, id);
            list_insert(proc->threads.end, th);
        }
    }

    closedir(taskdir);

    interrupt_all_threads(proc);
    capture_state(proc);
    return proc;
}

void dbg_detach(struct process *proc) {
    /* ptrace requires that threads be stopped before calling PTRACE_DETACH.
     * In some cases, a blocked thread can possibly have a pending trap which
     * would crash the process if encountered after ptracing.  We cycle an extra
     * continue/interrupt to consume this trap event ourselves before
     * detaching. */
    interrupt_all_threads(proc);
    continue_all_threads(proc);
    interrupt_all_threads(proc);

    /* Supplement to PTRACE_O_EXITKILL */
    if (proc->child) {
        kill(proc->id, SIGKILL);
    }

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

    free_breakpoints(proc);
    list_remove(proc);
    free(proc);
}

int dbg_free(struct thread *th) {
    if (th->id <= 0) {
        free_states(th);
        list_remove(th);
        free(th);
        return 0;
    }
    return -1;
}

int dbg_sync(struct process *proc) {
    struct thread *acted = NULL;

    struct list *threads = &proc->threads;
    for (struct thread *th = threads->head; th != threads->end; th = th->next) {
        if (th->id > 0) {
            if (/*!th->stopped &&*/ wait_thread(th)) {
                acted = th;
            } else if (th->stopped && th->doing) {
                acted = th;
            }
        }
    }

    if (acted) {
        interrupt_all_threads(proc);
        uninstall_breakpoints(acted);
        capture_state(proc);
        resume_threads(proc);
    }

    return acted != NULL;
}

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

    /* todo: move this into dbg_sync? */
    /* probably not */
    ptrace(PTRACE_INTERRUPT, th->id, NULL, NULL);
    th->doing = 0;
    th->donext = 0;
    return 0;
}

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

    th->doing = PTRACE_SINGLESTEP;
    th->donext = PTRACE_CONT;
    th->clearstates = 1;
    return 0;
}

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

    th->doing = PTRACE_SINGLESTEP;
    th->donext = PTRACE_SYSCALL;
    th->clearstates = 1;
    return 0;
}

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

    if (th->state != th->states.tail) {
        th->state = th->state->next;
        return 0;
    }

    if (th->id <= 0) {
        return -1;
    }

    th->doing = PTRACE_SINGLESTEP;
    th->donext = 0;
    th->clearstates = 0;
    return 0;
}

int dbg_stepover(struct thread *th) {
    if (!th->stopped) {
        return -1;
    }

    if (th->state != th->states.tail) {
        th->state = th->state->next;
        return 0;
    }

    if (th->id <= 0) {
        return -1;
    }

    struct archinfo archinfo;
    struct iovec regs = { &th->state->regs, th->state->regsize };
    architecture_info(&archinfo, &regs);

    csh handle;
    cs_open(archinfo.cs_arch, archinfo.cs_mode, &handle);
    cs_insn *insn = cs_malloc(handle);

    uint64_t address = archinfo.progmctr;
    size_t size = 128;
    const uint8_t *code = deref(th, address, size);

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

    if (insn->id == archinfo.cs_call) {
        struct breakpoint *b = add_breakpoint(th->proc, address);
        b->user = 0;
        b->stack = archinfo.stackptr;
        b->tid = th->id;
        b->enabled = -1;

        th->doing = PTRACE_SINGLESTEP;
        th->donext = PTRACE_CONT;
        th->clearstates = 0;
    } else {
        th->doing = PTRACE_SINGLESTEP;
        th->donext = 0;
        th->clearstates = 0;
    }

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

int dbg_stepback(struct thread *th) {
    if (!th->stopped) {
        return -1;
    }

    if (th->state != th->states.head) {
        th->state = th->state->prev;
        return 0;
    }

    return -1;
}

void *deref(struct thread *th, unsigned long address, size_t size) {
    /* todo: size, and mapping breaks */
    (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 <= address && address < m->end) {
            return (unsigned char *)m->data + (address - m->start);
        }
    }

    return NULL;
}






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