//  ================================================================
//
//    stackless_coroutine.c
//
//  ----------------------------------------------------------------
//
//    (C) 2025 Mitya Selivanov <guattari.tech>
//
//    Any use of this code is prohibited.
//
//  ================================================================

#ifndef STACKLESS_COROUTINE_HEADER_GUARD_
#define STACKLESS_COROUTINE_HEADER_GUARD_

//  ================================================================
//
//    Options
//
//  ================================================================

#ifndef ENABLE_TESTING
#define ENABLE_TESTING 0
#endif

//  ================================================================

typedef signed             i32;
typedef unsigned char      b8;

#ifndef NULL
#define NULL ((void *) 0)
#endif

#ifdef __GNUC__
#define CORO_FALLTHROUGH_ __attribute__((fallthrough));
#else
#define CORO_FALLTHROUGH_
#endif

#ifdef __cplusplus
extern "C" {
#endif

typedef struct {
  i32 _;
} Coro_Void_;

typedef void (*Coro_State_Machine)(void *self_void_);

#define CORO_PROMISE_DATA_             \
  struct {                             \
    i32                _index;         \
    i32                _id;            \
    Coro_State_Machine _state_machine; \
  }

typedef struct {
  CORO_PROMISE_DATA_;
} Coro_Promise_Void;

#if defined(__GNUC__) || defined(__clang__)
#  pragma GCC diagnostic push
#  pragma GCC diagnostic ignored "-Wunused-function"
#  pragma GCC diagnostic ignored "-Wunknown-pragmas"
#endif

#define CORO_PROMISE_FROM(coro_) (*((Coro_Promise_Void *) (coro_)))

//  TODO: Only use static dispatch?

#ifdef ENABLE_CORO_CUSTOM_DISPATCH
//  Application should implement this procedure if custom dispatch is enabled.
void dispatch_stackless_coroutine(void *promise);
#else
static void dispatch_stackless_coroutine(void *promise) {
  //  Dynamic dispatch by default.
  CORO_PROMISE_FROM(promise)._state_machine(promise);
}
#endif

#if defined(__GNUC__) || defined(__clang__)
#  pragma GCC diagnostic pop
#endif

#define CORO_PROMISE_STRUCT_(coro_, Ret_Type_, ...) \
  struct Coro_Promise_##coro_ {                     \
    CORO_PROMISE_DATA_;                             \
    Ret_Type_ return_value;                         \
    __VA_ARGS__                                     \
  }

#define CORO_PROC_DECL_(coro_) void coro_(void *self_void_)

#define CORO_LINE_() __LINE__

#define CORO_IMPL(coro_)                                                            \
  CORO_PROC_DECL_(coro_) {                                                          \
    b8 _Is_Inside_Coro = 1;                                                         \
    struct Coro_Promise_##coro_ *self = (struct Coro_Promise_##coro_ *) self_void_; \
    switch (self->_index) {                                                         \
      case 0: (void) _Is_Inside_Coro;

#define CORO_END            \
    }                       \
    (void) _Is_Inside_Coro; \
    self->_index = -1;      \
  }

#define CORO_DECL(coro_, Ret_Type_, ...)               \
  CORO_PROMISE_STRUCT_(coro_, Ret_Type_, __VA_ARGS__); \
  CORO_PROC_DECL_(coro_)

#define CORO_DECL_STATIC(coro_, Ret_Type_, ...)        \
  CORO_PROMISE_STRUCT_(coro_, Ret_Type_, __VA_ARGS__); \
  static CORO_PROC_DECL_(coro_)

#define CORO(coro_, Ret_Type_, ...)                    \
  CORO_PROMISE_STRUCT_(coro_, Ret_Type_, __VA_ARGS__); \
  CORO_IMPL(coro_)

#define CORO_DECL_VOID(coro_, ...)          \
  CORO_DECL(coro_, Coro_Void_, __VA_ARGS__)

#define CORO_VOID(coro_, ...)          \
  CORO(coro_, Coro_Void_, __VA_ARGS__)

