#include "elna/backend/riscv.hpp"
#include <cassert>
#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);
    }

    visitor::visitor(std::shared_ptr<source::writer> writer,
            std::shared_ptr<source::symbol_table> table)
        : writer(writer), table(table)
    {
    }

    void visitor::generate_intrinsics()
    {
        this->writer->sink("printf");
        {
            auto format_string = this->writer->sink(reinterpret_cast<const std::byte *>("%c\n\0"), 4);

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

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

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

            epilogue(8);

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

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

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

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

            epilogue(8);

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

    void visitor::relocate(std::string_view name, address_t target)
    {
        this->references.push_back(reference());
        this->references.back().name = name;
        this->references.back().offset = writer->size() + instructions.size() * 4;
        this->references.back().target = target;
    }

    void visitor::visit(source::declaration *declaration)
    {
    }

    void visitor::visit(source::constant_definition *definition)
    {
    }

    void visitor::prologue()
    {
        this->variable_counter = 1;

        this->instructions.push_back(instruction(base_opcode::opImm));
        this->instructions.push_back(instruction(base_opcode::store));
        this->instructions.push_back(instruction(base_opcode::store));
        this->instructions.push_back(instruction(base_opcode::opImm));
   }

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

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

    void visitor::visit(source::procedure_definition *definition)
    {
        prologue();

        auto main_symbol =
            std::dynamic_pointer_cast<source::procedure_info>(this->table->lookup(definition->identifier()));
        this->table = main_symbol->scope();
        definition->body().accept(this);
        this->table = main_symbol->scope()->scope();

        // Set the return value (0).
        this->instructions.push_back(instruction(base_opcode::op)
            .r(x_register::a0, funct3_t::_and, x_register::zero, x_register::zero));

        epilogue(static_cast<std::uint32_t>(this->variable_counter * 4 + 8 + main_symbol->stack_size()));
        this->writer->sink(definition->identifier(),
                reinterpret_cast<const std::byte *>(this->instructions.data()),
                this->instructions.size() * sizeof(instruction));
        this->instructions.clear();
    }

    void visitor::visit(source::block *block)
    {
        block->body().accept(this);
    }

    void visitor::visit(source::program *program)
    {
        generate_intrinsics();
        for (auto& definition : program->definitions())
        {
            definition->accept(this);
        }
        prologue();

        auto main_symbol =
            std::dynamic_pointer_cast<source::procedure_info>(this->table->lookup("main"));
        program->body().accept(this);

        // Set the return value (0).
        this->instructions.push_back(instruction(base_opcode::op)
            .r(x_register::a0, funct3_t::_and, x_register::zero, x_register::zero));

        epilogue(static_cast<std::uint32_t>(this->variable_counter * 4 + 8 + main_symbol->local_stack_size));
        this->writer->sink("main", reinterpret_cast<const std::byte *>(this->instructions.data()),
                this->instructions.size() * sizeof(instruction));
    }

    void visitor::visit(source::call_statement *statement)
    {
        std::size_t argument_offset{ 0 };

        for (auto& argument : statement->arguments())
        {
            argument->accept(this);
            const auto free_register = this->register_in_use ? x_register::a0 : x_register::t0;
            this->instructions.push_back(instruction(base_opcode::store)
                    .s(argument_offset, funct3_t::sw, x_register::sp, free_register));
            argument_offset += 4;
        }
        relocate(statement->name(), address_t::text);
        this->instructions.push_back(instruction(base_opcode::auipc).u(x_register::ra, 0));
        this->instructions.push_back(instruction(base_opcode::jalr)
            .i(x_register::ra, funct3_t::jalr, x_register::ra, 0));
    }

    void visitor::visit(source::compound_statement *statement)
    {
        for (auto& nested_statement : statement->statements())
        {
            nested_statement->accept(this);
        }
    }

    void visitor::visit(source::assign_statement *statement)
    {
        const auto free_register = this->register_in_use ? x_register::a0 : x_register::t0;
        auto symbol = table->lookup(statement->lvalue());
        auto variable_symbol = std::dynamic_pointer_cast<source::variable_info>(symbol);

        statement->rvalue().accept(this);

        this->instructions.push_back(instruction(base_opcode::store)
            .s(variable_symbol->offset, funct3_t::sw, x_register::s0, x_register::a0));
    }

    void visitor::visit(source::if_statement *statement)
    {
        statement->prerequisite().accept(this);

        const auto free_register = this->register_in_use ? x_register::a0 : x_register::t0;

        auto before_branch = instructions.size();
        instructions.push_back(instruction(base_opcode::branch));
        statement->body().accept(this);
        instructions[before_branch]
            .b((instructions.size() - before_branch) * 4 - 4, funct3_t::beq, x_register::zero, free_register);
    }

    void visitor::visit(source::while_statement *statement)
    {
        statement->prerequisite().accept(this);
        statement->body().accept(this);
    }

    void visitor::visit(source::type_expression *type)
    {
    }

    void visitor::visit(source::variable_expression *variable)
    {
        const auto free_register = this->register_in_use ? x_register::a0 : x_register::t0;

        auto symbol = table->lookup(variable->name());
        if (auto constant_symbol = std::dynamic_pointer_cast<source::constant_info>(symbol))
        {
            this->instructions.push_back(
                instruction(base_opcode::opImm) // movl $x, %eax; where $x is a number.
                    .i(free_register, funct3_t::addi, x_register::zero, constant_symbol->value())
            );
        }
        else if (auto variable_symbol = std::dynamic_pointer_cast<source::variable_info>(symbol))
        {
            this->instructions.push_back(
                instruction(base_opcode::load)
                    .i(free_register, funct3_t::lw, x_register::s0, variable_symbol->offset)
            );
        }
        else if (auto parameter_symbol = std::dynamic_pointer_cast<source::parameter_info>(symbol))
        {
            this->instructions.push_back(
                instruction(base_opcode::load)
                    .i(free_register, funct3_t::lw, x_register::s0, parameter_symbol->offset)
            );
        }
    }

    void visitor::visit(source::binary_expression *expression)
    {
        const auto lhs_register = this->register_in_use ? x_register::a0 : x_register::t0;

        this->register_in_use = true;
        expression->lhs().accept(this);
        auto lhs_stack_position = this->variable_counter * 4;
        ++this->variable_counter;

        this->instructions.push_back(
            instruction(base_opcode::store)
                .s(static_cast<std::uint32_t>(lhs_stack_position), funct3_t::sw, x_register::sp, x_register::a0)
        );

        this->register_in_use = false;
        expression->rhs().accept(this);
        this->register_in_use = lhs_register == x_register::a0; // Restore.

         this->instructions.push_back(instruction(base_opcode::load)
            .i(x_register::a0, funct3_t::lw, x_register::sp,
                static_cast<std::int8_t>(lhs_stack_position))
        );

        // Calculate the result and assign it to a variable on the stack.
        switch (expression->operation())
        {
            case source::binary_operator::sum:
                this->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:
                this->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:
                this->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:
                this->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:
                this->instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::sub, x_register::a0, x_register::t0, funct7_t::sub));
                this->instructions.push_back(instruction(base_opcode::opImm)
                    .i(lhs_register, funct3_t::sltiu, lhs_register, 1));
                break;
            case source::binary_operator::not_equals:
                this->instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::sub, x_register::a0, x_register::t0, funct7_t::sub));
                this->instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::sltu, x_register::zero, lhs_register));
                break;
            case source::binary_operator::less:
                this->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:
                this->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:
                this->instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::slt, x_register::a0, x_register::t0));
                this->instructions.push_back(instruction(base_opcode::opImm)
                    .i(lhs_register, funct3_t::xori, lhs_register, 1));
                break;
            case source::binary_operator::less_equal:
                this->instructions.push_back(instruction(base_opcode::op)
                    .r(lhs_register, funct3_t::slt, x_register::t0, x_register::a0));
                this->instructions.push_back(instruction(base_opcode::opImm)
                    .i(lhs_register, funct3_t::xori, lhs_register, 1));
                break;
        }
    }

    void visitor::visit(source::unary_expression *expression)
    {
        const auto free_register = this->register_in_use ? x_register::a0 : x_register::t0;
        auto operand_identifier = dynamic_cast<source::variable_expression&>(expression->operand()).name();
        auto variable_symbol =
            std::dynamic_pointer_cast<source::variable_info>(this->table->lookup(operand_identifier));

        this->instructions.push_back(instruction(base_opcode::opImm)
            .i(free_register, funct3_t::addi, x_register::s0, variable_symbol->offset));
    }

    void visitor::visit(source::integer_literal *number)
    {
        const auto free_register = this->register_in_use ? x_register::a0 : x_register::t0;

        this->instructions.push_back(
            instruction(base_opcode::opImm) // movl $x, %eax; where $x is a number.
                .i(free_register, funct3_t::addi, x_register::zero, number->number())
        );
    }

    void visitor::visit(source::boolean_literal *number)
    {
        const auto free_register = this->register_in_use ? x_register::a0 : x_register::t0;

        this->instructions.push_back(
            instruction(base_opcode::opImm) // movl $x, %eax; where $x is a number.
                .i(free_register, funct3_t::addi, x_register::zero, number->boolean())
        );
    }
}