#pragma once

#include <fmt/core.h>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "aisa/aisa.h"

namespace isa::fib {

    namespace Reg {
        enum {
            ENV_TOP,
            ENV_FIB,

            AREG,
            ATMP,

            PC,

            A,
            B,
            Q,
        };

        const char *disasm(aisa::regnum_t x)
        {
            switch (x) {
                case ENV_TOP: return "ENV_TOP";
                case ENV_FIB: return "ENV_FIB";
                case AREG: return "AREG";
                case ATMP: return "ATMP";
                case PC: return "PC";
                case A: return "A";
                case B: return "B";
                case Q: return "Q";
                default: return "???";
            }
        }
    }

    struct StepAdd : public aisa::Step {
        StepAdd(aisa::regnum_t q, aisa::regnum_t a, aisa::regnum_t b)
        {
            source_regs.resize(2);
            source_regs[0] = a;
            source_regs[1] = b;

            destination_regs.resize(1);
            destination_regs[0] = q;
        }

        std::string disasm(const aisa::Wires *w) const override
        {
            if (w)
                return fmt::format("{}({}) <- {}({}) + {}({})", Reg::disasm(destination_regs[0]), w->destination_vals[0], Reg::disasm(source_regs[0]), w->source_vals[0], Reg::disasm(source_regs[1]), w->source_vals[1]);
            return fmt::format("{} <- {} + {}", Reg::disasm(destination_regs[0]), Reg::disasm(source_regs[0]), Reg::disasm(source_regs[1]));
        }

        void evaluate(aisa::Wires &w) const override
        {
            w.destination_vals.resize(1);
            w.destination_vals[0] = w.source_vals[0] + w.source_vals[1];
        }
    };

    struct StepAddImmediate : public aisa::Step {
        aisa::regval_t x;

        StepAddImmediate(aisa::regnum_t q, aisa::regnum_t a, aisa::regval_t x)
            :x(x)
        {
            source_regs.resize(1);
            source_regs[0] = a;

            destination_regs.resize(1);
            destination_regs[0] = q;
        }

        std::string disasm(const aisa::Wires *w) const override
        {
            if (w)
                return fmt::format("{}({}) <- {}({}) + {}", Reg::disasm(destination_regs[0]), w->destination_vals[0], Reg::disasm(source_regs[0]), w->source_vals[0], x);
            return fmt::format("{} <- {} + {}", Reg::disasm(destination_regs[0]), Reg::disasm(source_regs[0]), x);
        }

        void evaluate(aisa::Wires &w) const override
        {
            w.destination_vals.resize(1);
            w.destination_vals[0] = w.source_vals[0] + x;
        }
    };

    template<unsigned int WORD> struct StepLoad : public aisa::Step {
        StepLoad(aisa::regnum_t q, aisa::regnum_t a)
        {
            source_regs.resize(1);
            source_regs[0] = a;

            destination_regs.resize(1);
            destination_regs[0] = q;

            mop = aisa::MOp::LOAD;
        }

        std::string disasm(const aisa::Wires *w) const override
        {
            if (w)
                return fmt::format("{}({}) <- [{}({})]", Reg::disasm(destination_regs[0]), w->destination_vals[0], Reg::disasm(source_regs[0]), w->source_vals[0]);
            return fmt::format("{} <- [{}]", Reg::disasm(destination_regs[0]), Reg::disasm(source_regs[0]));
        }

        aisa::MemInfo meminfo(const aisa::Wires &w) const override
        {
            aisa::MemInfo mi;
            mi.physical_addr = w.source_vals[0] * WORD;
            mi.size = WORD;
            return mi;
        }

        void evaluate(aisa::Wires &w) const override
        {
            w.destination_vals.resize(1);
            auto &q = w.destination_vals[0];
            q = 0;
            for (unsigned int i = 0; i < WORD; ++i)
                q |= static_cast<aisa::regval_t>(w.memory_val[i]) << (8 * i);
        }
    };

    struct StepMove : public aisa::Step {
        StepMove(aisa::regnum_t q, aisa::regnum_t a)
        {
            source_regs.resize(1);
            source_regs[0] = a;

            destination_regs.resize(1);
            destination_regs[0] = q;
        }

        std::string disasm(const aisa::Wires *w) const override
        {
            if (w)
                return fmt::format("{}({}) <- {}({})", Reg::disasm(destination_regs[0]), w->destination_vals[0], Reg::disasm(source_regs[0]), w->source_vals[0]);
            return fmt::format("{} <- {}", Reg::disasm(destination_regs[0]), Reg::disasm(source_regs[0]));
        }

        void evaluate(aisa::Wires &w) const override
        {
            w.destination_vals = w.source_vals;
        }
    };

    struct StepMoveImmediate : public aisa::Step {
        aisa::regval_t x;