#define CORO_STATIC(coro_, Ret_Type_, ...)             \
  CORO_PROMISE_STRUCT_(coro_, Ret_Type_, __VA_ARGS__); \
  static CORO_IMPL(coro_)

#define CORO_STATIC_VOID(coro_, ...)          \
  CORO_STATIC(coro_, Coro_Void_, __VA_ARGS__)

#define coro_resume(promise_) \
  dispatch_stackless_coroutine(&(promise_))

#define coro_next(promise_) \
  (dispatch_stackless_coroutine(&(promise_)), (promise_).return_value)

#define yield(...)                     \
  do {                                 \
    (void) _Is_Inside_Coro;            \
    self->_index       = CORO_LINE_(); \
    self->return_value = __VA_ARGS__;  \
    return;                            \
  case CORO_LINE_():;                  \
  } while (0)

#define yield_void               \
  do {                           \
    (void) _Is_Inside_Coro;      \
    self->_index = CORO_LINE_(); \
    return;                      \
  case CORO_LINE_():;            \
  } while (0)

#define async_return(...)             \
  do {                                \
    (void) _Is_Inside_Coro;           \
    self->_index       = -1;          \
    self->return_value = __VA_ARGS__; \
    return;                           \
  } while (0)

#define async_return_void   \
  do {                      \
    (void) _Is_Inside_Coro; \
    self->_index = -1;      \
    return;                 \
  } while (0)

#define await(promise_)                          \
  do {                                           \
    (void) _Is_Inside_Coro;                      \
    CORO_FALLTHROUGH_                            \
  case CORO_LINE_():                             \
    if ((promise_)._index != -1) {               \
      self->_index = CORO_LINE_();               \
      dispatch_stackless_coroutine(&(promise_)); \
    }                                            \
    if ((promise_)._index != -1)                 \
      return;                                    \
  } while (0)

#define yield_await(promise_)                       \
  do {                                              \
    (void) _Is_Inside_Coro;                         \
    CORO_FALLTHROUGH_                               \
  case CORO_LINE_():                                \
    if ((promise_)._index != -1) {                  \
      self->_index = CORO_LINE_();                  \
      dispatch_stackless_coroutine(&(promise_));    \
      self->return_value = (promise_).return_value; \
      return;                                       \
    }                                               \
  } while (0)

#define Promise_Of(coro_) struct Coro_Promise_##coro_

#define CORO_INITIAL_(id_, coro_)                      \
  ._index = 0, ._id = (id_), ._state_machine = (coro_)

#define CORO_INIT(promise_, coro_, ...)                                      \
  do {                                                                       \
    promise_ = (Promise_Of(coro_)) { CORO_INITIAL_(0, coro_), __VA_ARGS__ }; \
  } while (0)

#define CORO_INIT_ID(promise_, id_, ...)                                       \
  do {                                                                         \
    promise_ = (Promise_Of(coro_)) { CORO_INITIAL_(id_, coro_), __VA_ARGS__ }; \
  } while (0)

#define coro_finished(promise_) ((promise_)._index == -1)

#define coro_finished_n(return_, num_promises_, promises_) \
  do {                                                     \
    i32 num_ = (num_promises_);                            \
    i32 index_;                                            \
    (return_) = 1;                                         \
    for (index_ = 0; index_ < (num_); ++index_)            \
      if (!coro_finished((promises_)[index_])) {           \
        (return_) = 0;                                     \
        break;                                             \
      }                                                    \
  } while (0)

#define coro_finished_all(return_, promises_)                                         \
  coro_finished_n((return_), sizeof(promises_) / sizeof((promises_)[0]), (promises_))

#ifdef __cplusplus
}
#endif

#endif // STACKLESS_COROUTINE_HEADER_GUARD_

//  ================================================================
//
//    Test suite
//
//  ================================================================

#ifndef STACKLESS_COROUTINE_HEADER
#ifndef STACKLESS_COROUTINE_IMPL_GUARD_
#define STACKLESS_COROUTINE_IMPL_GUARD_

#if ENABLE_TESTING

#define TEST_FILE stackless_coroutine
#include "test.c"

