#include "elna/backend/riscv.hpp"
#include <cassert>
#include <functional>
#include <memory>

namespace elna::riscv
{
    instruction::instruction(base_opcode opcode)
    {
        this->representation = static_cast<std::underlying_type<base_opcode>::type>(opcode);
    }

    instruction& instruction::i(x_register rd, funct3_t funct3, x_register rs1, std::uint32_t immediate)
    {
        this->representation |= (static_cast<std::underlying_type<x_register>::type>(rd) << 7)
            | (static_cast<std::underlying_type<funct3_t>::type>(funct3) << 12)
            | (static_cast<std::underlying_type<x_register>::type>(rs1) << 15)
            | (immediate << 20);

        return *this;
    }

    instruction& instruction::s(std::uint32_t imm, funct3_t funct3, x_register rs1, x_register rs2)
    {
        this->representation |= ((imm & 0x1f) << 7)
            | (static_cast<std::underlying_type<funct3_t>::type>(funct3) << 12)
            | (static_cast<std::underlying_type<x_register>::type>(rs1) << 15)
            | (static_cast<std::underlying_type<x_register>::type>(rs2) << 20)
            | ((imm & 0xfe0) << 20);

        return *this;
    }

    instruction& instruction::b(std::uint32_t imm, funct3_t funct3, x_register rs1, x_register rs2)
    {
        this->representation |= ((imm & 0x800) >> 4) | ((imm & 0x1e) << 7)
            | (static_cast<std::underlying_type<funct3_t>::type>(funct3) << 12)
            | (static_cast<std::underlying_type<x_register>::type>(rs1) << 15)
            | (static_cast<std::underlying_type<x_register>::type>(rs2) << 20)
            | ((imm & 0x7e0) << 20) | ((imm & 0x1000) << 19);

        return *this;
    }

    instruction& instruction::r(x_register rd, funct3_t funct3, x_register rs1, x_register rs2, funct7_t funct7)
    {
        this->representation |= (static_cast<std::underlying_type<x_register>::type>(rd) << 7)
            | (static_cast<std::underlying_type<funct3_t>::type>(funct3) << 12)
            | (static_cast<std::underlying_type<x_register>::type>(rs1) << 15)
            | (static_cast<std::underlying_type<x_register>::type>(rs2) << 20)
            | (static_cast<std::underlying_type<funct7_t>::type>(funct7) << 25);

        return *this;
    }

    instruction& instruction::u(x_register rd, std::uint32_t imm)
    {
        this->representation |= (static_cast<std::underlying_type<x_register>::type>(rd) << 7) | (imm << 12);

        return *this;
    }

    instruction& instruction::j(x_register rd, std::uint32_t imm)
    {
        this->representation |= (static_cast<std::underlying_type<x_register>::type>(rd) << 7)
            | (imm & 0xff000) | ((imm & 0x800) <<  9) | ((imm & 0x7fe) << 20) | ((imm & 0x100000) << 11);

        return *this;
    }

    const std::byte *instruction::cbegin() const
    {
        return reinterpret_cast<const std::byte *>(&this->representation);
    }

    const std::byte *instruction::cend() const
    {
        return reinterpret_cast<const std::byte *>(&this->representation) + sizeof(this->representation);
    }

    static void relocate(std::string_view name, address_t target, std::vector<reference>& references,
            std::vector<instruction>& instructions, std::shared_ptr<source::writer<std::byte>> writer)
    {
        references.push_back(reference());
        references.back().name = name;
        references.back().offset = writer->size() + instructions.size() * 4;
        references.back().target = target;
    }

    static void prologue(std::vector<instruction>& instructions)
    {
        instructions.push_back(instruction(base_opcode::opImm));
        instructions.push_back(instruction(base_opcode::store));
        instructions.push_back(instruction(base_opcode::store));
        instructions.push_back(instruction(base_opcode::opImm));
   }

