#import "LPStackTraceGenerator.h"

#include <execinfo.h>
#include <mach/mach.h>

#if defined(__x86_64__) || defined(__arm64__)

// To get the FP_LINK_OFFSET and ISALIGNED for the new architecture,
// please check the lastest version of
// https://opensource.apple.com/source/Libc/Libc-1272.250.1/gen/thread_stack_pcs.c.auto.html
#if defined(__x86_64__) || defined(__arm64__)
#define FP_LINK_OFFSET 1
#endif

#if defined(__x86_64__)
#define ISALIGNED(a) ((((uintptr_t)(a)) & 0xf) == 0)
#elif defined(__arm64__)
#define ISALIGNED(a) ((((uintptr_t)(a)) & 0x1) == 0)
#endif

#define ISVALID(a) \
  ((a) >= stackbot && (a) <= stacktop && frame != NULL && (uintptr_t)a >= 4096 && ISALIGNED(a))

static _STRUCT_MCONTEXT GetContext(pthread_t pthread);
static void *GetFP(_STRUCT_MCONTEXT ctx);
static void *GetPC(_STRUCT_MCONTEXT ctx);
static bool SafeReadMemory(vm_address_t src, void *dest, size_t len);
/**
 * Stack unwind by using frame pointer.
 *
 * @param buffer Output. Array of addresses.
 * @param max Size of buffer.
 * @param *frames Output. Number of acutal frames.
 * @param thread Target Thread.
 * @param frame Frame pointer of target thread.
 */
static void TheadStackUnwind(vm_address_t *buffer, unsigned max, unsigned *frames, pthread_t thread,
                             void *frame);
static NSString *getThreadName(pthread_t thread);

static inline _STRUCT_MCONTEXT GetContext(pthread_t pthread) {
  thread_t thread = pthread_mach_thread_np(pthread);
  _STRUCT_MCONTEXT ctx;
#if defined(__x86_64__)
  mach_msg_type_number_t count = x86_THREAD_STATE64_COUNT;
  thread_get_state(thread, x86_THREAD_STATE64, (thread_state_t)&ctx.__ss, &count);
#elif defined(__arm64__)
  mach_msg_type_number_t count = ARM_THREAD_STATE64_COUNT;
  thread_get_state(thread, ARM_THREAD_STATE64, (thread_state_t)&ctx.__ss, &count);
#endif

  return ctx;
}

static inline void *GetFP(_STRUCT_MCONTEXT ctx) {
#if defined(__x86_64__)
  void *fp = (void *)ctx.__ss.__rbp;
#elif defined(__arm64__)
  void *fp = (void *)ctx.__ss.__fp;
#endif
  return fp;
}

static inline void *GetPC(_STRUCT_MCONTEXT ctx) {
#if defined(__x86_64__)
  void *pc = (void *)ctx.__ss.__rip;
#elif defined(__arm64__)
  void *pc = (void *)ctx.__ss.__pc;
#endif
  return pc;
}

static bool SafeReadMemory(vm_address_t src, void *dest, size_t len) {
  vm_size_t readSize = len;
  return vm_read_overwrite(mach_task_self(), src, len, (pointer_t)dest, &readSize) == KERN_SUCCESS;
}

static void TheadStackUnwind(vm_address_t *buffer, unsigned max, unsigned *frames, pthread_t thread,
                             void *frame) {
  void *next;
  void *stacktop = pthread_get_stackaddr_np(thread);
  void *stackbot = stacktop - pthread_get_stacksize_np(thread);

  *frames = 0;

  // make sure return address is never out of bounds
  stacktop -= (FP_LINK_OFFSET + 1) * sizeof(void *);

  if (!ISVALID(frame)) return;

  while (max--) {
    buffer[*frames] = *(vm_address_t *)(((void **)frame) + FP_LINK_OFFSET);
    (*frames)++;
    if (!SafeReadMemory((vm_address_t)frame, &next, sizeof(next))) {
      NSLog(@"Unable to read memory in stack unwinding. Address: %llu", (uint64_t)frame);
      return;
    }
    if (!ISVALID(next) || next <= frame) return;
    frame = next;
  }
}

static NSString *getThreadName(pthread_t thread) {
  char name[128];
  if (pthread_getname_np(thread, name, sizeof(name)) != 0) {
    return @"";
  }
  return [NSString stringWithCString:name encoding:NSASCIIStringEncoding];
}

int LPSnapshotThreadStackTrace(void **buffer, int size, pthread_t thread) {
  if (buffer == NULL || size == 0) {
    return 0;
  }

  unsigned int num_frames;
  _STRUCT_MCONTEXT ctx = GetContext(thread);
  // Current instruction.
  void *frame = GetFP(ctx);
  buffer[0] = GetPC(ctx);

  TheadStackUnwind((vm_address_t *)buffer + 1, --size, &num_frames, thread, frame);

  // Index 0 is used by the current instruction.
  num_frames += 1;

  while (num_frames >= 1 && buffer[num_frames - 1] == NULL) {
    num_frames -= 1;
  }
  return num_frames;
}

int LPSnapshotAllThreadsStackTrace(LPThreadStackTrace *buffer, int size) {
  if (buffer == NULL || size == 0) {
    return 0;
  }

  mach_msg_type_number_t threads_count;
  thread_act_array_t threads_list;
  task_threads(mach_task_self(), &threads_list, &threads_count);
  int count;

  for (count = 0; count < threads_count && count < size; ++count) {
    pthread_t pthread = pthread_from_mach_thread_np(threads_list[count]);
    int frames = LPSnapshotThreadStackTrace(buffer[count].stacktrace,
                                            kLPThreadStackTraceMaxFramesCount, pthread);
    buffer[count].frames = frames;
    buffer[count].thread = pthread;
  }

  return count;
}

NSArray<NSString *> *LPStackTraceString(void **buffer, int size) {
  if (buffer == NULL || size == 0) {
    return @[];
  }

  char **strings = backtrace_symbols(buffer, size);

  if (strings == NULL) {
    return @[];
  }

  NSMutableArray<NSString *> *result = [NSMutableArray array];
  for (int i = 0; i < size; i++) {
    [result addObject:[NSString stringWithCString:strings[i] encoding:NSASCIIStringEncoding]];
  }
  free(strings);
  return result;
}

NSArray<NSArray<NSString *> *> *LPStackTracesString(LPThreadStackTrace *buffer, int size) {
  if (buffer == NULL || size == 0) {
    return @[];
  }

  NSMutableArray<NSArray<NSString *> *> *result = [NSMutableArray array];
  for (int i = 0; i < size; i++) {
    NSMutableArray<NSString *> *threadResult = [NSMutableArray array];
    [threadResult addObject:[NSString stringWithFormat:@"Thread %d name: %@", i,
                                                       getThreadName(buffer[i].thread)]];
    [threadResult addObjectsFromArray:LPStackTraceString(buffer[i].stacktrace, buffer[i].frames)];

    [result addObject:threadResult];
  }

  return result;
}
#endif