#include <nxos.h>
#include <string.h>
#include <stdlib.h>
#include <nxos/ptrace.h>
#include "gdb.h"

const struct reg_descriptor g_register_descriptors[] = {
    {GDB_CPU_RISCV_REG_X0, "x0"},
    {GDB_CPU_RISCV_REG_RA, "ra"},
    {GDB_CPU_RISCV_REG_SP, "sp"},
    {GDB_CPU_RISCV_REG_GP, "gp"},
    {GDB_CPU_RISCV_REG_TP, "tp"},
    {GDB_CPU_RISCV_REG_T0, "t0"},
    {GDB_CPU_RISCV_REG_T1, "t1"},
    {GDB_CPU_RISCV_REG_T2, "t2"},
    {GDB_CPU_RISCV_REG_S0, "s0"},
    {GDB_CPU_RISCV_REG_S1, "s1"},
    {GDB_CPU_RISCV_REG_A0, "a0"},
    {GDB_CPU_RISCV_REG_A1, "a1"},
    {GDB_CPU_RISCV_REG_A2, "a2"},
    {GDB_CPU_RISCV_REG_A3, "a3"},
    {GDB_CPU_RISCV_REG_A4, "a4"},
    {GDB_CPU_RISCV_REG_A5, "a5"},
    {GDB_CPU_RISCV_REG_A6, "a6"},
    {GDB_CPU_RISCV_REG_A7, "a7"},
    {GDB_CPU_RISCV_REG_S2, "s2"},
    {GDB_CPU_RISCV_REG_S3, "s3"},
    {GDB_CPU_RISCV_REG_S4, "s4"},
    {GDB_CPU_RISCV_REG_S5, "s5"},
    {GDB_CPU_RISCV_REG_S6, "s6"},
    {GDB_CPU_RISCV_REG_S7, "s7"},
    {GDB_CPU_RISCV_REG_S8, "s8"},
    {GDB_CPU_RISCV_REG_S9, "s9"},
    {GDB_CPU_RISCV_REG_S10, "s10"},
    {GDB_CPU_RISCV_REG_S11, "s11"},
    {GDB_CPU_RISCV_REG_T3, "t3"},
    {GDB_CPU_RISCV_REG_T4, "t4"},
    {GDB_CPU_RISCV_REG_T5, "t5"},
    {GDB_CPU_RISCV_REG_T6, "t6"},
    {GDB_CPU_RISCV_REG_PC, "pc"},
};

static gdb_bkp_or_watch_t gdb_bkp_buf[__GDB_MAX_BREAKPOINTS_NUM_];
static gdb_bkp_or_watch_t gdb_watch_buf[__GDB_MAX_WATCHPOINTS_NUM_];

int isChildStop = 0;

int WaitForChild(){
    while(isChildStop == 0){
        NX_Printf("child process running\n");
        NX_ThreadSleep(1000);
    }
    isChildStop = 0;
    return 0;
}

int SetRegisterValue(NX_Solt tid, int reg, unsigned long long value){
    if(reg <= 0 || reg >= GDB_CPU_NUM_REGISTERS){
        return -1;    
    }
    Registers regs;
    if(NX_Ptrace(PTRACE_GETREGS, tid, NULL, &regs) != 0){
        return -1;
    }
    regs.reg[reg] = value;
    if(NX_Ptrace(PTRACE_SETREGS, tid, NULL, &regs) != 0){
        return -1;
    }
    return 0;
}

int GetRegister(NX_Solt tid, Registers *regs){
    if(NX_Ptrace(PTRACE_GETREGS, tid, NULL, regs) != 0){
        return -1;
    }
    return 0;
}


void ReadCmd(char *cmd, int fd, int length){
    memset(cmd, 0, length);
    char tmp = 0;
    char *p = cmd;
    int i = 0;
    while(tmp != '\n' && i < length){
        NX_FileRead(fd, &tmp, 1);
        if(tmp != 8){
            NX_FileWrite(fd, &tmp, 1);
            *p++ = tmp;
            i++;
        }else{
            //backspace
            if(p != cmd){
                p--;
                i--;
                NX_FileWrite(fd, &tmp, 1);
                NX_FileWrite(fd, " ", 1);
                NX_FileWrite(fd, &tmp, 1);
            }
        }
    }
}

