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

#include <capstone/capstone.h>

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

static PANEL *left, *right;
static struct console cons;
static int mode = 0;

static void describe_status(struct tracee *_dbg, PANEL *pan) {
    struct tracee dbg = *_dbg;
    int status = dbg.status;
    struct ptrace_syscall_info info;

    if (WIFEXITED(status)) {
        pprintw(pan, "exited with code: %u\n", WEXITSTATUS(status));
    } else if (WIFSIGNALED(status)) {
        pprintw(pan, "terminated by signal: %s\n", strsignal(WTERMSIG(status)));
    } else {
        unsigned long msg;
        if (ptrace(PTRACE_GETEVENTMSG, dbg.id, NULL, &msg) < 0) {
            perror("PTRACE_GETEVENTMSG");
        }

        if (WIFSTOPPED(status)) {
            switch (status >> 8) {
                case ((PTRACE_EVENT_VFORK << 8) | SIGTRAP):
                    pprintw(pan, "child vfork()'d to: %lu\n", msg);
                    break;
                case ((PTRACE_EVENT_FORK << 8) | SIGTRAP):
                    pprintw(pan, "child fork()'d to: %lu\n", msg);
                    break;
                case ((PTRACE_EVENT_CLONE << 8) | SIGTRAP):
                    pprintw(pan, "child clone()'d to: %lu\n", msg);
                    break;
                case ((PTRACE_EVENT_VFORK_DONE << 8) | SIGTRAP):
                    pprintw(pan, "finished vfork\n");
                    break;
                case ((PTRACE_EVENT_EXEC << 8) | SIGTRAP):
                    pprintw(pan, "child exec()'d from: %lu\n", msg);
                    break;
                case ((PTRACE_EVENT_EXIT << 8) | SIGTRAP):
                    pprintw(pan, "exiting with code: %lu\n", msg);
                    break;
                case ((PTRACE_EVENT_STOP << 8) | SIGTRAP):
                    pprintw(pan, "child stopped\n");
                    break;
                case ((PTRACE_EVENT_SECCOMP << 8) | SIGTRAP):
                    pprintw(pan, "child seccomp event\n");
                    break;

                case (SIGTRAP | 0x80):
                    ptrace(PTRACE_GET_SYSCALL_INFO, dbg.id, sizeof(info), &info);
                    pprintw(pan, "child entering syscall: %llu (%llx, %llx, %llx, %llx, %llx, %llx)\n",
                            info.entry.nr,
                            info.entry.args[0],
                            info.entry.args[1],
                            info.entry.args[2],
                            info.entry.args[3],
                            info.entry.args[4],
                            info.entry.args[5]);
                    break;

                case SIGSTOP:
                case SIGTSTP:
                case SIGTTIN:
                case SIGTTOU:
                    pprintw(pan, "child group-stopped\n");
                    break;

                default:
                    pprintw(pan, "received signal: %s\n", strsignal(WSTOPSIG(status)));
                    break;
            }
        } else {
            pprintw(pan, "child stop event unrecognized\n");
        }
    }
}

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

        for (struct breakpoint *bp=breaks->head; bp!=breaks->end; bp=bp->next) {
            pprintw(pan, "0x%lx  ", bp->address);
        }
        pprintw(pan, "\n");
    }
}

static void describe_states(struct tracee *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 tracee *dbg, PANEL *pan) {
    struct user_regs_struct *regs = &dbg->state->regs;
    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);
}

static void dump_stack(struct tracee *dbg, PANEL *pan) {
    unsigned long sp = dbg->state->regs.rsp;
    for (size_t i = 0; i < 16; i++, sp += 8) {
        unsigned long word = *(unsigned long *)deref(dbg, sp,sizeof(unsigned long));
        pprintw(pan, "0x%lx:\t0x%lx\n", sp, word);
    }
}