    static void epilogue(const std::size_t stack_size, std::vector<instruction>& instructions)
    {
        instructions[0].i(x_register::sp, funct3_t::addi, x_register::sp, -stack_size);
        instructions[1].s(0, funct3_t::sw, x_register::sp, x_register::s0);
        instructions[2].s(4, funct3_t::sw, x_register::sp, x_register::ra);
        instructions[3].i(x_register::s0, funct3_t::addi, x_register::sp, stack_size);

        // Epilogue.
        instructions.push_back(instruction(base_opcode::load)
            .i(x_register::s0, funct3_t::lw, x_register::sp, 0));
        instructions.push_back(instruction(base_opcode::load)
            .i(x_register::ra, funct3_t::lw, x_register::sp, 4));
        instructions.push_back(instruction(base_opcode::opImm)
            .i(x_register::sp, funct3_t::addi, x_register::sp, stack_size));
        instructions.push_back(instruction(base_opcode::jalr)
            .i(x_register::zero, funct3_t::jalr, x_register::ra, 0));
    }

    static void generate_intrinsics(std::shared_ptr<source::writer<std::byte>> writer,
            std::vector<reference>& references)
    {
        writer->sink("printf");
        {
            std::vector<instruction> instructions;
            auto format_string = writer->sink(reinterpret_cast<const std::byte *>("%c\n\0"), 4);

            prologue(instructions);
            instructions.push_back(instruction(base_opcode::opImm)
                    .i(x_register::a1, funct3_t::addi, x_register::zero, 't'));
            instructions.push_back(instruction(base_opcode::branch)
                    .b(8, funct3_t::bne, x_register::zero, x_register::a0));
            instructions.push_back(instruction(base_opcode::opImm)
                    .i(x_register::a1, funct3_t::addi, x_register::zero, 'f'));

            relocate(format_string, address_t::high20, references, instructions, writer);
            instructions.push_back(instruction(base_opcode::lui).u(x_register::a5, 0));
            relocate(format_string, address_t::lower12i, references, instructions, writer);
            instructions.push_back(instruction(base_opcode::opImm)
                    .i(x_register::a0, funct3_t::addi, x_register::a5, 0));

            relocate("printf", address_t::text, references, instructions ,writer);
            instructions.push_back(instruction(base_opcode::auipc).u(x_register::ra, 0));
            instructions.push_back(instruction(base_opcode::jalr)
                    .i(x_register::ra, funct3_t::jalr, x_register::ra, 0));

            epilogue(8, instructions);

            writer->sink("writeb", reinterpret_cast<const std::byte *>(instructions.data()),
                    instructions.size() * sizeof(instruction));
        }
        {
            std::vector<instruction> instructions;
            auto format_string = writer->sink(reinterpret_cast<const std::byte *>("%d\n\0"), 4);

            prologue(instructions);
            instructions.push_back(instruction(base_opcode::opImm)
                    .i(x_register::a1, funct3_t::addi, x_register::a0, 0));

            relocate(format_string, address_t::high20, references, instructions, writer);
            instructions.push_back(instruction(base_opcode::lui).u(x_register::a5, 0));
            relocate(format_string, address_t::lower12i, references, instructions, writer);
            instructions.push_back(instruction(base_opcode::opImm)
                    .i(x_register::a0, funct3_t::addi, x_register::a5, 0));

            relocate("printf", address_t::text, references, instructions, writer);
            instructions.push_back(instruction(base_opcode::auipc).u(x_register::ra, 0));
            instructions.push_back(instruction(base_opcode::jalr)
                    .i(x_register::ra, funct3_t::jalr, x_register::ra, 0));

            epilogue(8, instructions);

            writer->sink("writei", reinterpret_cast<const std::byte *>(instructions.data()),
                    instructions.size() * sizeof(instruction));
        }
    }

