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

#include <capstone/capstone.h>

#include "debugger.h"
#include "helpers.h"

#define BREAKPOINT_INSN 0xcc

//struct list global_processes = {0};
//struct thread *global_thread = NULL;

static const int PTRACE_OPTIONS =
    PTRACE_O_TRACECLONE |
    PTRACE_O_TRACEEXIT  |
    PTRACE_O_TRACESYSGOOD;

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

    struct user_regs_struct regs;
    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) {
        if (b->installed) {
            if (!check) {
                if (regs.rip - 1 == b->address) {
                    regs.rip--;
                    ptrace(PTRACE_SETREGS, th->id, NULL, &regs);
                    check = 1;
                    b->hits++;

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

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

                    if (!b->enabled) {
                        restart = 1;
                    }
                }
            }

            /* this needs moved, we have to finish the for-loop
             * to actually know if restart is true or false... */
            if (b->enabled < 0 && !restart) {
                struct breakpoint *del = b;
                b = b->prev;
                list_remove(del);
                free(del);
            }
        }
    }

    return restart;
}

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

            word = (word & ~0xff) | BREAKPOINT_INSN;
            ptrace(PTRACE_POKETEXT, proc->id, b->address, word);
            b->installed = 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->installed) {
            ptrace(PTRACE_POKETEXT, proc->id, b->address, b->text);
            b->installed = 0;
        }
    }
}

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 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_thread(struct thread *th) {
    if (th->clearstates) {
        free_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 void capture_state(struct thread *th, int all) {
    if (all) {
        struct list *threads = &th->proc->threads;
        for (struct thread *t = threads->head; t != threads->end; t = t->next) {
            capture_state_thread(t);
        }
    } else {
        capture_state_thread(th);
    }
}

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

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

static int interrupt_all_threads(struct process *proc) {
    int stopped = 0;
    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, 1)) {}
            stopped = 1;
        }
    }
    return stopped;
}

void add_breakpoint(struct process *proc, unsigned long address, unsigned long stack, pid_t tid, int enabled) {
    struct breakpoint *b = xmalloc(sizeof(*b));
    b->address = address;
    b->text = 0;
    b->installed = 0;
    b->hits = 0;
    b->stack = stack;
    b->tid = tid;
    b->enabled = enabled;
    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_attach(pid_t pid, int child) {
struct process *dbg_attach(pid_t pid, int child) {
    //if (global_processes.end == NULL) {
    //    list_init(&global_processes);
    //}

    struct process *proc = new_process(pid, child);
    //list_insert(global_processes.end, proc);

    int options = PTRACE_OPTIONS;
    if (child) {
        options |= PTRACE_O_EXITKILL;
    }

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

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

    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, options) < 0) {
                closedir(taskdir);
                dbg_detach(proc);
                return NULL;
            }

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

    closedir(taskdir);

    //global_thread = proc->threads.head;
    interrupt_all_threads(proc);
    return proc;
}

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

    if (proc->child) {
        kill(proc->id, SIGKILL);
    }

    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);
        }
        free_states(th);
        list_remove(th);
        free(th);
    }

    list_remove(proc);
    free(proc);
    return 0;
}

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

    /* todo: check what happens if we call on an already known
     * stopped thread? */
    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;
    }

    int stopped;
    struct thread *newth;

    unsigned long eventmsg;
    ptrace(PTRACE_GETEVENTMSG, th->id, NULL, &eventmsg);

    if (WIFSTOPPED(status)) {
        switch (status >> 8) {
            /* todo: other ptrace event stops */
            case SIGTRAP | (PTRACE_EVENT_CLONE << 8):
                newth = new_thread(th->proc, eventmsg);
                list_insert(th->proc->threads.end, newth);
                while (!dbg_wait(newth, 1)) {}

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

                if (!recursion) {
                    stopped = interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                    capture_state(th, stopped);
                }

                return 1;

            case SIGTRAP | (PTRACE_EVENT_EXIT << 8):
                th->stopped = 1;
                th->signal = 0; /* eventmsg has exit code, but would inject sig */
                th->cont = 0;
                th->status = "EXIT EVENT";

                if (!recursion) {
                    stopped = interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                    capture_state(th, stopped);
                }

                return 1;

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

                if (!recursion) {
                    stopped = interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                    capture_state(th, stopped);
                }

                return 1;

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

                if (!recursion) {
                    stopped = interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                    capture_state(th, stopped);
                }

                return 1;

            case SIGTRAP:
                th->stopped = 1;
                th->signal = 0;
                th->status = "STEP/BREAKPOINT";

                if (th->cont != 0) {
                    /* gdb this portion.  are there race conditions
                     * that matter?? */
                    install_breakpoints(th->proc);
                    ptrace(th->cont, th->id, NULL, NULL);
                    th->cont = 0;
                    th->stopped = 0;
                    th->status = "RUNNING";
                    return 0;
                }

                /* todo: Test two threads hitting a breakpoint at
                 * the same time. */
                int restart = detect_breakpoint(th);

                if (!recursion) {
                    stopped = interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                    if (!restart) {
                        capture_state(th, stopped);
                    }
                }

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

                return 1;

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

                if (!recursion) {
                    stopped = interrupt_all_threads(th->proc);
                    uninstall_breakpoints(th->proc);
                    capture_state(th, stopped);
                }

                return 1;
        }
    }

    return -1;
}

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

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

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

    return 0;
}

int dbg_step(struct thread *th, int stepover) {
    // todo: support step-out
    //if (th->id < 0 || !th->stopped) {
    //    return -1;
    //}

    if (!th->stopped) {
        return -1;
    }

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

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

    csh cshandle;
    cs_open(CS_ARCH_X86, CS_MODE_64, &cshandle);
    cs_insn *insn = cs_malloc(cshandle);

    uint64_t address = th->state->regs.rip;
    size_t size = 128;
    const uint8_t *code = deref(th, address, size);

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

    if (insn->id == X86_INS_CALL && stepover) {
        add_breakpoint(th->proc, address, th->state->regs.rsp, th->id, -1);
        th->cont = PTRACE_CONT;
    } else {
        th->cont = 0;
    }

    ptrace(PTRACE_SINGLESTEP, th->id, NULL, th->signal);

    th->stopped = 0;
    th->signal = 0;
    th->status = "RUNNING";

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

int dbg_pets(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;
}





///* 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.
// */

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