#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <signal.h>
#include <syscall.h>

#include <unistd.h>
#include <errno.h>
#include <assert.h>
#include <fcntl.h>

#include <sys/uio.h>
#include <sys/types.h>
#include <sys/user.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <sys/personality.h>

typedef uint8_t  u8;
typedef uint16_t u16;
typedef uint32_t u32;
typedef uint64_t u64;

typedef int8_t  s8;
typedef int16_t s16;
typedef int32_t s32;
typedef int64_t s64;

#define NT_PRSTATUS 1
#define DIE(string) do { fprintf(stderr, string); exit(1); } while (0)

#include "structs.h"

struct mi_process {
    pid_t pid;
    
    char *source;
    u64 source_size;
    
    u8 *elf;
    u64 elf_size;
    
    u64 base_address;
    u64 main_address;
};

struct mi_registers {
    u64 rbp;
    u64 rbx;
    u64 rax;
    u64 rcx;
    u64 rdx;
    u64 rsi;
    u64 rdi;
    u64 rip;
    u64 rsp;
    
    struct user_regs_struct _sys;
};

#include "dwarf.c"


static void
print_word_at_address(int child, u64 address)
{
    long word = ptrace(PTRACE_PEEKDATA, child, address, NULL);
    u8 nb[4];
    memcpy(nb, &word, 4);
    printf("word at %#018lx: {%x %x %x %x}\n", address, nb[0], nb[1], nb[2], nb[3]);
}

static void
print_current_line(struct mi_process proc, u64 line)
{
    char *source = proc.source;
    u64 size = proc.source_size;
    u64 current_line = 1;
    
    for (u32 i = 0; i < size; ++i) {
        if (current_line == line) {
            int len = 0;
            int start = 0;
            int leading = 1;
            
            for (u32 j = i; j < size; ++j) {
                if (leading && (source[j] == '\t' || source[j] == ' ')) {
                    ++start;
                } else {
                    leading = 0;
                }
                
                if (source[j] == '\n') {
                    break;
                }
                
                ++len;
            }
            
            
            printf("%.*s\n", len - start, source + i + start);
            return;
        }
        
        if (source[i] == '\n') {
            current_line++;
        }
    }
}