int HandleContinue(NX_Solt solt){
    // enable all breakpoints
    int ret = 0;
    int i;
    int c_ebreak = gdb_c_ebreak_instruction;
    int ebreak = gdb_ebreak_instruction;
    for(i = 0; i < __GDB_MAX_BREAKPOINTS_NUM_; i++){
        if(gdb_bkp_buf[i].en == 1){
            if(gdb_bkp_buf[i].len == 2){
                if((ret = NX_Ptrace(PTRACE_POKEINT16, solt, (void*)gdb_bkp_buf[i].addr, &c_ebreak)) != 0){
                    break;
                }
            }else if(gdb_bkp_buf[i].len == 4){
                if((ret = NX_Ptrace(PTRACE_POKEINT32, solt, (void*)gdb_bkp_buf[i].addr, &ebreak)) != 0){
                    break;
                }
            }
        }
    }
    if(ret == 0){
        NX_Ptrace(PTRACE_CONT, solt, NULL, NULL);
        WaitForChild();
        for(i = 0; i < __GDB_MAX_BREAKPOINTS_NUM_; i++){
            if(gdb_bkp_buf[i].en == 1){
                if(gdb_bkp_buf[i].len == 2){
                    if((ret = NX_Ptrace(PTRACE_POKEINT16, solt, (void*)gdb_bkp_buf[i].addr, &gdb_bkp_buf[i].instruction)) != 0){
                        break;
                    }
                }else if(gdb_bkp_buf[i].len == 4){
                    if((ret = NX_Ptrace(PTRACE_POKEINT32, solt, (void*)gdb_bkp_buf[i].addr, &gdb_bkp_buf[i].instruction)) != 0){
                        break;
                    }
                }
            }
        }
        return 0;
    }
    NX_Printf("continue error!\n");
    return -1;
}

int DiskplayBreakPoint(){
    NX_Printf("Num\tType\tDisp\tEnb\tAddress\n");
    int num = 1;
    for(int i = 0; i < __GDB_MAX_BREAKPOINTS_NUM_; i++){
        if(gdb_bkp_buf[i].addr != 0){
            NX_Printf("%d\tbreakpoint\tkeep\t%s\t0x%p\n", num++, gdb_bkp_buf[i].en ? "y" : "n", gdb_bkp_buf[i].addr);
        }
    }
    return 0;
}

int HandleInfo(NX_Solt child, char*argv[8]){
    if(IsPrefix(argv[1], "registers") || IsPrefix(argv[1], "reg")){
        NX_Printf("info registers\n");
        Registers regs;
        GetRegister(child, &regs);
        NX_Printf("%s: 0x%p\n", g_register_descriptors[0].name, 0);
        for(int i = 1; i < GDB_CPU_NUM_REGISTERS -1; i++){
            NX_Printf("%s: 0x%p\n", g_register_descriptors[i].name, regs.reg[i]);
        }
        NX_Printf("%s: 0x%p\n", g_register_descriptors[GDB_CPU_NUM_REGISTERS -1].name, regs.reg[0]);
    }
    else if(IsPrefix(argv[1], "breakpoints") || (argv[1][0] == 'b' && argv[1][1] == 'p')){
        DiskplayBreakPoint();
    }
    else if(IsPrefix(argv[1], "threads")){
        NX_Printf("info threads\n");
    }
    else if(IsPrefix(argv[1], "frame")){
        NX_Printf("info frame\n");
    }
    else if(IsPrefix(argv[1], "locals")){
        NX_Printf("info locals\n");
    }
    else if(IsPrefix(argv[1], "args")){
        NX_Printf("info args\n");
    }
    else{
        return -1;
    }
    return 0;
}

int HandleSetRegister(NX_Solt child, char* argv[8]){
    char * command = argv[1];
    command++;
    int i = 0;
    for(;i<GDB_CPU_NUM_REGISTERS;i++){
        if(IsPrefix(command, g_register_descriptors[i].name)){
            break;
        }
    }
    if(i == GDB_CPU_NUM_REGISTERS){
        return -1;
    }

    command += 2;
    if(*command == 0){
        command = argv[2];
    }

    if(!IsPrefix(command, "=")){
        return -1;
    }

    command++;
    if(*command == 0){
        command = argv[3];
    }

    if(command == NULL){
        return -1;
    }

    char* endptr;
    unsigned long long value = strtoul(command, &endptr, 0);

    return SetRegisterValue(child, i, value);
    
}

