#pragma once

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

#include "trace.h"

namespace async {

struct task_final_suspend {
    bool await_ready() noexcept(true) { return false; }
    void await_suspend(std::coroutine_handle<> h) noexcept(true) {
        if (parent) {
            TRACE(tracing::TraceEvent::kAsyncCallParent);
            parent();
            TRACE(tracing::TraceEvent::kAsyncCallParentDone);

            if (parent && parent.done()) {
                TRACE(tracing::TraceEvent::kAsyncDestroy);
                parent.destroy();
            }
        }
    }
    void await_resume() noexcept(true) {}

    std::coroutine_handle<> parent;
};

template <typename T>
struct gimme {
    // child interface
    bool await_ready() { return false; }
    void await_suspend(std::coroutine_handle<> h) {
        ha = h;
        waiting = true;
        TRACE(tracing::TraceEvent::kAsyncGimmeWaiting);
    }
    T await_resume() {
        waiting = false;
        TRACE(tracing::TraceEvent::kAsyncGimmeResume);
        return std::move(stuff);
    }

    // parent interface
    void feed(T&& s) {
        if (!waiting) {
           __builtin_trap();
        }
        if (!ha) {
           __builtin_trap();
        }
        stuff = s;
        ha.resume();
    }

    bool waiting = false;
    std::coroutine_handle<> ha;
    T stuff;
};

template <typename T = void>
struct task;

template <>
struct task<void> {
    struct promise_type;
    using handle_type = std::coroutine_handle<promise_type>;

    struct promise_type {
        task get_return_object() {
            return {.h = handle_type::from_promise(*this)};
        }
        std::suspend_always initial_suspend() noexcept { return {}; }
        task_final_suspend final_suspend() noexcept {
            return {.parent = parent};
        }
        void return_void() {}
        void unhandled_exception() {
            TRACE(tracing::TraceEvent::kAsyncException);
        }

        std::coroutine_handle<> parent;
    };

    // awaitable
    bool await_ready() {
        TRACE(tracing::TraceEvent::kAsyncCoAwait);
        h();
        if (h.done()) {
            TRACE(tracing::TraceEvent::kAsyncDestroy);
            h.destroy();
            return true;
        }
        return false;
    }
    void await_suspend(std::coroutine_handle<> ha) {
        TRACE(tracing::TraceEvent::kAsyncSuspend);
        h.promise().parent = ha;
    }
    void await_resume() {}

    std::coroutine_handle<promise_type> h;
};

template <typename T>
struct task {
    struct promise_type;
    using handle_type = std::coroutine_handle<promise_type>;

    struct promise_type {
        task get_return_object() {
            return {.h = handle_type::from_promise(*this)};
        }
        std::suspend_always initial_suspend() noexcept { return {}; }
        task_final_suspend final_suspend() noexcept {
            return {.parent = parent};
        }
        void return_value(T&& value) { ret_value = std::move(value); }
        void unhandled_exception() {
            TRACE(tracing::TraceEvent::kAsyncException);
        }
        template<std::convertible_to<T> From>
        task_final_suspend yield_value(From&& value) {
            ret_value = std::forward<From>(value);
            result_ready = true;
            return {.parent = parent};
        }

        T ret_value;
        bool result_ready = false;
        std::coroutine_handle<> parent;
    };

    // awaitable
    bool await_ready() {
        h.promise().parent = {};
        TRACE(tracing::TraceEvent::kAsyncCoAwait);
        h();
        if (h.promise().result_ready) {
            return true;
        }
        if (h.done()) {
            destroyed = true;
            ret_value = std::move(h.promise().ret_value);
            TRACE(tracing::TraceEvent::kAsyncDestroy); h.destroy();
            return true;
        }
        return false;
    }
    void await_suspend(std::coroutine_handle<> ha) {
        TRACE(tracing::TraceEvent::kAsyncSuspend);
        h.promise().parent = ha;
    }
    T await_resume() {
        if (!destroyed) {
            h.promise().result_ready = false;
            return std::move(h.promise().ret_value);
        }
        return std::move(ret_value);
    }

    bool destroyed = false;
    T ret_value;
    std::coroutine_handle<promise_type> h;
};

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

    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(bool (*idle_function)());
void step();

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