#include <signal.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/ptrace.h>
#include <unistd.h>

#include <capstone/capstone.h>

#include "console.h"
#include "debugger.h"
#include "helpers.h"
#include "list.h"

static PANEL *left, *right;

static pid_t parse_pid(const char *nptr) {
    char *endptr;
    pid_t pid = strtoul(nptr, &endptr, 0);
    return (*endptr ? 0 : pid);
}

static pid_t dofork(char **argv, struct console *cons) {
    pid_t pid = fork();
    if (pid < 0) {
        return -1;
    }

    if (pid == 0) {
        usleep(10000);
        console_configslave(cons);
        //close_range(STDERR_FILENO+1, ~0U, CLOSE_RANGE_UNSHARE);
        //raise(SIGSTOP); // ptrace(PTRACE_TRACEME, 0, NULL, NULL);
        execvp(argv[0], argv);
        exit(EXIT_FAILURE);
    }

    return pid;
}

static void list_breakpoints(struct thread *dbg, PANEL *pan) {
    struct list *breaks = &dbg->proc->breakpoints;
    if (breaks->head != breaks->end) {
        pprintw(pan, "---\n");

        for (struct breakpoint *bp=breaks->head; bp!=breaks->end; bp=bp->next) {
            pprintw(pan, "0x%lx (%c) (%i)   stack=0x%lx tid=%lu enabled=%d\n",
                    bp->address,
                    (bp->installed ? '*' : ' '),
                    bp->hits,
                    bp->stack, bp->tid, bp->enabled);
        }
    }
}

static void describe_states(struct thread *dbg, PANEL *pan) {
    struct list *states = &dbg->states;
    for (struct state *s = states->head; s != states->end; s = s->next) {
        pprintw(pan, "%c", (s == dbg->state ? '#' : '-'));
    }
    pprintw(pan, "\n");
}

static void dump_registers(struct thread *dbg, PANEL *pan) {
#ifdef ARCH_X86
    if (dbg->state->regsize == sizeof(struct user_regs_32)) {
        struct user_regs_32 *regs = &dbg->state->regs.x86_32;
        pprintw(pan, "orig_eax = 0x%08x\n", regs->orig_eax);
        pprintw(pan, "eax = 0x%08x\n", regs->eax);
        pprintw(pan, "ebx = 0x%08x\n", regs->ebx);
        pprintw(pan, "ecx = 0x%08x\n", regs->ecx);
        pprintw(pan, "edx = 0x%08x\n", regs->edx);
        pprintw(pan, "edi = 0x%08x\n", regs->edi);
        pprintw(pan, "esi = 0x%08x\n", regs->esi);
        pprintw(pan, "esp = 0x%08x\n", regs->esp);
        pprintw(pan, "ebp = 0x%08x\n", regs->ebp);
        pprintw(pan, "eip = 0x%08x\n", regs->eip);
    } else {
        struct user_regs_64 *regs = &dbg->state->regs.x86_64;
        pprintw(pan, "orig_rax = 0x%016llx\n", regs->orig_rax);
        pprintw(pan, "rax = 0x%016llx\n", regs->rax);
        pprintw(pan, "rbx = 0x%016llx\n", regs->rbx);
        pprintw(pan, "rcx = 0x%016llx\n", regs->rcx);
        pprintw(pan, "rdx = 0x%016llx\n", regs->rdx);
        pprintw(pan, "rdi = 0x%016llx\n", regs->rdi);
        pprintw(pan, "rsi = 0x%016llx\n", regs->rsi);
        pprintw(pan, "rsp = 0x%016llx\n", regs->rsp);
        pprintw(pan, "rbp = 0x%016llx\n", regs->rbp);
        pprintw(pan, "rip = 0x%016llx\n", regs->rip);
    }
#elif defined ARCH_AARCH64
    struct user_regs_64 *regs = &dbg->state->regs.arm64;
    pprintw(pan, "x0 = 0x%016llx\n", regs->regs[0]);
    pprintw(pan, "x1 = 0x%016llx\n", regs->regs[1]);
    pprintw(pan, "x2 = 0x%016llx\n", regs->regs[2]);
    pprintw(pan, "x3 = 0x%016llx\n", regs->regs[3]);
    pprintw(pan, "x4 = 0x%016llx\n", regs->regs[4]);
    pprintw(pan, "x5 = 0x%016llx\n", regs->regs[5]);
    pprintw(pan, "x6 = 0x%016llx\n", regs->regs[6]);
    pprintw(pan, "x7 = 0x%016llx\n", regs->regs[7]);
    pprintw(pan, "sp = 0x%016llx\n", regs->sp);
    pprintw(pan, "pc = 0x%016llx\n", regs->pc);
#endif
}

