#ifndef CONTEXT_SWITCHER_CONTEXT_SWITCHER_H_
#define CONTEXT_SWITCHER_CONTEXT_SWITCHER_H_

#include <utility>
#include <cstdio>
#include <cstdint>
#include <new>

#include "context.h"
#include "continuation.h"

#ifdef THREAD_SANITIZER
#include <sanitizer/tsan_interface.h>
#endif

/**
 * Exposes main interface to start a stackful coroutine and to
 * switch between different context.
 *
 * Both cscontext and fcontext implementations require no global state
 * (everything is passed over in the context transfer routines).
 *
 * Thread sanitizer support uses global state to cache coroutine handles,
 * as recycling them frequently causes problems. To do this coroutines
 * started with the same stack frame are 'recycled'. For the end user this
 * means that the stack memory used by the coroutines MUST NOT be used by anything
 * else between usages (you can pool your stack memory as long as you always use
 * the same stacks to create coroutines).
 */
namespace context_switcher {
// cscontext implementation.
// Creates coroutines on the fly while entering them, making it very fast for
// extremely short lived routines (e.g. to implement a cactus stack like construct).
#ifdef CS_USE_CSCONTEXT
continuation switch_context(continuation &&cont);

template<typename F>
struct lambda_capture {
  template<typename FARG>
  explicit lambda_capture(FARG &&lambda) : lambda_{std::forward<FARG>(lambda)} {}

  continuation operator()(continuation_t cont) {
    return lambda_(continuation{cont});
  }

 private:
  F lambda_;
};

template<typename F>
static lambda_capture<typename std::remove_reference<F>::type> *place_lambda_capture(F &&lambda, char *memory) {
  return new(memory) lambda_capture<typename std::remove_reference<F>::type>(std::forward<F>(lambda));
}

template<typename T>
cscontext::continuation_t lambda_capture_callback(cscontext::continuation_t continuation_pointer,
                                                  void *lambda_capture_param) {
  // Perform Call
  T *lambda_capture = reinterpret_cast<T *>(lambda_capture_param);
  continuation cont = (*lambda_capture)(continuation_pointer);

  // Free resources and switch to result_continuation (this execution thread is finished with the return)
  lambda_capture->~T();
  return cont.consume();
}

template<typename F>
continuation enter_context(stack_pointer_t stack_memory, size_t stack_size, F &&lambda) {
  stack_pointer_t lambda_memory = stack_memory + stack_size - sizeof(lambda_capture<F>);
  auto *captured_lambda = place_lambda_capture(std::forward<F>(lambda), lambda_memory);

  stack_pointer_t stack_base = lambda_memory - 64;
  stack_pointer_t stack_limit = stack_memory;

  cscontext::callback_t callback = lambda_capture_callback<lambda_capture<F>>;

  void *result = cscontext::__cs_enter_context(stack_base, captured_lambda, callback, stack_limit);
  return continuation{result};
}
#endif

// fcontext implementation (no thread sanitizer).
// Creates a new coroutine, then jumps control to it.
// This creation followed by a jump is slightly slower than the cscontext direct
// creation and jump, but the boost.context assembly is available for more plaforms.
#if defined(CS_USE_FCONTEXT) and not defined(THREAD_SANITIZER)
continuation switch_context(continuation &&cont);

template<typename F>
struct lambda_capture {
  template<typename FARG>
  explicit lambda_capture(FARG &&lambda) : lambda_{std::forward<FARG>(lambda)} {}

  continuation operator()(continuation_t cont) {
    return lambda_(continuation{cont});
  }

