#pragma once

#include <chrono>
#include <coroutine>
#include <utility>

namespace async {
 
template<typename T = void>
struct task {
  struct promise_type;
  using handle_type = std::coroutine_handle<promise_type>;
 
  struct maybe_suspend {
    bool await_ready() noexcept(true) { return !parent; }
    void await_suspend(std::coroutine_handle<>) noexcept(true) {
      if (parent) {
        parent();
      }
    }
    void await_resume() noexcept(true) { }
 
    std::coroutine_handle<> parent;
  };
 
  struct promise_type {
    task get_return_object() { return {.h = handle_type::from_promise(*this)}; }
    std::suspend_always initial_suspend() noexcept { return {}; }
    maybe_suspend final_suspend() noexcept { return { .parent = parent }; }
    void return_value(T&& value) { ret_value = std::move(value); }
    void unhandled_exception() {}
 
    T ret_value;
    std::coroutine_handle<> parent;
  };
 
  // awaitable
  bool await_ready() { return h.done(); }
  void await_suspend(std::coroutine_handle<> ha) {
    h();
    h.promise().parent = ha;
  }
  T await_resume() { return std::move(h.promise().ret_value); }
 
  std::coroutine_handle<promise_type> h;
};

template<>
struct task<void> {
  struct promise_type;
  using handle_type = std::coroutine_handle<promise_type>;
 
  struct maybe_suspend {
    bool await_ready() noexcept(true) { return !parent; }
    void await_suspend(std::coroutine_handle<>) noexcept(true) {
      if (parent) {
        parent();
      }
    }
    void await_resume() noexcept(true) { }
 
    std::coroutine_handle<> parent;
  };
 
  struct promise_type {
    task get_return_object() { return {.h = handle_type::from_promise(*this)}; }
    std::suspend_always initial_suspend() noexcept { return {}; }
    maybe_suspend final_suspend() noexcept { return { .parent = parent }; }
    void return_void() {}
    void unhandled_exception() {}
 
    std::coroutine_handle<> parent;
  };
 
  // awaitable
  bool await_ready() { return h.done(); }
  void await_suspend(std::coroutine_handle<> ha) {
    h();
    h.promise().parent = ha;
  }
  void await_resume() {}
 
  std::coroutine_handle<promise_type> h;
};

enum class AwaitableType {
    kUnknown = 0,
    kUartRx = 1,

    kNumTypes
};

void schedule(std::coroutine_handle<> h, int ms = 0);
void enqueue(std::coroutine_handle<> h, AwaitableType type);
void resume(AwaitableType type);  // typically called from an ISR

void main_loop(void (*idle_function)());

inline auto await(AwaitableType type) {
  struct awaitable {
    AwaitableType type;
 
    bool await_ready() { return false; };
    void await_suspend(std::coroutine_handle<> h) {
      enqueue(h, type);
    }
    void await_resume() {}
  };
 
  return awaitable{type};
}

inline auto delay(int ms) {
  struct awaitable {
    int ms;
 
    bool await_ready() { return false; };
    void await_suspend(std::coroutine_handle<> h) {
      schedule(h, ms);
    }
    void await_resume() {}
  };
 
  return awaitable{ms};
}

}  // namespace async