#pragma once

#include <coroutine>
#include <optional>
#include <utility>
#include <memory>
#include <vector>

#include "aisa/aisa.h"
#include "aisa/coroutine.h" // IWYU pragma: export

namespace aisa {

    template<typename CRTP> struct AsyncEval {
        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 std::suspend_always{};
            }
        }

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

        task<void> async_fetch_mem(byte_t *bytes, addr_t physical_addr, addr_t size)
        {
            while (true) {
                if (crtp().fetch_mem(bytes, physical_addr, size))
                    co_return;
                co_await std::suspend_always{};
            }
        }

        task<void> async_store_mem(addr_t physical_addr, const byte_t *bytes, addr_t size)
        {
            while (true) {
                if (crtp().store_mem(physical_addr, bytes, size))
                    co_return;
                co_await std::suspend_always{};
            }
        }

        task<bool> async_predicate(const Step &step)
        {
            if (step.predicate.has_value()) {
                regval_t pval = co_await crtp().async_load_reg(step.predicate->first);
                co_return pval == step.predicate->second;
            }
            co_return true;
        }

        task<void> async_load_source_registers(const Step &step, Wires &w)
        {
            w.source_vals.resize(step.source_regs.size());
            for (unsigned int i = 0; i < step.source_regs.size(); ++i)
                w.source_vals[i] = co_await crtp().async_load_reg(step.source_regs[i]);
        }

        task<void> async_load_source_memory(const Step &step, const MemInfo &mi, Wires &w)
        {
            w.memory_val.resize(mi.size);
            co_await crtp().async_fetch_mem(w.memory_val.data(), mi.physical_addr, mi.size);
        }

        task<void> async_load_sources(const Step &step, Wires &w)
        {
            co_await crtp().async_load_source_registers(step, w);

            if (step.mop == MOp::LOAD) {
                auto mi = step.meminfo(w);
                co_await crtp().async_load_source_memory(step, mi, w);
            }
        }

        task<void> async_write_destination_registers(const Step &step, const Wires &w)
        {
            for (unsigned int i = 0; i < step.destination_regs.size(); ++i)
                co_await crtp().async_store_reg(step.destination_regs[i], w.destination_vals[i]);
        }

        task<void> async_write_destination_memory(const Step &step, const MemInfo &mi, const Wires &w)
        {
            co_await crtp().async_store_mem(mi.physical_addr, w.memory_val.data(), mi.size);
        }

        task<void> async_write_destinations(const Step &step, const Wires &w)
        {
            if (w.aborted)
                co_return;

            co_await crtp().async_write_destination_registers(step, w);

            if (step.mop == MOp::STORE) {
                auto mi = step.meminfo(w);
                co_await crtp().async_write_destination_memory(step, mi, w);
            }
        }

        task<void> async_push_task(std::unique_ptr<const Task> &&task)
        {
            while (true) {
                if (crtp().push_task(std::move(task)))
                    co_return;
                co_await std::suspend_always{};
            }
        }

        task<void> async_new_task(std::unique_ptr<const Task> &&task, regval_t environment_val)
        {
            auto rn = task->environment;
            co_await crtp().async_push_task(std::move(task));
            co_await crtp().async_store_reg(rn, environment_val);
        }

        task<Wires> async_run_step(const Step &step)
        {
            Wires w;

            if (!co_await crtp().async_predicate(step))
                co_return std::move(w);

            co_await crtp().async_load_sources(step, w);

            step.evaluate(w);

            co_await crtp().async_write_destinations(step, w);

            if (w.new_task.has_value())
                co_await crtp().async_new_task(std::move(w.new_task->first), w.new_task->second);

            co_return std::move(w);
        }

        task<const Task *> async_top_task()
        {
            while (true) {
                if (auto rv = crtp().top_task(); rv.has_value())
                    co_return &**rv;
                co_await std::suspend_always{};
            }
        }

        task<void> async_pop_task()
        {
            while (true) {
                if (crtp().pop_task())
                    co_return;
                co_await std::suspend_always{};
            }
        }

        task<std::optional<std::unique_ptr<const Step>>> async_fetch_step()
        {
            auto task = co_await crtp().async_top_task();

            auto rn = task->environment;
            auto rv = co_await crtp().async_load_reg(rn);

            auto step = task->step(rv);
            if (step.has_value()) {
                co_await crtp().async_store_reg(rn, step->second);
                co_return std::move(step->first);
            } else {
                co_await crtp().async_pop_task();
                co_return {};
            }
        }

        task<std::pair<std::unique_ptr<const Step>, Wires>> async_fetch_and_run_step()
        {
            while (true) {
                if (auto step = co_await crtp().async_fetch_step(); step.has_value()) {
                    auto wires = co_await crtp().async_run_step(**step);
                    co_return {std::move(*step), std::move(wires)};
                }
            }
        }

        task<void> async_run_subtree()
        {
            while (true) {
                if (auto step = co_await crtp().async_fetch_step(); step.has_value())
                    co_await crtp().async_run_step_and_subtree(**step);
                else
                    break;
            }
        }

        task<Wires> async_run_step_and_subtree(const Step &step)
        {
            auto w = co_await crtp().async_run_step(step);

            if (w.new_task.has_value())
                co_await crtp().async_run_subtree();

            co_return std::move(w);
        }

        task<void> async_setup_initial_task(const ISA &isa)
        {
            auto task = isa.initial_task();

            co_await crtp().async_new_task(std::move(task.first), task.second);
        }

        task<void> async_run(const ISA &isa)
        {
            co_await crtp().async_setup_initial_task(isa);

            co_await crtp().async_run_subtree();
        }

    };

}