int SetMemoryGetAddrData(char* argv[8], char *command, unsigned long long *addr, long long *value){
    int pargv = 2;
    if(*command == 0){
        command = argv[pargv++];
    }
    * addr = strtoul(command, &command, 0);
    if(addr == 0){
        return -1;
    }
    if(*command == 0){
        command = argv[pargv++];
    }
    if(!IsPrefix(command, "=")){
        return -1;
    }
    command++;
    if(*command == 0){
        command = argv[pargv++];
    }
    if(command == NULL){
        return -1;
    }
    *value = strtol(command, &command, 0);
    return 0;
}

int HandleSet(NX_Solt child, char* argv[8]){
    char* command;
    command = argv[1];
    if(IsPrefix(command, "$")){
        HandleSetRegister(child, argv);
    }else if(IsPrefix(command, "char")){
        unsigned long long addr;
        long long value;
        command += strlen("char");
        if(SetMemoryGetAddrData(argv, command, &addr, &value) == -1){
            return -1;
        }
        NX_Ptrace(PTRACE_POKEINT8, child, (void*)addr, (void*)&value);
    }else if (IsPrefix(command, "short")) {
        unsigned long long addr;
        long long value;
        command += strlen("short");
        if(SetMemoryGetAddrData(argv, command, &addr, &value) == -1){
            return -1;
        }
        NX_Ptrace(PTRACE_POKEINT16, child, (void*)addr, (void*)&value);
    }else if(IsPrefix(command, "int")){
        unsigned long long addr;
        long long value;
        command += strlen("int");
        if(SetMemoryGetAddrData(argv, command, &addr, &value) == -1){
            return -1;
        }
        NX_Ptrace(PTRACE_POKEINT32, child, (void*)addr, (void*)&value);
    }else if(IsPrefix(command, "long")){
        unsigned long long addr;
        long long value;
        command += strlen("long");
        if(SetMemoryGetAddrData(argv, command, &addr, &value) == -1){
            return -1;
        }
        NX_Ptrace(PTRACE_POKEINT64, child, (void*)addr, (void*)&value);
    }else if(IsPrefix(command, "float")){
        //todo
        return -1;
    }else if (IsPrefix(command, "double")) {
        //todo
        return -1;
    }
    else{
        return -1;
    }
    return 0;
}

int GetData(NX_Solt child, long long addr, char* buf, int size, int stepSize){
    memset(buf, 0, size * stepSize);
    int i = 0;
    int ret = 0;
    if(stepSize == 1){
        for(i = 0; i < size; ++i){
            ret = NX_Ptrace(PTRACE_PEEKINT8, child, (void*)(addr + i), buf + i);
            if(ret == -1){
                NX_Printf("PTRACE PEEKDATA failed\n");
                return -1;
            }
        }
    }else if(stepSize == 2){
        for(i = 0; i < size; ++i){
            ret = NX_Ptrace(PTRACE_PEEKINT16, child, (void*)(addr + i * stepSize), buf + i * stepSize);
            if(ret == -1){
                NX_Printf("PTRACE PEEKDATA failed\n");
                return -1;
            }
        }
    }else if(stepSize == 4){
        for(i = 0; i < size; ++i){
            ret = NX_Ptrace(PTRACE_PEEKINT32, child, (void*)(addr + i * stepSize), buf + i * stepSize);
            if(ret == -1){
                NX_Printf("PTRACE PEEKDATA failed\n");
                return -1;
            }
        }
    }else if(stepSize == 8){
        for(i = 0; i < size; ++i){
            ret = NX_Ptrace(PTRACE_PEEKINT64, child, (void*)(addr + i * stepSize), buf + i * stepSize);
            if(ret == -1){
                NX_Printf("PTRACE PEEKDATA failed\n");
                return -1;
            }
        }
    }else{
        NX_Printf("do not support\n");
        return -1;
    }
    
    return 0;
}