static void dump_stack(struct thread *dbg, PANEL *pan) {
    struct archinfo archinfo;
    struct iovec regs = { &dbg->state->regs, dbg->state->regsize };
    architecture_info(&archinfo, &regs);

    unsigned long sp = archinfo.stackptr;
    for (size_t i = 0; i < 16; i++, sp += archinfo.wordsize) {
        if (dbg->state->regsize == sizeof(struct user_regs_32)) {
            unsigned int *word = deref(dbg, sp, sizeof(unsigned int));
            pprintw(pan, "0x%lx:\t0x%08x\n", sp, *word);
        } else {
            unsigned long *word = deref(dbg, sp, sizeof(unsigned long));
            pprintw(pan, "0x%lx:\t0x%016lx\n", sp, *word);
        }
    }
}

static int rip_visited(struct thread *th, unsigned long rip) {
    struct list *states = &th->states;
    for (struct state *s = states->head; s != states->end; s = s->next) {
        struct archinfo archinfo;
        struct iovec regs = { &s->regs, s->regsize };
        architecture_info(&archinfo, &regs);

        if (rip == archinfo.progmctr) {
            return 1;
        }
    }
    return 0;
}

static void disasm(struct thread *dbg, PANEL *pan) {
    struct archinfo archinfo;
    struct iovec regs = { &dbg->state->regs, dbg->state->regsize };
    architecture_info(&archinfo, &regs);

    csh handle;
    cs_insn *insn;

    if (cs_open(archinfo.cs_arch, archinfo.cs_mode, &handle) != CS_ERR_OK) {
        //perror("capstone open");
    } else {
        uint64_t address = archinfo.progmctr;
        size_t codez = 128;
        const uint8_t *code = deref(dbg, address, codez);
        insn = cs_malloc(handle);

        for (size_t i = 0; i < 32; i++) {
            if (!cs_disasm_iter(handle, &code, &codez, &address, insn)) {
                break;
            }
            if (dbg->stopped) {
                if (insn->address == archinfo.progmctr) {
                    pattron(pan, COLOR_PAIR(1));
                } else if (get_breakpoint(dbg->proc, insn->address)) {
                    pattron(pan, COLOR_PAIR(3));
                } else if (rip_visited(dbg, insn->address)) {
                    pattron(pan, COLOR_PAIR(2));
                }
            }
            pprintw(pan, "0x%"PRIx64":\t%s %s\n", insn->address, insn->mnemonic, insn->op_str);
            if (dbg->stopped) {
                if (insn->address == archinfo.progmctr) {
                    pattroff(pan, COLOR_PAIR(1));
                } else if (get_breakpoint(dbg->proc, insn->address)) {
                    pattroff(pan, COLOR_PAIR(3));
                } else if (rip_visited(dbg, insn->address)) {
                    pattroff(pan, COLOR_PAIR(2));
                }
            }
        }

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

static void info_update(struct thread *th, PANEL *pan) {
    pclear(pan);
    pprintw(pan, "TID: %li\n", (long)th->id);
    pprintw(pan, "%s (%i: %s)\n", th->status, th->signal, strsignal(th->signal));
    list_breakpoints(th, pan);
    describe_states(th, pan);
    dump_registers(th, pan);
    pprintw(pan, "---\n");
    dump_stack(th, pan);
    pprintw(pan, "---\n");
    disasm(th, pan);
}

static int wait_all_threads(struct list *processes) {
    int action = 0;
    for (struct process *proc = processes->head; proc != processes->end; proc = proc->next) {
        action |= dbg_sync(proc);
    }
    return action;
}

static void layout(struct list *processes, struct thread *th) {
    int w = COLS/2;
    reset_panel(left, LINES-2, w, 1, 0);
    reset_panel(right, LINES-2, COLS-w, 1, w);

    clear();

    for (struct process *proc = processes->head; proc != processes->end; proc = proc->next) {
        if (th->proc == proc) {
            attron(COLOR_PAIR(1));
            printw("{ ");
        } else {
            attron(COLOR_PAIR(4));
            printw("{ ");
            attroff(COLOR_PAIR(4));
        }

        struct list *threads = &proc->threads;
        for (struct thread *t = threads->head; t != threads->end; t = t->next) {
            if (t == th) {
                printw("**");
            }
            printw("%li (%s)", (long)t->id, t->status);
            if (t == th) {
                printw("**");
            }

            if (t->next != threads->end) {
                printw("   ");
            }
        }


        if (th->proc == proc) {
            printw(" }  ");
            attroff(COLOR_PAIR(1));
        } else {
            attron(COLOR_PAIR(4));
            printw(" }  ");
            attroff(COLOR_PAIR(4));
        }
    }
}

int main(int argc, char **argv) {
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <pid | cmdline>\n", argv[0]);
        return EXIT_FAILURE;
    }

    getchar();

    struct console cons = {0};
    console_init(&cons);

    struct list processes = {0};
    struct thread *th = NULL;
    list_init(&processes);

    int child = 0;
    pid_t pid = parse_pid(argv[1]);
    argv[argc] = NULL;

    if (pid == 0) {
        pid = dofork(argv+1, &cons);
        child = 1;
    }

    struct process *proc = dbg_attach(pid, child);
    if (!proc) {
        fprintf(stderr, "Failed to attach to process %li\n", (long)pid);
        return EXIT_FAILURE;
    }
    list_insert(processes.end, proc);
    th = proc->threads.head;

    if (child) {
        dbg_cont(th);
    }

    cursinit();
    left = newpan(0, 0, 0, 0);
    right = newpan(0, 0, 0, 0);
    layout(&processes, th);

    int quit = 0;
    int mode = 0;
    int dirty = 1;
    int pressed = 0;
    while (!quit) {
        int ch = getch();
        pressed = 1;

        if (mode == 0) {
            switch (ch) {
                case KEY_RESIZE:
                    layout(&processes, th);
                    break;
                case 'q':
                    quit = 1;
                    break;
                case 'i':
                    mode = 1;
                    console_enter(&cons, right);
                    break;
                case 'j':
                    dbg_stepover(th);
                    break;
                case 'k':
                    dbg_stepback(th);
                    break;
                case 'l':
                    dbg_stepin(th);
                    break;
                case 'h':
                    /* todo: step out */
                    break;
                case 'g':
                    if (th->stopped) {
                        th->state = th->states.head;
                    }
                    break;
                case 'G':
                    if (th->stopped) {
                        th->state = th->states.tail;
                    }
                    break;
                case 's':
                    dbg_syscall(th);
                    break;
                case 'c':
                    dbg_cont(th);
                    break;
                case 'p':
                    dbg_intr(th);
                    break;
                case 't':
                    proc = th->proc;
                    th = th->next;
                    if (th == proc->threads.end) {
                        do {
                            proc = proc->next;
                        } while (proc == processes.end);
                        th = proc->threads.head;
                    }
                    layout(&processes, th);
                    break;
                case 'T':
                    proc = th->proc;
                    th = th->prev;
                    if (th == proc->threads.end) {
                        do {
                            proc = proc->prev;
                        } while (proc == processes.end);
                        th = proc->threads.tail;
                    }
                    layout(&processes, th);
                    break;
                case 'd':
                    proc = th->proc;
                    if (proc->child) {
                        break;
                    }
                    if (proc->prev == proc->next) {
                        break;
                    }
                    struct process *del = proc;
                    do {
                        proc = proc->next;
                    } while (proc == processes.end);
                    th = proc->threads.head;
                    dbg_detach(del);
                    break;
                case ':':
                    mvprintw(LINES-1, 0, ":");
                    curs_set(TRUE);
                    echo();
                    timeout(-1);
                    char cmd[128] = {0};
                    getstr(cmd);
                    curs_set(FALSE);
                    noecho();
                    timeout(25);
                    clear();
                    layout(&processes, th);

                    char *t = cmd;
                    int en = 1;
                    pid_t tid = 0;
                    if (t[0] == '!') {
                        en = -1;
                        t++;
                    } else if (t[0] == '#') {
                        en = 0;
                        t++;
                    } else if (t[0] == '@') {
                        tid = th->id;
                        t++;
                    }
                    unsigned long address = strtoul(t, NULL, 0);
                    struct breakpoint *b = add_breakpoint(th->proc, address);
                    b->enabled = en;
                    b->tid = tid;
                    break;
                case ERR:
                    pressed = 0;
                    break;
            }
        } else {
            switch (ch) {
                case KEY_RESIZE:
                    layout(&processes, th);
                    break;
                case KEY_ESCAPE:
                    mode = 0;
                    console_leave(&cons, right);
                    break;
                case ERR:
                    pressed = 0;
                    break;
                default:
                    console_input(&cons, ch);
                    break;
            }
        }

        dirty |= wait_all_threads(&processes);
        dirty |= console_update(&cons, right);

        if (dirty || pressed) {
            layout(&processes, th);
            info_update(th, left);
            cursupdate();
            dirty = 0;
        }
    }

    for (struct process *p = processes.head; p != processes.end; p = p->next) {
        struct process *del = p;
        p = p->prev;
        dbg_detach(del);
    }

    endwin();
    return EXIT_SUCCESS;
}