#pragma once

#include <coroutine>
#include <cstdint>
#include <initializer_list>
#include <iostream>
#include <optional>
#include <utility>
#include <vector>

#include "aisa/coroutine.h"

namespace aisa {

    using regnum_t = std::uint_fast64_t;
    using regval_t = std::uint64_t;

    template<typename CRTP> struct EvalState {
        CRTP & crtp() noexcept { return static_cast<CRTP &>(*this); }

        task<regval_t> async_load_reg(regnum_t rn)
        {
            while (true) {
                if (auto rv = crtp().load_reg(rn); rv.has_value())
                    co_return *rv;
                co_await suspend();
            }
        }

        task<void> async_store_reg(regnum_t rn, regval_t rv)
        {
            while (true) {
                if (crtp().store_reg(rn, rv))
                    co_return;
                co_await suspend();
            }
        }
    };

    struct Step {
        std::optional<std::pair<regnum_t, regval_t>> predicate;
        std::vector<regnum_t> source_regs;
        std::vector<regnum_t> destination_regs;

        std::optional<regnum_t> predicate_reg() const
        {
            if (predicate.has_value())
                return predicate->first;
            return {};
        }

        std::optional<regnum_t> expected_predicate_val() const
        {
            if (predicate.has_value())
                return predicate->second;
            return {};
        }

        virtual std::vector<regval_t> compute_destinations(const std::vector<regval_t> &source_vals) const = 0;

        template<typename State> task<void> eval(State &state) const
        {
            if (predicate.has_value()) {
                regval_t pval = co_await state.async_load_reg(predicate->first);
                if (pval != predicate->second)
                    co_return;
            }
            std::vector<regval_t> source_vals;
            source_vals.reserve(source_regs.size());
            for (unsigned int i = 0; i < source_regs.size(); ++i)
                source_vals.emplace_back(co_await state.async_load_reg(source_regs[i]));
            auto destination_vals = compute_destinations(source_vals);
            for (unsigned int i = 0; i < destination_regs.size(); ++i)
                co_await state.async_store_reg(destination_regs[i], destination_vals[i]);
        }
    };

}