int PrintFormatX(char * buf, int size, int step, long long addr){
    if(step == 1){
        NX_Printf("0x%p: ", addr);
        for(int i = 0; i < size; i++){
            NX_Printf("%02p ", buf[i]);
            if(i % 16 == 15){
                NX_Printf("\n0x%p: ", addr + i + 1);
            }
        }
    }else if(step == 2){
        short* dest = (short*)buf;
        NX_Printf("0x%p: ", addr);
        for(int i = 0; i < size; i++){
            NX_Printf("%04p ", dest[i]);
            if(i % 8 == 7){
                NX_Printf("\n0x%p: ", addr + (i + 1) * 2);
            }
        }
    }else if(step == 4){
        int* dest = (int*)buf;
        NX_Printf("0x%p: ", addr);
        for(int i = 0; i < size; i++){
            NX_Printf("%08p ", dest[i]);
            if(i % 4 == 3){
                NX_Printf("\n0x%p: ", addr + (i + 1) * 4);
            }
        }
    }else if(step == 8){
        long long* dest = (long long*)buf;
        NX_Printf("0x%p: ", addr);
        for(int i = 0; i < size; i++){
            NX_Printf("%016p ", dest[i]);
            if(i % 2 == 1){
                NX_Printf("\n0x%p: ", addr + (i + 1) * 8);
            }
        }
    }else{
        return -1;
    }
    return 0;
}

int PrintXData(char f, int stepSize, char* buf, int size, long long addr){
    if(f == 'x'){
        PrintFormatX(buf, size, stepSize, addr);
    }else{
        NX_Printf("do not support %c.\n", f);
        return -1;
    }
    return 0;
}

int HandleX(NX_Solt child, char* argv[8]){
    char* command;
    command = argv[1];
    if(IsPrefix(command, "/") == 0){
        NX_Printf("Unkown command: %s, expected /.\n", command);
        return -1;
    }
    command++;

    int num = strtol(command, &command, 0);
    if(num == 0){
        NX_Printf("Unkown command: %s, expected num.\n", command);
        return -1;
    }

    if(*command == '0'){
        return -1;
    }
    char f = *command;
    command++;

    if(*command == '0'){
        return -1;
    }
    char u = *command;

    long long addr = strtol(argv[2], NULL, 0);
    if(addr == 0){
        return 0;
    }

    char * buf;
    int stepSize;
    if(u == 'b'){
        buf = (char*)malloc(num);
        stepSize = 1;
    }else if(u == 'h'){
        buf = (char*)malloc(num * 2);
        stepSize = 2;
    }else if(u == 'w'){
        buf = (char*)malloc(num * 4);
        stepSize = 4;
    }else if(u == 'g'){
        buf = (char*)malloc(num * 8);
        stepSize = 8;
    }else{
        NX_Printf("do not support %c.\n", u);
        return -1;
    }
    if(buf == NULL){
        NX_Printf("Malloc failed.\n");
        return -1;
    }
    GetData(child, addr, buf, num, stepSize);
    PrintXData(f, stepSize, buf, num, addr);
    free(buf);
    return 0;
}

int AddBreakPoint(NX_U64 addr, NX_U64 len, NX_U64 type, NX_U32 instruction){
    for(int i = 0; i < __GDB_MAX_BREAKPOINTS_NUM_; i++){
        if(gdb_bkp_buf[i].en == 0 && gdb_bkp_buf[i].addr == 0){
            gdb_bkp_buf[i].en = 1;
            gdb_bkp_buf[i].addr = addr;
            gdb_bkp_buf[i].len = len;
            gdb_bkp_buf[i].type = type;
            gdb_bkp_buf[i].instruction = instruction;
            return 0;
        }
    }
    return -1;
}

int HandleBreakPoint(NX_Solt child, char* argv[8]){
    char* command = argv[1];
    unsigned long long addr = 0;
    if(command == NULL){
        return -1;
    }
    addr = strtoul(command, &command, 0);
    if(addr == 0 || addr % 2){
        return -1;
    }
    int instruction;
    if(NX_Ptrace(PTRACE_PEEKINT32, child, (void*)addr, &instruction) == -1){
        return -1;
    }
    int length = 0;
    if((instruction & 0x3) != 0x3){
        length = 2;
        instruction = (instruction & 0x0000ffff) ;
    }else{
        length = 4;
    }
    AddBreakPoint(addr, length, insert_break_or_watchpoint_software_breakpoint, instruction);
    return 0;
}

