#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 "tsan_support.h"
#include <mutex>
#include <unordered_map>
#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();
  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);
void context_loop(fcontext::transfer_t initial_transfer);

struct cached_fiber {
  continuation_t entry_cont_;
  void *tsan_fiber_;
};
extern std::mutex cache_mtx;
extern std::unordered_map<char *, cached_fiber> cached_fibers;

class lambda_capture_base {
 public:
  virtual continuation run(continuation_t cont) = 0;
  virtual void deconstruct() = 0;
};

template<typename F>
class lambda_capture : public lambda_capture_base {
 public:
  explicit lambda_capture(F lambda) : lambda_{lambda}, previous_tsan_fiber_{nullptr} {}

  void deconstruct() override {
    //TODO: re-add deconstructor call
//    ~lambda_capture<F>();
  }

  continuation run(continuation_t cont) override {
    return lambda_(continuation{cont, previous_tsan_fiber_});
  }

  void set_previous_tsan_fiber(void *tsan_fiber) {
    previous_tsan_fiber_ = tsan_fiber;
  }

 private:
  F lambda_;
  void *previous_tsan_fiber_;
};

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_pointer_t stack_base = stack_memory + stack_size;
  stack_pointer_t stack_limit = stack_memory;

  auto *captured_lambda = place_lambda_capture(lambda, lambda_memory);
  continuation_t fiber_cont;
  void *fiber_tsan;
  {
    std::lock_guard<std::mutex> lock{cache_mtx};

    auto cache_result = cached_fibers.find(stack_memory);
    if (cache_result == cached_fibers.end()) {
      // No entry found...create it
      fcontext::callback_t callback = context_loop;
      fcontext::continuation_t
          initial_context = fcontext::make_fcontext(stack_base, stack_base - stack_limit, callback);
      void *new_fiber = __tsan_create_fiber(0);
      void *old_fiber = __tsan_get_current_fiber();

      __tsan_switch_to_fiber(new_fiber, 0);
      fcontext::transfer_t fiber_context = fcontext::jump_fcontext(initial_context, old_fiber);

      cached_fibers[stack_memory].tsan_fiber_ = new_fiber;
      cached_fibers[stack_memory].entry_cont_ = fiber_context.continuation;
    }

    fiber_tsan = cached_fibers[stack_memory].tsan_fiber_;
    fiber_cont = cached_fibers[stack_memory].entry_cont_;
  }

  captured_lambda->set_previous_tsan_fiber(__tsan_get_current_fiber());
  __tsan_switch_to_fiber(fiber_tsan, 0);
  fcontext::transfer_t transfer = fcontext::jump_fcontext(fiber_cont, captured_lambda);

  if (transfer.data) {
    return continuation{transfer.continuation, transfer.data};
  } else {
    return continuation{nullptr, nullptr};
  }
}
#endif

}

#endif //CONTEXT_SWITCHER_CONTEXT_SWITCHER_H_
