#include <algorithm>
#include <iostream>
#include <optional>
#include <set>
#include <string>
#include <utility>
#include <vector>

#include "sim/sim.h"
#include "sim/queue.h"
#include "uarch/fetch.h"
#include "uarch/types.h"
#include "util/assert.h"

namespace uarch {

    FetchStage::FetchStage(sim::Scheduler &scheduler, const aisa::ISA &isa, sim::Queue<FillReq> &fillreqq, sim::Queue<Fill> &fillq, sim::Queue<Uop> &uopq)
        : sim::Schedulable(scheduler)
        , fillreqq(fillreqq)
        , fillq(fillq)
        , uopq(uopq)
    {
        fillreqq.add_writer(this);
        fillq.add_reader(this);
        uopq.add_writer(this);

        auto [task, env] = isa.initial_task();
        store_reg(task->environment, env);
        push_task(std::move(task));
    }

    void FetchStage::clock()
    {
        bool sent_uop = false;
        while (!sent_uop) {
            while (!step) {
                const auto &task = *top_task().value();
                auto s = task.step(load_reg(task.environment).value());
                if (s.has_value()) {
                    step = std::move(s->first);
                    std::cout << "fetch step " << step->disasm() << "\n";
                    wires = {};
                    outstanding_fill = false;
                    fill_complete = false;
                    ASSERT(store_reg(task.environment, s->second), "Could not write next environment value");
                } else {
                    pop_task();
                }
            }

            if (step->predicate.has_value()) {
                auto x = load_reg(step->predicate->first);
                if (x.has_value()) {
                    if (*x != step->predicate->second)
                        goto nextstep;
                } else {
                    std::cout << "fetch stalls on predicate register";
                    break;
                }
            }

            wires.source_vals.reserve(step->source_regs.size());
            for (unsigned int i = wires.source_vals.size(); i < step->source_regs.size(); ++i) {
                auto x = load_reg(step->source_regs[i]);
                if (!x.has_value()) {
                    std::cout << "fetch stalls on source register";
                    break;
                }
                wires.source_vals.emplace_back(*x);
            }

            if (step->mop == aisa::MOp::LOAD && !fill_complete) {
                auto mi = step->meminfo(wires);
                if (outstanding_fill) {
                    for (unsigned int i = 0; i < fillq.available(); ++i) {
                        auto &fill = fillq.peek(i);
                        if (fill.physical_addr == mi.physical_addr) {
                            std::cout << "fetch fills\n";
                            wires.memory_val = std::move(fill.bytes);
                            outstanding_fill = false;
                            fill_complete = true;
                        }
                    }
                    if (!fill_complete)
                        break;
                } else {
                    std::cout << "fetch sends fill request\n";
                    FillReq req;
                    req.fillq = &fillq;
                    req.physical_addr = mi.physical_addr;
                    req.size = mi.size;
                    fillreqq.write(std::move(req));
                    outstanding_fill = true;
                    fill_complete = false;
                    break;
                }
            }

            step->evaluate(wires);

            for (unsigned int i = 0; i < step->destination_regs.size(); ++i)
                ASSERT(store_reg(step->destination_regs[i], wires.destination_vals[i]), "Could not write destination register");

            ASSERT(step->mop != aisa::MOp::STORE, "No stores allowed in fetch");

            if (wires.new_task.has_value()) {
                std::cout << "fetch sends subtask " << wires.new_task->first->disasm() << " downstream\n";
                Uop uop;
                static_cast<aisa::VectorRF &>(uop) = *this;
                ASSERT(uop.store_reg(wires.new_task->first->environment, wires.new_task->second), "Could not write initial environment to uop");
                ASSERT(uop.push_task(std::move(wires.new_task->first)), "Could not push subtask to uop");
                uopq.write(std::move(uop));
                sent_uop = true;
            }

nextstep:
            step = nullptr;
            wires = {};
            outstanding_fill = false;
            fill_complete = false;
        }

        while (fillq.available()) {
            std::cout << "fetch consumes fill\n";
            fillq.read();
        }
    }

}