        StepMoveImmediate(aisa::regnum_t q, aisa::regval_t x)
            : x(x)
        {
            destination_regs.resize(1);
            destination_regs[0] = q;
        }

        std::string disasm(const aisa::Wires *w) const override
        {
            if (w)
                return fmt::format("{}({}) <- {}", Reg::disasm(destination_regs[0]), w->destination_vals[0], x);
            return fmt::format("{} <- {}", Reg::disasm(destination_regs[0]), x);
        }

        void evaluate(aisa::Wires &w) const override
        {
            w.destination_vals.resize(1);
            w.destination_vals[0] = x;
        }
    };

    template<typename Subtask> struct StepSpawn : public aisa::Step {
        void evaluate(aisa::Wires &w) const override
        {
            w.new_task = {std::make_unique<Subtask>(), 0};
        }

        std::string disasm(const aisa::Wires *w) const override
        {
            return fmt::format("spawn");
        }
    };

    template<unsigned int WORD> struct StepStore : public aisa::Step {
        StepStore(aisa::regnum_t a, aisa::regnum_t d)
        {
            source_regs.resize(2);
            source_regs[0] = a;
            source_regs[1] = d;

            mop = aisa::MOp::STORE;
        }

        std::string disasm(const aisa::Wires *w) const override
        {
            if (w)
                return fmt::format("[{}({})] <- {}({})", Reg::disasm(source_regs[0]), w->source_vals[0], Reg::disasm(source_regs[1]), w->source_vals[1]);
            return fmt::format("[{}] <- {}", Reg::disasm(source_regs[0]), Reg::disasm(source_regs[1]));
        }

        aisa::MemInfo meminfo(const aisa::Wires &w) const override
        {
            aisa::MemInfo mi;
            mi.physical_addr = w.source_vals[0] * WORD;
            mi.size = WORD;
            return mi;
        }

        void evaluate(aisa::Wires &w) const override
        {
            w.memory_val.resize(WORD);
            for (unsigned int i = 0; i < WORD; ++i)
                w.memory_val[i] = w.source_vals[1] >> (8 * i);
        }
    };

    template<unsigned int WORD> struct TaskFib : public aisa::Task {
        TaskFib()
        {
            environment = Reg::ENV_FIB;
        }

        std::string disasm() const override
        {
            return "fib";
        }

        std::optional<std::pair<std::unique_ptr<const aisa::Step>, aisa::regval_t>> step(aisa::regval_t env) const override
        {
            switch (env) {
                case 0: return {{std::make_unique<StepLoad<WORD>>(Reg::A, Reg::PC), env+1}};
                case 1: return {{std::make_unique<StepAddImmediate>(Reg::PC, Reg::PC, 1), env+1}};
                case 2: return {{std::make_unique<StepLoad<WORD>>(Reg::B, Reg::PC), env+1}};
                case 3: return {{std::make_unique<StepAddImmediate>(Reg::AREG, Reg::PC, 1), env+1}};
                case 4: return {{std::make_unique<StepAdd>(Reg::Q, Reg::A, Reg::B), env+1}};
                case 5: return {{std::make_unique<StepStore<WORD>>(Reg::AREG, Reg::Q), env+1}};
            }

            return {};
        }
    };

    template<unsigned int WORD> struct TaskTop : public aisa::Task {
        TaskTop()
        {
            environment = Reg::ENV_TOP;
        }

        std::string disasm() const override
        {
            return "top";
        }

        std::optional<std::pair<std::unique_ptr<const aisa::Step>, aisa::regval_t>> step(aisa::regval_t env) const override
        {
            switch (env) {
                case 0: return {{std::make_unique<StepMoveImmediate>(Reg::PC, 0), env+1}};
                case 1: return {{std::make_unique<StepMoveImmediate>(Reg::AREG, 0), env+1}};
                case 2: return {{std::make_unique<StepMoveImmediate>(Reg::ATMP, 0), env+1}};
                case 3: return {{std::make_unique<StepStore<WORD>>(Reg::AREG, Reg::ATMP), env+1}};
                case 4: return {{std::make_unique<StepMoveImmediate>(Reg::AREG, 1), env+1}};
                case 5: return {{std::make_unique<StepMoveImmediate>(Reg::ATMP, 1), env+1}};
                case 6: return {{std::make_unique<StepStore<WORD>>(Reg::AREG, Reg::ATMP), env+1}};
                case 7: return {{std::make_unique<StepSpawn<TaskFib<WORD>>>(), env}};
            }

            return {};
        }
    };

    template<unsigned int WORD=1> struct Fib : public aisa::ISA {
        std::pair<std::unique_ptr<const aisa::Task>, aisa::regval_t> initial_task() const override
        {
            auto t = std::make_unique<TaskTop<WORD>>();
            return {std::move(t), 0};
        }
    };

}