#pragma once
#include <coroutine>
#include <variant>
#include <memory>
#include <type_traits>
#include <functional>

namespace dae
{
    struct WaitForSeconds { float seconds; };
    struct WaitUntil { std::function<bool()> predicate; };
    using YieldInstruction = std::variant<std::nullptr_t, WaitForSeconds, WaitUntil>;

    struct Task {
        struct promise_type {
            YieldInstruction current_yield = nullptr;

            Task get_return_object() {
                return Task{ std::coroutine_handle<promise_type>::from_promise(*this) };
            }
            std::suspend_never initial_suspend() noexcept { return {}; }
            std::suspend_always final_suspend() noexcept { return {}; }
            void unhandled_exception() { std::terminate(); }
            void return_void() {}

            std::suspend_always yield_value(std::nullptr_t) noexcept {
                current_yield = nullptr;
                return {};
            }
            std::suspend_always yield_value(WaitForSeconds wfs) noexcept {
                current_yield = wfs;
                return {};
            }
            std::suspend_always yield_value(WaitUntil wu) noexcept {
                current_yield = std::move(wu);
                return {};
            }
        };

        std::coroutine_handle<promise_type> handle;

        Task(std::coroutine_handle<promise_type> h) : handle(h) {}
        Task(Task&& other) noexcept : handle(std::exchange(other.handle, nullptr)) {}
        Task& operator=(Task&& other) noexcept {
            if (this != &other) {
                if (handle) handle.destroy();
                handle = std::exchange(other.handle, nullptr);
            }
            return *this;
        }
        Task(const Task&) = delete;
        Task& operator=(const Task&) = delete;
        ~Task()
        {
            if (handle) handle.destroy();
        }
    };

    template <typename F>
    concept TaskCallable = requires(F&& f) {
        { std::forward<F>(f)() } -> std::same_as<Task>;
    };

    struct CoroutineId
    {
        uint64_t id {0};
        bool operator==(const CoroutineId& other) const { return id == other.id; }
        bool operator!=(const CoroutineId& other) const { return id != other.id; }
        explicit operator bool() const { return id != 0; }
    };

    class CoroutineRunner final
    {
        class CoroutineRunnerImpl;
        std::unique_ptr<CoroutineRunnerImpl> impl;
    public:
        CoroutineRunner();
        ~CoroutineRunner();
        CoroutineRunner(const CoroutineRunner&) = delete;
        CoroutineRunner& operator=(const CoroutineRunner&) = delete;
        CoroutineRunner(CoroutineRunner&&) = delete;
        CoroutineRunner& operator=(CoroutineRunner&&) = delete;

        CoroutineId StartCoroutine(Task task);

        template <TaskCallable Callable>
        CoroutineId StartCoroutine(Callable&& callable) {
            return StartCoroutine(callable());
        }

        bool StopCoroutine(CoroutineId id);
        void StopAllCoroutines();
        void Update(float deltaTime);
    };
}