#include "elna/backend/riscv.hpp"
#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 imm1, funct3_t funct3, x_register rs1, x_register rs2)
    {
        this->representation |= ((imm1 & 0b11111) << 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)
            | ((imm1 & 0b111111100000) << 20);

        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;
    }

    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);
    }

    void visitor::visit(source::definition *definition)
    {
        constants[definition->identifier()] = definition->body().number();
    }

    void visitor::visit(source::block *block)
    {
        for (const auto& block_definition : block->definitions())
        {
            block_definition->accept(this);
        }
        block->body().accept(this);

        // Prologue.
        const uint stackSize = static_cast<std::uint32_t>(variable_counter * 4 + 12);

        std::vector<instruction> prologue{
            instruction(base_opcode::opImm)
                .i(x_register::sp, funct3_t::addi, x_register::sp, -stackSize),
            instruction(base_opcode::store)
                .s(stackSize - 4, funct3_t::sw, x_register::sp, x_register::s0),
            instruction(base_opcode::store)
                .s(stackSize - 8, funct3_t::sw, x_register::sp, x_register::ra),
            instruction(base_opcode::opImm)
                .i(x_register::s0, funct3_t::addi, x_register::sp, stackSize)
        };
        this->instructions.insert(this->instructions.cbegin(), prologue.begin(), prologue.end());

        // Print the result.
        this->instructions.push_back(instruction(base_opcode::opImm)
            .i(x_register::a1, funct3_t::addi, x_register::a0, 0));
        this->references[0] = reference();
        this->references[0].name = ".CL0";
        this->references[0].offset = instructions.size() * 4;
        this->references[0].target = address_t::high20;
        this->instructions.push_back(instruction(base_opcode::lui).u(x_register::a5, 0));
        this->references[1] = reference();
        this->references[1].name = ".CL0";
        this->references[1].offset = instructions.size() * 4;
        this->references[1].target = address_t::lower12i;

        this->instructions.push_back(instruction(base_opcode::opImm)
            .i(x_register::a0, funct3_t::addi, x_register::a5, 0));
        this->references[2] = reference();
        this->references[2].name = "printf";
        this->references[2].offset = instructions.size() * 4;
        this->references[2].target = 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));
        // 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.
        this->instructions.push_back(instruction(base_opcode::load)
            .i(x_register::s0, funct3_t::lw, x_register::sp, stackSize - 4));
        this->instructions.push_back(instruction(base_opcode::load)
            .i(x_register::ra, funct3_t::lw, x_register::sp, stackSize - 8));
        this->instructions.push_back(instruction(base_opcode::opImm)
                .i(x_register::sp, funct3_t::addi, x_register::sp, stackSize));
        this->instructions.push_back(instruction(base_opcode::jalr)
            .i(x_register::zero, funct3_t::jalr, x_register::ra, 0));
    }

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

    void visitor::visit(source::variable_expression *variable)
    {
        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, constants[variable->name()])
        );
    }

    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::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);

        this->instructions.push_back( // movl %eax, -x(%rbp); where x is a number.
            instruction(base_opcode::store)
                .s(static_cast<std::uint32_t>(this->variable_counter * 4), funct3_t::sw, x_register::sp, x_register::a0)
        );
        auto lhs_stack_position = ++this->variable_counter;

        this->register_in_use = false;
        expression->rhs().accept(this);

         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 - 1) * 4))
        );

        // 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;
        }
    }
}