int DeleteBreakPoint(NX_Solt child, char* argv[8]){
    int num = strtol(argv[1], NULL, 0);
    if(num == 0){
        return -1;
    }
    int i = 0;
    for(; i < __GDB_MAX_BREAKPOINTS_NUM_; i++){
        if(gdb_bkp_buf[i].addr != 0){
            num--;
            if(num == 0){
                gdb_bkp_buf[i].en = 0;
                gdb_bkp_buf[i].addr = 0;
                gdb_bkp_buf[i].len = 0;
                gdb_bkp_buf[i].type = 0;
                gdb_bkp_buf[i].instruction = 0;
                return 0;
            }
        }
    }
    return 0;
}

int HandleStepi(NX_Solt child){
    int ret = NX_Ptrace(PTRACE_SINGLESTEP, child, NULL, NULL);
    if(ret == -1){
        NX_Printf("PTRACE_SINGLESTEP failed.\n");
        return -1;
    }else if(ret == 1){
        WaitForChild();    
    }
    return 0;
}

void DbgHandleCommand(char *cmd, NX_Solt child){
    char *lcmd = strdup(cmd);
    char *argv[command_token_length] = { 0 };    
    char *command;

    argv[0] = strtok(lcmd, " ");
    for(int i = 1; i < command_token_length; i++){
        argv[i] = strtok(NULL, " ");
        if(argv[i] == NULL) break;
    }
    command = argv[0];
    if(command == NULL) return;
    if(IsPrefix(command, "continue")){
        if(HandleContinue(child) == -1){
            NX_Printf("Continue failed.\n");
        }else{
            //todo
        }
    }
    else if(IsPrefix(command, "info")){
        if(HandleInfo(child, argv) == -1){
            goto unknown;
        }
    }
    else if(IsPrefix(command, "set")){
        if(HandleSet(child, argv) == -1){
            goto unknown;
        }
    }
    else if(IsPrefix(command, "x")){
        if(HandleX(child, argv) == -1){
            goto unknown;
        }
    }
    else if(IsPrefix(command, "break") || (strlen(command) == 1 && *command == 'b')){
        if(HandleBreakPoint(child, argv) == -1){
            goto unknown;
        }
    }
    else if(IsPrefix(command, "delete")){
        if(DeleteBreakPoint(child, argv) == -1){
            goto unknown;
        }
    }
    else if(IsPrefix(command, "stepi")){
        HandleStepi(child);
    }
    else{
        goto unknown;
    }
    free(lcmd);
    return ;
unknown:
    NX_Printf("Unkown command: %s.\n", cmd);
    free(lcmd);
}


void DbgRun(int fd, NX_Solt child){
    char cmd[100];
    
    /*UI for start up*/
    while (1) {
        NX_Printf("\ngdb$ :");
        ReadCmd(cmd, fd, 100);
        DbgHandleCommand(cmd, child);
    }
}

int IsPrefix(const char *s, const char *ss){
    if(s == NULL || ss == NULL) return 0;
    if(strlen(ss) > strlen(s)) return 0;

    return !strncmp(s, ss, strlen(ss));
}

void SignalHandler(NX_SignalInfo * info)
{
    NX_Printf("child process stopped\n");
    isChildStop = 1;
}



// NX_Error NX_Main(char *cmdline, char *envline)
int main(int argc, char *argv[])
{
    char *name = strdup(argv[1]);
    if(name == NULL)
    {
        NX_Printf("Usage: gdb <program>\n");
        return -NX_EINVAL;
    }

    NX_SignalAttr attr;
    NX_Error err;
    attr.handler = SignalHandler;
    err = NX_SignalSetAttr(NX_SIGNAL_CHILD, &attr);
    NX_Printf("signal set attr:%s\n", NX_ErrorToString(err));

    NX_U32 exitCode;
    NX_Solt child = NX_ProcessLaunch(name, NX_THREAD_CREATE_DEBUG, &exitCode, NULL, NULL);
    if(child == NX_SOLT_INVALID_VALUE){
        NX_Printf("Launch %s failed.\n", name);
        return -NX_EINVAL;
    }

    WaitForChild();

    int fd = -1;
    fd = NX_FileOpen("/System/Device/console", NX_FILE_RDWR, 0);
    if(fd == -1){
        NX_Printf("Open console failed.\n");
        return -NX_EINVAL;
    }
    
    memset(gdb_bkp_buf, 0, sizeof(gdb_bkp_buf));
    memset(gdb_watch_buf, 0, sizeof(gdb_watch_buf));
    DbgRun(fd, child);
    return NX_EOK;
}