#pragma once

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

#include "trace.h"

namespace async {

struct continuation : std::suspend_always {
    void await_suspend(std::coroutine_handle<>) noexcept {
        if (parent) {
            parent.resume();
        }
    }
    std::coroutine_handle<> parent;
};

template <typename T = void>
struct task;

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

    ~task() {
        if (h) {
            h.destroy();
        }
    }

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

        std::coroutine_handle<> parent;
    };

    // awaitable
    bool await_ready() {
        h.promise().parent = {};
        h.resume();
        if (h.done()) {
            return true;
        }
        return false;
    }
    void await_suspend(std::coroutine_handle<> ha) { h.promise().parent = ha; }
    void await_resume() {}

    handle_type h;
};

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

    ~task() {
        if (h) {
            h.destroy();
        }
    }

    struct promise_type {
        task get_return_object() {
            return {.h = handle_type::from_promise(*this)};
        }
        std::suspend_always initial_suspend() noexcept { return {}; }
        continuation final_suspend() noexcept { return {.parent = parent}; }
        void return_value(T&& value) { ret_value = std::move(value); }
        void unhandled_exception() {}
        template <std::convertible_to<T> From>
        continuation 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 = {};
        h.resume();
        if (h.promise().result_ready || h.done()) {
            return true;
        }
        return false;
    }
    void await_suspend(std::coroutine_handle<> ha) { h.promise().parent = ha; }
    T await_resume() {
        h.promise().result_ready = false;
        return std::move(h.promise().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};
}

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);
        if (parent) {
            schedule(parent);
        }
    }
    T await_resume() {
        waiting = false;
        TRACE(tracing::TraceEvent::kAsyncGimmeResume);
        return std::move(stuff);
    }

    // parent interface
    auto feed(T&& s) {
        struct awaitable {
            bool await_ready() {
                g.parent = {};
                g.ha.resume();
                return g.waiting;
            }
            void await_suspend(std::coroutine_handle<> h) { g.parent = h; }
            void await_resume() {}

            gimme<T>& g;
        };

        if (!waiting) {
            __builtin_trap();
        }
        if (!ha) {
            __builtin_trap();
        }
        stuff = std::move(s);
        return awaitable{.g = *this};
    }

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

}  // namespace async