 private:
  F lambda_;
};

template<typename T>
void lambda_capture_callback(fcontext::transfer_t transfer) {
  // Perform Call
  T *lambda_capture = reinterpret_cast<T *>(transfer.data);
  continuation cont = (*lambda_capture)(transfer.continuation);

  // Free resources and switch to result_continuation (this execution thread is finished with the return)
  lambda_capture->~T();

  continuation_t cont_pointer = cont.consume();
  if (cont_pointer == nullptr) {
    printf("Error!!!\n");
  }
  fcontext::jump_fcontext(cont_pointer, (void *) 0);
}

template<typename F>
static lambda_capture<typename std::remove_reference<F>::type> *place_lambda_capture(F &&lambda, char *memory) {
  return new(memory) lambda_capture<typename std::remove_reference<F>::type>(std::forward<F>(lambda));
}

template<typename F>
continuation enter_context(stack_pointer_t stack_memory, size_t stack_size, F &&lambda) {
  stack_pointer_t lambda_memory = stack_memory + stack_size - sizeof(lambda_capture<F>);
  auto *captured_lambda = place_lambda_capture(std::forward<F>(lambda), lambda_memory);

  stack_pointer_t stack_base = lambda_memory - 64;
  stack_pointer_t stack_limit = stack_memory;

  fcontext::callback_t callback = lambda_capture_callback<lambda_capture<F>>;

  fcontext::continuation_t new_context = fcontext::make_fcontext(stack_base, stack_base - stack_limit, callback);
  fcontext::transfer_t transfer = fcontext::jump_fcontext(new_context, captured_lambda);
  if (transfer.data) {
    return continuation{transfer.continuation};
  } else {
    return continuation{nullptr};
  }
}
#endif

// thread sanitizer implementation.
// Thread sanitizer, while having fiber support since mid 2019, has trouble
// with the creation/destruction of a high amount of short lived coroutines.
// To circumvent this we create one coroutine per stack base address and cache them in global variables.
// Please use with care and only for testing purposes.
#if defined(THREAD_SANITIZER)
continuation switch_context(continuation &&cont);

template<typename F>
struct lambda_capture {
  template<typename FARG>
  explicit lambda_capture(FARG &&lambda) : lambda_{std::forward<FARG>(lambda)} {}

  continuation operator()(continuation &&cont) {
    return lambda_(std::move(cont));
  }

 private:
  F lambda_;
};

template<typename T>
void lambda_capture_callback(fcontext::transfer_t transfer) {
  fcontext::fiber_data received_fiber_data = *(fcontext::fiber_data *) transfer.data;

  // Perform Call
  T *lambda_capture = reinterpret_cast<T *>(received_fiber_data.data);
  continuation cont = (*lambda_capture)({transfer.continuation, received_fiber_data.old_fiber});

  // Free resources and switch to result_continuation (this execution thread is finished with the return)
  lambda_capture->~T();

  continuation_t cont_pointer = cont.consume();
  void *last_fiber = __tsan_get_current_fiber();
  fcontext::fiber_data send_fiber_data{last_fiber, nullptr, true};

  __tsan_switch_to_fiber(cont.get_tsan_fiber(), 0);
  fcontext::jump_fcontext(cont_pointer, &send_fiber_data);
}

template<typename F>
static lambda_capture<typename std::remove_reference<F>::type> *place_lambda_capture(F &&lambda, char *memory) {
  return new(memory) lambda_capture<typename std::remove_reference<F>::type>(std::forward<F>(lambda));
}

template<typename F>
continuation enter_context(stack_pointer_t stack_memory, size_t stack_size, F &&lambda) {
  stack_pointer_t lambda_memory = stack_memory + stack_size - sizeof(lambda_capture<F>);
  auto *captured_lambda = place_lambda_capture(std::forward<F>(lambda), lambda_memory);

  stack_pointer_t stack_base = lambda_memory - 64;
  stack_pointer_t stack_limit = stack_memory;

  fcontext::callback_t callback = lambda_capture_callback<lambda_capture<F>>;
  fcontext::continuation_t new_context = fcontext::make_fcontext(stack_base, stack_base - stack_limit, callback);

  void *next_fiber = __tsan_create_fiber(0);
  void *last_fiber = __tsan_get_current_fiber();
  fcontext::fiber_data send_fiber_data{last_fiber, captured_lambda, false};

  __tsan_switch_to_fiber(next_fiber, 0);
  fcontext::transfer_t transfer = fcontext::jump_fcontext(new_context, &send_fiber_data);
  fcontext::fiber_data received_fiber_data = *(fcontext::fiber_data *) transfer.data;

  if (received_fiber_data.finished) {
    __tsan_destroy_fiber(received_fiber_data.old_fiber);
    return continuation{nullptr, nullptr};
  } else {
    return continuation{transfer.continuation, received_fiber_data.old_fiber};
  }
}
#endif

}

#endif //CONTEXT_SWITCHER_CONTEXT_SWITCHER_H_
