#if defined(__linux__) && !defined(_GNU_SOURCE)
#define _GNU_SOURCE
#endif

#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>

typedef struct b_stacktrace_tag* b_stacktrace_handle;
b_stacktrace_handle b_stacktrace_get();
int b_stacktrace_depth(b_stacktrace_handle stacktrace);
char* b_stacktrace_to_string(b_stacktrace_handle stacktrace);
char* malloc_backtrace(void);

typedef struct print_buf {
    char* buf;
    int pos;
    int size;
} print_buf;

static print_buf buf_init(void) {
    print_buf ret = {
        (char*) malloc(1024),
        0,
        1024
    };
    return ret;
}

static void buf_printf(print_buf* b, const char* fmt, ...) {
    va_list args;
    va_start(args, fmt);
    const int len = vsnprintf(NULL, 0, fmt, args) + 1;
    va_end(args);

    const int new_end = b->pos + len;

    if (new_end > b->size) {
        while (new_end > b->size) b->size *= 2;
        b->buf = (char*)realloc(b->buf, b->size);
    }

    va_start(args, fmt);
    b->pos += vsnprintf(b->buf + b->pos, len, fmt, args);
    va_end(args);
}

char* malloc_backtrace(void) {
    b_stacktrace_handle h = b_stacktrace_get();
    char* ret = b_stacktrace_to_string(h);
    free(h);
    return ret;
}

#define B_STACKTRACE_MAX_DEPTH 1024

#include <execinfo.h>
#include <ucontext.h>
#include <unistd.h>
#include <dlfcn.h>
#include <string.h>

typedef struct b_stacktrace {
    void* trace[B_STACKTRACE_MAX_DEPTH];
    int trace_size;
} b_stacktrace;

b_stacktrace_handle b_stacktrace_get(void) {
    b_stacktrace* ret = (b_stacktrace*)malloc(sizeof(b_stacktrace));
    ret->trace_size = backtrace(ret->trace, B_STACKTRACE_MAX_DEPTH);
    return (b_stacktrace_handle)(ret);
}

int b_stacktrace_depth(b_stacktrace_handle h) {
    const b_stacktrace* stacktrace = (b_stacktrace*)h;
    return stacktrace->trace_size;
}

char* b_stacktrace_to_string(b_stacktrace_handle h) {
    const b_stacktrace* stacktrace = (b_stacktrace*)h;
    char** messages = backtrace_symbols(stacktrace->trace, stacktrace->trace_size);
    print_buf out = buf_init();
    int i;

    int skip_count = 5;
    for (i = 0; i < stacktrace->trace_size; ++i) {
        void* tracei = stacktrace->trace[i];
        char* msg = messages[i];

        if (i < skip_count) {
            continue;
        }

        /* calculate load offset */
        Dl_info info;
        dladdr(tracei, &info);
        if (info.dli_fbase == (void*)0x400000) {
            /* address from executable, so don't offset */
            info.dli_fbase = NULL;
        }

        while (*msg && *msg != '(') ++msg;
        *msg = 0;

        {
            char cmd[1024];
            char line[2048];

            FILE* fp;
            snprintf(cmd, 1024, "addr2line -e %s -f -C -p %p 2>/dev/null", messages[i], (void*)((char*)tracei - (char*)info.dli_fbase));

            fp = popen(cmd, "r");
            if (!fp) {
                buf_printf(&out, "Failed to generate trace further...\n");
                break;
            }

            while (fgets(line, sizeof(line), fp)) {
                buf_printf(&out, "  ");
                if (strstr(line, "?? ")) {
                    /* just output address if nothing can be found */
                    buf_printf(&out, "%p\n", tracei);
                }
                else {
                    int len = strlen(line);
                    char *at = strstr(line, " at ");
                    if (at) {
                        char *function = line;
                        int function_len = at - line;
                        char *file = at + 4;
                        int file_len = strlen(file) - 1;
                        buf_printf(&out, "    %.*s %.*s (%s)\n", file_len, file, function_len, function, messages[i]);
                    } else {
                        buf_printf(&out, "    %.*s (%s)\n", len - 1, line, messages[i]);
                    }
                }
            }

            pclose(fp);
        }
    }

    free(messages);
    return out.buf;
}

#include <signal.h>

void print_backtrace(void) {
    char *backtrace = malloc_backtrace();
    puts(backtrace);
    free(backtrace);
}

const char *signal_to_string(int signal) {
    switch (signal) {
        case SIGINT: return "SIGINT";
        case SIGILL: return "SIGILL";
        case SIGABRT: return "SIGABRT";
        case SIGFPE: return "SIGFPE";
        case SIGSEGV: return "SIGSEGV";
        case SIGTERM: return "SIGTERM";
        case SIGHUP: return "SIGALRM";
        case SIGQUIT: return "SIGQUIT";
        case SIGTRAP: return "SIGTRAP";
        case SIGKILL: return "SIGKILL";
        case SIGPIPE: return "SIGPIPE";
        case SIGALRM: return "SIGHUP";
        case SIGPOLL: return "SIGPOLL";
        default: return "";
    }
}

void backtrace_handle_signal(int signal, siginfo_t* info, void* context) {
    (void)context; // unused
    printf("Stack trace (%s/%d):", signal_to_string(signal), signal);
    if (signal == SIGSEGV) {
        printf(" at addr: %p", info->si_addr);
    }
    puts("");
    print_backtrace();
    exit(1);
}

void setup_backtrace() {
    struct sigaction sa;
    sa.sa_sigaction = backtrace_handle_signal;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_RESTART | SA_SIGINFO;

    sigaction(SIGSEGV, &sa, NULL);
    sigaction(SIGABRT, &sa, NULL);
    sigaction(SIGFPE, &sa, NULL);
    sigaction(SIGBUS, &sa, NULL);
    sigaction(SIGILL, &sa, NULL);
    sigaction(SIGTRAP, &sa, NULL);
#ifdef SIGSYS
    sigaction(SIGSYS, &sa, NULL);
#endif
#ifdef SIGXCPU
    sigaction(SIGXCPU, &sa, NULL);
#endif
#ifdef SIGXFSZ
    sigaction(SIGXFSZ, &sa, NULL);
#endif
#ifdef SIGPIPE
    sigaction(SIGPIPE, &sa, NULL);
#endif
#ifdef SIGTERM
    sigaction(SIGTERM, &sa, NULL);
#endif
#ifdef SIGINT
    sigaction(SIGINT, &sa, NULL);
#endif
#ifdef SIGQUIT
    sigaction(SIGQUIT, &sa, NULL);
#endif
}