    static void load_in_register(std::shared_ptr<source::operand> operand, const x_register target,
            std::shared_ptr<source::procedure_info> procedure_info, std::vector<instruction>& instructions)
    {
        std::shared_ptr<source::integer_operand> integer_operand{ nullptr };
        std::shared_ptr<source::variable_operand> variable_operand{ nullptr };
        std::shared_ptr<source::temporary_variable> temporary_variable{ nullptr };
        std::shared_ptr<source::variable_info> variable_symbol{ nullptr };
        std::shared_ptr<source::parameter_info> parameter_symbol{ nullptr };

        if ((integer_operand = std::dynamic_pointer_cast<source::integer_operand>(operand)) != nullptr)
        {
            instructions.push_back(instruction(base_opcode::opImm)
                    .i(target, funct3_t::addi, x_register::zero, integer_operand->value()));
        }
        else if ((variable_operand = std::dynamic_pointer_cast<source::variable_operand>(operand)) != nullptr)
        {
            const auto& name = procedure_info->scope()->lookup(variable_operand->name());
            if ((variable_symbol = std::dynamic_pointer_cast<source::variable_info>(name)) != nullptr)
            {
                instructions.push_back(instruction(base_opcode::load)
                        .i(target, funct3_t::lw, x_register::s0, variable_symbol->offset));
            }
            else if ((parameter_symbol = std::dynamic_pointer_cast<source::parameter_info>(name)) != nullptr)
            {
                instructions.push_back(instruction(base_opcode::load)
                        .i(target, funct3_t::lw, x_register::s0, parameter_symbol->offset));
            }
        }
        else if ((temporary_variable = std::dynamic_pointer_cast<source::temporary_variable>(operand)) != nullptr)
        {
            instructions.push_back(instruction(base_opcode::load));
            instructions.back().i(target, funct3_t::lw, x_register::s0,
                    procedure_info->local_stack_size + 4 * temporary_variable->counter());
        }
    }

    static void store_from_register(std::shared_ptr<source::operand> destination, const x_register target,
            std::shared_ptr<source::procedure_info> procedure_info, std::vector<instruction>& instructions)
    {
        std::shared_ptr<source::variable_operand> variable_operand{ nullptr };
        std::shared_ptr<source::temporary_variable> temporary_variable{ nullptr };
        std::shared_ptr<source::variable_info> variable_symbol{ nullptr };

        if ((variable_operand = std::dynamic_pointer_cast<source::variable_operand>(destination)) != nullptr)
        {
            variable_symbol = std::dynamic_pointer_cast<source::variable_info>(
                    procedure_info->scope()->lookup(variable_operand->name()));
            instructions.push_back(instruction(base_opcode::store)
                    .s(variable_symbol->offset, funct3_t::sw, x_register::s0, target));
        }
        else if ((temporary_variable = std::dynamic_pointer_cast<source::temporary_variable>(destination)) != nullptr)
        {
            instructions.push_back(instruction(base_opcode::store));
            instructions.back().s(procedure_info->local_stack_size + 4 * temporary_variable->counter(),
                    funct3_t::sw, x_register::s0, target);
        }
    }