static struct mi_process
process_create(char *path)
{
    struct mi_process result = { 0 };
    
    /* disable ASLR */
    syscall(SYS_personality, ADDR_NO_RANDOMIZE);
    
    pid_t pid = fork();
    
    if (pid == -1) {
        DIE("fork error\n");
    }
    
    if (pid == 0) {
        int rt = ptrace(PTRACE_TRACEME, 0, 0, 0);
        assert(rt == 0);
        char *args[] = { path, NULL };
        rt = execvp(path, args);
        if (rt == -1) {
            DIE("exevp error\n");
        }
    }
    
    u64 source_size;
    u64 elf_size;
    u8 *elf;
    char *source;
    struct stat s;
    int fd;
    
    {
        fd = open(path, O_RDONLY);
        assert(fd != -1);
        fstat(fd, &s);
        elf_size = s.st_size;
        elf = mmap(0, s.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
        close(fd);
    }
    
    {
        fd = open("/code/trace-test/traceme.c", O_RDONLY); // TODO
        assert(fd != -1);
        fstat(fd, &s);
        source_size = s.st_size;
        source = mmap(0, source_size, PROT_READ, MAP_PRIVATE, fd, 0);
        close(fd);
    }
    
    result.elf = elf;
    result.elf_size = elf_size;
    result.source = source;
    result.source_size = source_size;
    result.pid = pid;
    
    return(result);
}

static struct mi_registers
get_process_registers(struct mi_process proc)
{
    struct user_regs_struct regs = { 0 };
    struct iovec iov = { &regs, sizeof(regs) };
    struct mi_registers result = { 0 };
    
    ptrace(PTRACE_GETREGSET, proc.pid, NT_PRSTATUS, &iov);
    
    result.rax = regs.rax;
    result.rip = regs.rip;
    result.rbp = regs.rbp;
    result.rbx = regs.rbx;
    result.rcx = regs.rcx;
    result.rdx = regs.rdx;
    result.rsi = regs.rsi;
    result.rdi = regs.rdi;
    result.rsp = regs.rsp;
    
    result._sys = regs;
    
    return(result);
}

static void
set_process_registers(struct mi_process proc, struct mi_registers regs)
{
    struct iovec iov = { &regs._sys, sizeof(regs._sys) };
    ptrace(PTRACE_GETREGSET, proc.pid, NT_PRSTATUS, &iov);
}

static void
print_current_instruction(struct mi_process proc)
{
    struct mi_registers regs = get_process_registers(proc);
    
    long wordb = ptrace(PTRACE_PEEKDATA, proc.pid, regs.rip - 4, NULL);
    long word = ptrace(PTRACE_PEEKDATA, proc.pid, regs.rip, NULL);
    long worda = ptrace(PTRACE_PEEKDATA, proc.pid, regs.rip + 4, NULL);
    
    u8 nb[12];
    memcpy(nb, &wordb, 4);
    memcpy(nb + 4, &word, 4);
    memcpy(nb + 8, &worda, 4);
    
    printf("PC = %#018lx: %x %x %x %x [%x] %x %x %x %x %x %x %x\n", regs.rip, 
           nb[0], nb[1], nb[2], nb[3], nb[4], nb[5], nb[6], nb[7],
           nb[8], nb[9], nb[10], nb[11]);
}


static void
command_regs(struct mi_process proc)
{
    struct mi_registers regs = get_process_registers(proc);
    printf("rax = %#018lx rcx = %#018lx\n"
           "rbp = %#018lx rbx = %#018lx\n"
           "rdx = %#018lx rsi = %#018lx\n"
           "rdi = %#018lx rip = %#018lx\n"
           "rsp = %#018lx\n", 
           regs.rax, regs.rcx, regs.rbp, regs.rbx, 
           regs.rdx, regs.rsi, regs.rdi, regs.rip, 
           regs.rsp);
}

static void
command_step(struct mi_process proc)
{
    ptrace(PTRACE_SINGLESTEP, proc.pid, 0, 0);
    waitpid(proc.pid, 0, 0);
    
#if 1
    print_current_instruction(proc);
    struct mi_registers regs = get_process_registers(proc);
    u64 line_number = pc_to_line_number(proc.elf, regs.rip - proc.base_address);
    print_current_line(proc, line_number);
#endif
}

static void
command_start(struct mi_process proc)
{
    u64 main_address = proc.base_address + proc.main_address;
    
    long saved_instruction = ptrace(PTRACE_PEEKDATA, proc.pid, main_address, NULL);
    long int3_byte = 0x000000cc;
    long int3_word = (saved_instruction & 0xFFFFFF00) | int3_byte;
    
    /* write 0xcc */
    ptrace(PTRACE_POKEDATA, proc.pid, main_address, int3_word);
    
    /* wait till child hits the interrupt */
    ptrace(PTRACE_CONT, proc.pid, 0, 0);
    waitpid(proc.pid, 0, 0);
    
    /* restore original instrucion */
    ptrace(PTRACE_POKEDATA, proc.pid, main_address, saved_instruction);
    
    /* step back to original instruction */
    struct mi_registers regs = get_process_registers(proc);
    regs.rip -= 1;
    set_process_registers(proc, regs);
}

static void
command_next(struct mi_process proc)
{
    struct mi_registers regs = get_process_registers(proc);
    
    u64 line_number = pc_to_line_number(proc.elf, regs.rip - proc.base_address);
    u64 next_line_number;
    
    do {
        ptrace(PTRACE_SINGLESTEP, proc.pid, 0, 0);
        waitpid(proc.pid, 0, 0);
        regs = get_process_registers(proc);
        //printf("%#lx\n", regs.rip - proc.base_address);
        next_line_number= pc_to_line_number(proc.elf, regs.rip - proc.base_address);
        // TODO: skip call, repXXX, etc
    } while (next_line_number == line_number);
    
    print_current_line(proc, next_line_number);
}

static void
command_list(struct mi_process proc)
{
    struct mi_registers regs = get_process_registers(proc);
    u64 line_number = pc_to_line_number(proc.elf, regs.rip - proc.base_address);
    print_current_line(proc, line_number);
}

static void
command_cont(struct mi_process proc)
{
    ptrace(PTRACE_CONT, proc.pid, 0, 0);
    // TODO
}

static void
command_stop(struct mi_process proc)
{
    kill(proc.pid, SIGINT);
    // TODO
}

int
main(int argc, char *argv[])
{
    if (argc != 2) {
        printf("Usage: %s executable\n", argv[0]);
        return(1);
    }
    
    struct mi_process process = process_create(argv[1]);
    printf("pid of child = %d\n", process.pid);
    
    size_t max_command_length = 1023;
    int command_length = 0;
    char *command = malloc(max_command_length + 1);
    
    process.base_address = 0x555555554000UL; // get_executable_base_address(file, proc.pid);
    process.main_address = get_address_of_subroutine(process.elf, "main");
    
    printf("Base address: %#lx\n", process.base_address);
    printf("Main address: %#lx\n", process.main_address);
    
    printf("> ");
    fflush(stdout);
    
    while ((command_length = getline(&command, &max_command_length, stdin))) {
        if (0 == strncmp(command, "exit\n", command_length) || 0 == strncmp(command, "q\n", command_length)) {
            break;
        } else if (strncmp(command, "regs\n", command_length) == 0) {
            command_regs(process);
        } else if (strncmp(command, "step\n", command_length) == 0) {
            command_step(process);
        } else if (strncmp(command, "start\n", command_length) == 0) {
            command_start(process);
        } else if (strncmp(command, "next\n", command_length) == 0) {
            command_next(process);
        } else if (strncmp(command, "line\n", command_length) == 0) {
            command_list(process);
        } else if (strncmp(command, "cont\n", command_length) == 0) {
            command_cont(process);
        } else if (strncmp(command, "stop\n", command_length) == 0) {
            command_stop(process);
        } else {
            printf("Unknown command: %*s", command_length, command);
        }
        
        int wstatus;
        int rt = waitpid(process.pid, &wstatus, WNOHANG);
        
        if (rt > 0 && WIFEXITED(wstatus)) {
            printf("Child exited\n");
            break;
        }
        
        printf("> ");
    }
    
    return(0);
}