/**
 * 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 <link.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 {                                        \
    fprintf(stderr, fmt "\n", ##__VA_ARGS__); \
  } while (0)

#define INFO(fmt, ...) LOG("[BT:\e[0;32mINFO\e[0m]: " fmt, ##__VA_ARGS__)
#define WARN(fmt, ...) LOG("[BT:\e[0;33mWARN\e[0m]: " fmt, ##__VA_ARGS__)
#define FAIL(fmt, ...) LOG("[BT:\e[0;31mFAIL\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("LIBBT_NOEXIT");
  if (noexit && noexit[0] == '1') {
    return;
  }
  INFO("exit after catching signal, set LIBBT_NOEXIT=1 to keep running");
  exit(42);
}

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

  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 ----------------------------------------------------------

static 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;
}

static void resolve_next_pthread_create() {
  if (NEXT_PTHREAD_CREATE) {
    return;
  }

  // The following implements dlsym(RTLD_NEXT, "pthread_create") manually.
  //
  // In the past I have observed a case where dlsym(RTLD_NEXT) returned an
  // address from a library earlier in the link map. This is problematic as
  // calling the received function pointer ends in an infinite recursion.
  //
  // This may be a bug or I did not fully understand the case yet as I have
  // only observed it in special occasions where there are other LD_PRELOAD
  // libraries implementing the same wrapper function and invoking the wrapper
  // before the static initializers of this library have been executed.
  //
  // So far this implementation yielded "better" results.

  void* ret_addr = __builtin_return_address(0);
  Dl_info _info;
  memset(&_info, 0, sizeof(_info));
  struct link_map* lmap;

  if (dladdr1(ret_addr, &_info, (void**)&lmap, RTLD_DL_LINKMAP) == 0) {
    FAIL("dladdr1 failed while resolving next pthread_create");
    exit(1);
  }

  for (lmap = lmap->l_next; lmap && !NEXT_PTHREAD_CREATE; lmap = lmap->l_next) {
    if (!lmap->l_name) {
      continue;
    }

    void* dso = dlopen(lmap->l_name, RTLD_LOCAL | RTLD_LAZY);
    if (!dso) {
      WARN("failed to open %s, continue with next", lmap->l_name);
      continue;
    }
    NEXT_PTHREAD_CREATE = dlsym(dso, "pthread_create");
    dlclose(dso);
  }

  if (!NEXT_PTHREAD_CREATE) {
    FAIL("failed to resolve next pthread_create function");
    exit(1);
  }
}

int pthread_create(pthread_t* thread,
                   const pthread_attr_t* attr,
                   void* (*start_routine)(void*),
                   void* arg) {
  resolve_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() {
  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;
}