#pragma once
#include <concepts>
#include <coroutine>
#include <exception>
#include <assert.h>

using namespace std;

enum TaskStates {
  STARTED, AWAIT, YIELD,
  EXCEPTION, RETURN,
  RETURN_VOID, DEAD
};

template<typename T>
struct Task {
  struct promise_type;

  using handle_type = std::coroutine_handle<promise_type>;

  struct promise_type {
    T value_;
    std::exception_ptr exception_;
    TaskStates state_{STARTED};

    Task get_return_object() {
      return Task(handle_type::from_promise(*this));
    }

    std::suspend_always initial_suspend() {
      state_ = AWAIT;
      return {};
    }

    std::suspend_always final_suspend() noexcept {
      return {};
    }

    void unhandled_exception() {
      state_ = EXCEPTION;
      exception_ = std::current_exception();
    }

    template<std::convertible_to<T> From> // C++20 concept
    void return_value(From &&from) {
      state_ = RETURN;
      value_ = std::forward<From>(from);
    }

    template<std::convertible_to<T> From> // C++20 concept
    std::suspend_always yield_value(From &&from) {
      state_ = YIELD;
      value_ = std::forward<From>(from);
      return {};
    }

    void return_void() {
      state_ = RETURN_VOID;
    }
  };

  handle_type h_;

  Task() {
  }

  Task(handle_type h) : h_(h) {
  }

  Task(const Task &t) : h_(t.h_) {
  }

  void destroy() {
    // this should make it safe to clal repeatedly
    if(h_.promise().state_ != DEAD) {
      h_.destroy();
      h_.promise().state_ = DEAD;
    }
  }

  T operator()() {
    assert(!h_.done());
    call();
    return std::move(h_.promise().value_);
  }

  bool done() {
    return h_.done();
  }

  TaskStates state() {
    return h_.promise().state_;
  }

  private:

  void call() {
    h_();

    if (h_.promise().exception_) {
      std::rethrow_exception(h_.promise().exception_);
    }
  }
};