static void disasm(struct tracee *dbg, PANEL *pan) {
    csh handle;
    cs_insn *insn;

    if (cs_open(CS_ARCH_X86, CS_MODE_64, &handle) != CS_ERR_OK) {
        perror("capstone open");
    } else {
        uint64_t address = dbg->state->regs.rip;
        size_t codez = 128;
        const uint8_t *code = deref(dbg, address, codez);
        insn = cs_malloc(handle);

        for (size_t i = 0; i < 16; i++) {
            if (!cs_disasm_iter(handle, &code, &codez, &address, insn)) {
                break;
            }
            pprintw(pan, "0x%"PRIx64":\t%s %s\n", insn->address, insn->mnemonic, insn->op_str);
        }

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

static void info_update(struct tracee *dbg, PANEL *pan) {
    pclear(pan);
    pprintw(pan, "PID: %li\n", (long)dbg->id);
    if (!dbg->stopped) {
        pprintw(pan, "Process is running...\n");
    } else {
        describe_status(dbg, pan);
        list_breakpoints(dbg, pan);
        describe_states(dbg, pan);
        dump_registers(dbg, pan);
        pprintw(pan, "---\n");
        dump_stack(dbg, pan);
        pprintw(pan, "---\n");
        disasm(dbg, pan);
    }
}

static void layout(void) {
    int w = COLS/2;
    reset_panel(left, LINES-1, w, 0, 0);
    reset_panel(right, LINES-1, COLS-w, 0, w);
}

int main(int argc, char **argv) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <stop> <command>\n", argv[0]);
        return 1;
    }

    cursinit();
    left = newpan(0, 0, 0, 0);
    right = newpan(0, 0, 0, 0);
    layout();

    console_init(&cons);

    argv[argc] = NULL;
    struct tracee dbg;
    if (dbg_new_process(&dbg, argv+2, &cons)) {
        pprintw(right, "Failed to start child\n");
    }

    unsigned long stop = strtoul(argv[1], NULL, 0);
    if (stop != 0) {
        struct breakpoint *b = xmalloc(sizeof(*b));
        b->address = stop;
        b->stack = 0;
        b->enabled = -1;
        b->active = 0;
        list_insert(dbg.breaks.end, b);
        dbg_cont(&dbg, PTRACE_CONT);
    }

    int quit = 0;
    while (!quit) {
        //if (dbg.stopped == NULL) {
            dbg_wait(&dbg);
        //}
        console_update(&cons, right);
        info_update(&dbg, left);
        cursupdate();

        int ch = getch();

        if (mode == 0) {
            switch (ch) {
                case KEY_RESIZE:
                    layout();
                    break;
                case 'q':
                    quit = 1;
                    break;
                case 'i':
                    mode = 1;
                    console_enter(&cons, right);
                    break;
                case 'j':
                    if (dbg.stopped) {
                        if (dbg.state != dbg.states.tail) {
                            dbg.state = dbg.state->next;
                        } else {
                            dbg_stepover(&dbg);
                        }
                    }
                    break;
                case 'k':
                    if (dbg.stopped) {
                        if (dbg.state != dbg.states.head) {
                            dbg.state = dbg.state->prev;
                        }
                    }
                    break;
                case 'l':
                    if (dbg.stopped) {
                        if (dbg.state != dbg.states.tail) {
                            //dbg.state = dbg.state->next;
                        } else {
                            dbg_stepin(&dbg);
                        }
                    }
                    break;
                //case 'h':
                //    if (dbg.stopped) {
                //        dbg_stepout(&dbg);
                //    }
                //    break;
                case 'g':
                    if (dbg.stopped) {
                        dbg.state = dbg.states.head;
                    }
                    break;
                case 'G':
                    if (dbg.stopped) {
                        dbg.state = dbg.states.tail;
                    }
                    break;
                case 's':
                    if (dbg.stopped) {
                        dbg_cont(&dbg, PTRACE_SYSCALL);
                    }
                    break;
                case 'c':
                    if (dbg.stopped) {
                        dbg_cont(&dbg, PTRACE_CONT);
                    }
                    break;
                case 'p':
                    if (!dbg.stopped) {
                        tgkill(dbg.id, dbg.id, SIGSTOP);
                    }
                    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();
                    refresh();

                    char *t = cmd;
                    struct breakpoint *b = xmalloc(sizeof(*b));
                    b->enabled = 1;
                    b->active = 0;
                    if (t[0] == '!') {
                        b->enabled = -1;
                        t++;
                    }
                    b->address = strtoul(t, NULL, 0);
                    b->stack = 0;
                    list_insert(dbg.breaks.end, b);
                    break;
            }
        } else {
            switch (ch) {
                case KEY_RESIZE:
                    layout();
                    break;
                case 0x1b:
                    mode = 0;
                    console_leave(&cons, right);
                    break;
                case ERR:
                    break;
                default:
                    console_input(&cons, ch);
                    break;
            }
        }
    }

    endwin();
    return 0;
}