CORO_DECL(
  test_foo,
  int,
);

CORO_IMPL(test_foo)
{
  async_return (42);
}
CORO_END

TEST("coroutine init") {
  Promise_Of(test_foo) promise;
  CORO_INIT(promise, test_foo, );
  REQUIRE(!coro_finished(promise));
}

TEST("coroutine init with value") {
  Promise_Of(test_foo) promise;
  CORO_INIT(promise, test_foo, .return_value = 42);
  REQUIRE(promise.return_value == 42);
  REQUIRE(!coro_finished(promise));
}

TEST("coroutine execute and return") {
  Promise_Of(test_foo) promise;
  CORO_INIT(promise, test_foo, );
  REQUIRE(coro_next(promise) == 42);
  REQUIRE(coro_finished(promise));
}

CORO_DECL_STATIC(
  test_bar,
  int,
);

CORO_IMPL(test_bar)
{
  yield_void;
  async_return (42);
}
CORO_END

TEST("coroutine execute two steps") {
  Promise_Of(test_bar) promise;
  CORO_INIT(promise, test_bar, .return_value = 0);
  coro_resume(promise);
  REQUIRE(promise.return_value == 0);
  coro_resume(promise);
  REQUIRE(promise.return_value == 42);
}

CORO_STATIC(
  test_gen,
  int,
  int i; int min; int max;
) {
  for (self->i = self->min; self->i < self->max; self->i++)
    yield (self->i);
  async_return (self->max);
}
CORO_END

TEST("coroutine generator") {
  Promise_Of(test_gen) promise;
  CORO_INIT(promise, test_gen, .min = 10, .max = 15);
  for (i32 i = 0; i <= 5; i++)
    REQUIRE(coro_next(promise) == 10 + i);
}

TEST("coroutine status finished") {
  Promise_Of(test_bar) promise;
  CORO_INIT(promise, test_bar, );
  REQUIRE(!coro_finished(promise));
  coro_resume(promise);
  REQUIRE(!coro_finished(promise));
  coro_resume(promise);
  REQUIRE(coro_finished(promise));
}

CORO_STATIC_VOID(test_task, )
{
  yield_void;
  yield_void;
  async_return_void;
}
CORO_END

TEST("coroutine task") {
  Promise_Of(test_task) promise;
  CORO_INIT(promise, test_task, );
  coro_resume(promise);
  REQUIRE(!coro_finished(promise));
  coro_resume(promise);
  REQUIRE(!coro_finished(promise));
  coro_resume(promise);
  REQUIRE(coro_finished(promise));
}

CORO_STATIC_VOID(
  test_nest_task,
  Promise_Of(test_task) promise;
) {
  CORO_INIT(self->promise, test_task, );
  await (self->promise);
  await (self->promise);
  await (self->promise);
}
CORO_END

TEST("coroutine nested task") {
  Promise_Of(test_nest_task) promise;
  CORO_INIT(promise, test_nest_task, );
  coro_resume(promise);
  REQUIRE(!coro_finished(promise));
  coro_resume(promise);
  REQUIRE(!coro_finished(promise));
  coro_resume(promise);
  REQUIRE(coro_finished(promise));
}

CORO_STATIC(
  test_nest_generator,
  int,
  Promise_Of(test_gen) promise;
) {
  CORO_INIT(self->promise, test_gen, .min = 1, .max = 3);
  yield_await (self->promise);
}
CORO_END

TEST("coroutine nested generator") {
  Promise_Of(test_nest_generator) promise;
  CORO_INIT(promise, test_nest_generator, );
  REQUIRE(coro_next(promise) == 1);
  REQUIRE(coro_next(promise) == 2);
  REQUIRE(coro_next(promise) == 3);
  REQUIRE(!coro_finished(promise));
  coro_resume(promise);
  REQUIRE(coro_finished(promise));
}

#undef TEST_FILE
#endif // ENABLE_TESTING

#endif // STACKLESS_COROUTINE_IMPL_GUARD_
#endif // STACKLESS_COROUTINE_HEADER