    static void perform_binary_operation(const source::binary_operator operation, std::shared_ptr<source::operand> lhs,
            std::shared_ptr<source::operand> rhs, std::shared_ptr<source::operand> destination,
            std::shared_ptr<source::procedure_info> procedure_info, std::vector<instruction>& instructions)
    {
        constexpr auto lhs_register = x_register::a0;
        std::shared_ptr<source::variable_operand> variable_operand{ nullptr };
        std::shared_ptr<source::temporary_variable> temporary_variable{ nullptr };
        std::shared_ptr<source::variable_info> variable_symbol{ nullptr };

        load_in_register(lhs, x_register::a0, procedure_info, instructions);
        load_in_register(rhs, x_register::t0, procedure_info, instructions);

        // Calculate the result and assign it to a variable on the stack.
        switch (operation)
        {
            case source::binary_operator::sum:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::add, x_register::a0, x_register::t0));
                break;
            case source::binary_operator::subtraction:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::sub, x_register::a0, x_register::t0, funct7_t::sub));
                break;
            case source::binary_operator::multiplication:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::mul, x_register::a0, x_register::t0, funct7_t::muldiv));
                break;
            case source::binary_operator::division:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::div, x_register::a0, x_register::t0, funct7_t::muldiv));
                break;
            case source::binary_operator::equals:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::sub, x_register::a0, x_register::t0, funct7_t::sub));
                instructions.push_back(instruction(base_opcode::opImm)
                    .i(lhs_register, funct3_t::sltiu, lhs_register, 1));
                break;
            case source::binary_operator::not_equals:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::sub, x_register::a0, x_register::t0, funct7_t::sub));
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::sltu, x_register::zero, lhs_register));
                break;
            case source::binary_operator::less:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::sltu, x_register::a0, x_register::t0));
                break;
            case source::binary_operator::greater_equal:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::sltu, x_register::t0, x_register::a0));
                break;
            case source::binary_operator::greater:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::slt, x_register::a0, x_register::t0));
                instructions.push_back(instruction(base_opcode::opImm)
                    .i(lhs_register, funct3_t::xori, lhs_register, 1));
                break;
            case source::binary_operator::less_equal:
                instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::slt, x_register::t0, x_register::a0));
                instructions.push_back(instruction(base_opcode::opImm)
                    .i(lhs_register, funct3_t::xori, lhs_register, 1));
                break;
        }
        store_from_register(destination, lhs_register, procedure_info, instructions);
    }

    std::vector<reference> generate(source::intermediate_code_generator generator,
            std::shared_ptr<source::symbol_table> table, std::shared_ptr<source::writer<std::byte>> writer)
    {
        std::vector<reference> references;

        generate_intrinsics(writer, references);

        for (auto& [identifier, intermediate_code] : generator)
        {
            std::vector<instruction> instructions;
            auto main_symbol = std::dynamic_pointer_cast<source::procedure_info>(table->lookup(identifier));
            std::size_t argument_offset{ 0 };
            const auto stack_size = static_cast<std::uint32_t>(
                    intermediate_code.variable_counter() * 4 + 8 + main_symbol->stack_size());
            std::unordered_map<std::size_t, std::function<void(std::size_t)>> missing_labels;
            std::unordered_map<std::size_t, std::function<void(std::size_t)>>::iterator missing_label =
                missing_labels.end();
            std::unordered_map<std::size_t, std::size_t> local_labels;

            for (auto& quadruple : intermediate_code)
            {
                switch (quadruple.operation())
                {
                    case source::quadruple_operator::start:
                        prologue(instructions);
                        break;
                    case source::quadruple_operator::stop:
                        epilogue(stack_size, instructions);
                        break;
                    case source::quadruple_operator::add:
                        perform_binary_operation(source::binary_operator::sum, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::sub:
                        perform_binary_operation(source::binary_operator::subtraction, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::mul:
                        perform_binary_operation(source::binary_operator::multiplication, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::div:
                        perform_binary_operation(source::binary_operator::division, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::eq:
                        perform_binary_operation(source::binary_operator::equals, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::neq:
                        perform_binary_operation(source::binary_operator::not_equals, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::lt:
                        perform_binary_operation(source::binary_operator::less, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::ge:
                        perform_binary_operation(source::binary_operator::greater_equal, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::gt:
                        perform_binary_operation(source::binary_operator::greater, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::le:
                        perform_binary_operation(source::binary_operator::less_equal, quadruple.operand1(),
                                quadruple.operand2(), quadruple.operand3(), main_symbol, instructions);
                        break;
                    case source::quadruple_operator::load:
                        {
                            auto operand_identifier =
                                std::dynamic_pointer_cast<source::variable_operand>(quadruple.operand1())->name();
                            auto variable_symbol = std::dynamic_pointer_cast<source::variable_info>(
                                    main_symbol->scope()->lookup(operand_identifier));

                            load_in_register(quadruple.operand1(), x_register::a0, main_symbol, instructions);
                            instructions.push_back(instruction(base_opcode::load)
                                .i(x_register::a0, funct3_t::lw, x_register::a0, 0));
                            store_from_register(quadruple.operand3(), x_register::a0, main_symbol, instructions);
                        }
                        break;
                    case source::quadruple_operator::ref:
                        {
                            auto operand_identifier =
                                std::dynamic_pointer_cast<source::variable_operand>(quadruple.operand1())->name();
                            auto variable_symbol = std::dynamic_pointer_cast<source::variable_info>(
                                    main_symbol->scope()->lookup(operand_identifier));

                            instructions.push_back(instruction(base_opcode::opImm)
                                .i(x_register::a0, funct3_t::addi, x_register::s0, variable_symbol->offset));
                            store_from_register(quadruple.operand3(), x_register::a0, main_symbol, instructions);
                        }
                        break;
                    case source::quadruple_operator::beqz:
                        load_in_register(quadruple.operand1(), x_register::a0, main_symbol, instructions);
                        instructions.push_back(instruction(base_opcode::branch));
                        missing_labels.emplace(
                                std::dynamic_pointer_cast<source::label_operand>(quadruple.operand3())->counter(),
                                [before_branch = instructions.size() - 1, &instructions](std::size_t target) {
                                    instructions[before_branch].b((target - before_branch) * 4,
                                            funct3_t::beq, x_register::zero, x_register::a0);
                                });
                        break;
                    case source::quadruple_operator::j:
                        {
                            auto local_label = local_labels.find(
                                    std::dynamic_pointer_cast<source::label_operand>(quadruple.operand3())->counter());
                            if (local_label != local_labels.end())
                            {
                                auto offset = -(instructions.size() - local_label->second) * 4;
                                instructions.push_back(instruction(base_opcode::auipc).u(x_register::a0, 0));
                                instructions.push_back(instruction(base_opcode::jalr)
                                        .i(x_register::zero, funct3_t::jalr, x_register::a0, offset));
                            }
                        }
                        break;
                    case source::quadruple_operator::label:
                        {
                            auto label_counter =
                                std::dynamic_pointer_cast<source::label_operand>(quadruple.operand3())->counter();
                            missing_label = missing_labels.find(label_counter);

                            if (missing_label != missing_labels.end())
                            {
                                missing_label->second(instructions.size());
                            }
                            local_labels[label_counter] = instructions.size();
                        }
                        break;
                    case source::quadruple_operator::assign:
                        load_in_register(quadruple.operand1(), x_register::a0, main_symbol, instructions);
                        store_from_register(quadruple.operand3(), x_register::a0, main_symbol, instructions);
                        break;
                    case source::quadruple_operator::param:
                        load_in_register(quadruple.operand1(), x_register::a0, main_symbol, instructions);
                        instructions.push_back(instruction(base_opcode::store)
                                .s(argument_offset, funct3_t::sw, x_register::sp, x_register::a0));
                        argument_offset += 4;
                        break;
                    case source::quadruple_operator::call:
                        relocate(std::dynamic_pointer_cast<source::variable_operand>(quadruple.operand1())->name(),
                                address_t::text, references, instructions, writer);
                        instructions.push_back(instruction(base_opcode::auipc).u(x_register::ra, 0));
                        instructions.push_back(instruction(base_opcode::jalr)
                            .i(x_register::ra, funct3_t::jalr, x_register::ra, 0));
                        argument_offset = 0;
                        break;
                }
            }
            writer->sink(identifier,
                    reinterpret_cast<const std::byte *>(instructions.data()),
                    instructions.size() * sizeof(instruction));
        }
        return references;
    }
}