/**
 * MIT License
 *
 * Copyright (c) 2024 Johannes Stölp
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 **/

#define _GNU_SOURCE
#include <dlfcn.h>
#include <signal.h>

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <unwind.h>  // from gcc, clang

#define LOG(fmt, ...)                                              \
  do {                                                             \
    char buf[1024];                                                \
    int cnt = snprintf(buf, sizeof(buf), fmt "\n", ##__VA_ARGS__); \
    write(STDERR_FILENO, buf, cnt);                                \
  } while (0)

#define INFO(fmt, ...) LOG("[\e[0;32mINFO:BT\e[0m]: " fmt, ##__VA_ARGS__)
#define WARN(fmt, ...) LOG("[\e[0;33mWARN:BT\e[0m]: " fmt, ##__VA_ARGS__)
#define FAIL(fmt, ...) LOG("[\e[0;31mFAIL:BT\e[0m]: " fmt, ##__VA_ARGS__)

// -- UNWIND HANDLER -----------------------------------------------------------

struct unwind_data {
  int num_frames;
  int frame;
  int skip;
};

static _Unwind_Reason_Code libbt_handle_unwind_frame(
    struct _Unwind_Context* uctx,
    void* arg) {
  struct unwind_data* data = (struct unwind_data*)arg;
  void* pc                 = (void*)_Unwind_GetIP(uctx);

  if (pc == NULL) {
    return _URC_END_OF_STACK;
  }

  if (data->skip > 0) {
    data->skip -= 1;
    return _URC_NO_REASON;
  }

  Dl_info info;
  memset(&info, 0, sizeof(info));
  if (dladdr(pc, &info) == 0) {
    FAIL("dladdr call failed pc=%p stop unwinding", pc);
    return _URC_END_OF_STACK;
  }

  if (info.dli_saddr) {
    LOG("#%2d [%16p] %s+0x%lx %s+0x%lx", data->frame, pc, info.dli_sname,
        pc - info.dli_saddr, info.dli_fname, info.dli_saddr - info.dli_fbase);
  } else {
    LOG("#%2d [%16p] %s+0x%lx", data->frame, pc, info.dli_fname,
        pc - info.dli_fbase);
  }

  if (data->frame++ < data->num_frames) {
    return _URC_NO_REASON;
  }
  return _URC_END_OF_STACK;
}

// -- SIGNAL HANDLING ----------------------------------------------------------

static struct sigaction OLD;
static struct sigaction NEW;

static void libbt_report(int sig, siginfo_t* info, void* ucontext) {
  INFO("caught signal %d", sig);

  struct unwind_data ctx = {
      .num_frames = 10,
      .frame      = 0,
      .skip       = 2,
  };
  _Unwind_Backtrace(libbt_handle_unwind_frame, &ctx);

  if (OLD.sa_flags & SA_SIGINFO && OLD.sa_sigaction) {
    OLD.sa_sigaction(sig, info, ucontext);
  } else if (OLD.sa_handler) {
    OLD.sa_handler(sig);
  }

  const char* noexit = getenv("BT_NOEXIT");
  if (noexit && noexit[0] == '1') {
    return;
  }
  INFO("exit after catching signal, set BT_NOEXIT=1 to keep running");
  exit(42);
}

static void* libbt_install_sigstack() {
  const size_t STACK_SIZE = 8192;

  void* sp = malloc(STACK_SIZE);
  if (sp == NULL) {
    WARN("failed to allocate alternative signal stack");
    return NULL;
  }

  stack_t stack = {
      .ss_sp    = sp,
      .ss_flags = 0,
      .ss_size  = STACK_SIZE,
  };
  if (sigaltstack(&stack, NULL) != 0) {
    WARN("failed to setup alternative signal stack");
  }
  return sp;
}

static int libbt_install_sighandler() {
  memset(&OLD, 0, sizeof(OLD));
  memset(&NEW, 0, sizeof(NEW));

  NEW.sa_flags     = SA_SIGINFO | SA_ONSTACK;
  NEW.sa_sigaction = libbt_report;

  return sigaction(SIGSEGV, &NEW, &OLD);
}

// -- THREAD HANDLING ----------------------------------------------------------

int (*NEXT_PTHREAD_CREATE)(pthread_t*,
                           const pthread_attr_t*,
                           void* (*)(void*),
                           void*);

static int THREAD_CREATE_SIGSTACK = 0;

struct thread_arg {
  void* (*start_routine)(void*);
  void* restrict arg;
};

static void* thread_entry(void* th_arg) {
  struct thread_arg arg = *(struct thread_arg*)th_arg;
  free(th_arg);

  void* sp  = THREAD_CREATE_SIGSTACK ? libbt_install_sigstack() : NULL;
  void* ret = arg.start_routine(arg.arg);
  free(sp);
  return ret;
}

int pthread_create(pthread_t* thread,
                   const pthread_attr_t* attr,
                   void* (*start_routine)(void*),
                   void* arg) {
  assert(NEXT_PTHREAD_CREATE);

  struct thread_arg* th_arg = malloc(sizeof(struct thread_arg));
  th_arg->start_routine     = start_routine;
  th_arg->arg               = arg;
  return NEXT_PTHREAD_CREATE(thread, attr, thread_entry, th_arg);
}

// -- GLOBAL INITIALIZATION ----------------------------------------------------

static void __attribute__((constructor)) init() {
  NEXT_PTHREAD_CREATE = dlsym(RTLD_NEXT, "pthread_create");
  if (NEXT_PTHREAD_CREATE == NULL) {
    FAIL("failed to get pthread_create");
    exit(1);
  }

  if (libbt_install_sighandler() != 0) {
    FAIL("failed to install signal handler");
    return;
  }
  INFO("installed signal handler");

  // Install signal stack for main thread, and leak it leak.
  libbt_install_sigstack();

  // Only create signal stacks in threads if the signal handler was installed.
  THREAD_CREATE_SIGSTACK = 1;
}