#include <cstdio>
#include <csetjmp>
#include <cstring>
#include <chrono>

#include "fcontext/fcontext.h"

#include "fiber_call.h"

using namespace std;

// Settings for stack and benchmark
const size_t NUM_RUNS = 1000000;
const size_t STACK_SIZE = 512 * 1;
const char MAGIC_NUMBER = (unsigned char) 0xAB;

// Memory for custom stack and continuation semantics
char custom_stack[STACK_SIZE] = {0};
jmp_buf buffer;

// Example callback function and declaration of our assembly stack switching routine
extern "C" {
void custom_stack_callback(void *);

void __attribute__ ((noinline)) callback() {
  static volatile int tmp;
  tmp = 0; // Force at least a single memory write
}
}

long measure_function_call() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    callback();
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

long measure_stack_switch() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    custom_stack_callback(&custom_stack[STACK_SIZE]);
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

long measure_continuation() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    if (setjmp(buffer) == 0) {
      custom_stack_callback(&custom_stack[STACK_SIZE]);
    }
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

long measure_continuation_and_jump() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    if (setjmp(buffer) == 0) {
      custom_stack_callback(&custom_stack[STACK_SIZE]);
      longjmp(buffer, 1);
    }
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

void fcontext_callback_fast(fcontext_transfer_t transfer) {
  for (;;) {
    callback();
    jump_fcontext(transfer.ctx, nullptr);
  }
}

long measure_fcontext_fast() {
  fcontext_t context = make_fcontext(&custom_stack[STACK_SIZE], STACK_SIZE, &fcontext_callback_fast);

  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    context = jump_fcontext(context, nullptr).ctx;
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

void fcontext_callback_clean(fcontext_transfer_t transfer) {
  callback();
  jump_fcontext(transfer.ctx, nullptr);
}

long measure_fcontext_clean() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    fcontext_t context = make_fcontext(&custom_stack[STACK_SIZE], STACK_SIZE, &fcontext_callback_clean);
    jump_fcontext(context, nullptr);
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

void fcontext_callcc(fcontext_transfer_t transfer) {
  callback();
  jump_fcontext(jump_fcontext(transfer.ctx, nullptr).ctx, nullptr);
}

long measure_fcontext_callcc() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    fcontext_t context = make_fcontext(&custom_stack[STACK_SIZE], STACK_SIZE, &fcontext_callcc);
    jump_fcontext(jump_fcontext(context, nullptr).ctx, nullptr);
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

long measure_custom() {
  using namespace pls::internal::base;

  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    fiber_call(custom_stack, STACK_SIZE, [](continuation_t continuation) {
      callback();
      return continuation;
    });
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

int main() {
  memset(custom_stack, MAGIC_NUMBER, STACK_SIZE);

  auto time_cont_jump = measure_continuation_and_jump();
  auto time_cont = measure_continuation();
  auto time_stack = measure_stack_switch();
  auto time_func = measure_function_call();
  auto time_fcontext_fast = measure_fcontext_fast();
  auto time_fcontext_clean = measure_fcontext_clean();
  auto time_fcontext_calcc = measure_fcontext_callcc();
  auto time_custom = measure_custom();

  printf("Base\n");
  printf("Function Call    : %10ld, %5.5f\n", time_func, ((float) time_func / NUM_RUNS));
  printf("Longjmp\n");
  printf("Stack Switching  : %10ld, %5.5f\n", time_stack, ((float) time_stack / NUM_RUNS));
  printf("Full Continuation: %10ld, %5.5f\n", time_cont, ((float) time_cont / NUM_RUNS));
  printf("Jump Continuation: %10ld, %5.5f\n", time_cont_jump, ((float) time_cont_jump / NUM_RUNS));
  printf("Boost\n");
  printf("FContext Fast    : %10ld, %5.5f\n", time_fcontext_fast, ((float) time_fcontext_fast / NUM_RUNS));
  printf("FContext Clean   : %10ld, %5.5f\n", time_fcontext_clean, ((float) time_fcontext_clean / NUM_RUNS));
  printf("FContext CallCC  : %10ld, %5.5f\n", time_fcontext_calcc, ((float) time_fcontext_calcc / NUM_RUNS));
  printf("Custom\n");
  printf("Custom Fast Call : %10ld, %5.5f\n", time_custom, ((float) time_custom / NUM_RUNS));

  for (unsigned int i = 0; i < STACK_SIZE; i++) {
    if (custom_stack[i] != MAGIC_NUMBER) {
      printf("\n\nUsed stack size about %u bytes.\n", (STACK_SIZE - i));
      break;
    }
  }

  return 0;
}