#include "elna/riscv.hpp"
#include <memory>
#include <cstring>

namespace elna
{
    Instruction::Instruction(BaseOpcode opcode)
    {
        this->instruction = static_cast<std::underlying_type<BaseOpcode>::type>(opcode);
    }

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

        return *this;
    }

    Instruction& Instruction::s(std::uint32_t imm1, Funct3 funct3, XRegister rs1, XRegister rs2)
    {
        this->instruction |= ((imm1 & 0b11111) << 7)
            | (static_cast<std::underlying_type<Funct3>::type>(funct3) << 12)
            | (static_cast<std::underlying_type<XRegister>::type>(rs1) << 15)
            | (static_cast<std::underlying_type<XRegister>::type>(rs2) << 20)
            | ((imm1 & 0b111111100000) << 20);

        return *this;
    }

    Instruction& Instruction::r(XRegister rd, Funct3 funct3, XRegister rs1, XRegister rs2, Funct7 funct7)
    {
        this->instruction |= (static_cast<std::underlying_type<XRegister>::type>(rd) << 7)
            | (static_cast<std::underlying_type<Funct3>::type>(funct3) << 12)
            | (static_cast<std::underlying_type<XRegister>::type>(rs1) << 15)
            | (static_cast<std::underlying_type<XRegister>::type>(rs2) << 20)
            | (static_cast<std::underlying_type<Funct7>::type>(funct7) << 25);

        return *this;
    }

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

        return *this;
    }

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

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

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

    void RiscVVisitor::visit(source::block *block)
    {
        for (const auto& block_definition : block->definitions())
        {
            block_definition->accept(this);
        }
        this->instructionsLength += 4;
        this->instructions = reinterpret_cast<Instruction *>(
                realloc(this->instructions, this->instructionsLength * sizeof(Instruction)));

        block->body().accept(this);

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

        this->instructions[0] = Instruction(BaseOpcode::opImm)
            .i(XRegister::sp, Funct3::addi, XRegister::sp, -stackSize);
        this->instructions[1] = Instruction(BaseOpcode::store)
            .s(stackSize - 4, Funct3::sw, XRegister::sp, XRegister::s0);
        this->instructions[2] = Instruction(BaseOpcode::store)
            .s(stackSize - 8, Funct3::sw, XRegister::sp, XRegister::ra);
        this->instructions[3] = Instruction(BaseOpcode::opImm)
            .i(XRegister::s0, Funct3::addi, XRegister::sp, stackSize);

        this->instructions = reinterpret_cast<Instruction*>(
            realloc(this->instructions, (this->instructionsLength + 10) * sizeof(Instruction)));

        // Print the result.
        this->instructions[instructionsLength++] = Instruction(BaseOpcode::opImm)
            .i(XRegister::a1, Funct3::addi, XRegister::a0, 0);
        this->references[0] = Reference();
        this->references[0].name = ".CL0";
        this->references[0].offset = instructionsLength * 4;
        this->references[0].target = Target::high20;
        this->instructions[instructionsLength++] = Instruction(BaseOpcode::lui).u(XRegister::a5, 0);
        this->references[1] = Reference();
        this->references[1].name = ".CL0";
        this->references[1].offset = instructionsLength * 4;
        this->references[1].target = Target::lower12i;

        this->instructions[instructionsLength++] = Instruction(BaseOpcode::opImm)
            .i(XRegister::a0, Funct3::addi, XRegister::a5, 0);
        this->references[2] = Reference();
        this->references[2].name = "printf";
        this->references[2].offset = instructionsLength * 4;
        this->references[2].target = Target::text;
        this->instructions[instructionsLength++] = Instruction(BaseOpcode::auipc).u(XRegister::ra, 0);
        this->instructions[instructionsLength++] = Instruction(BaseOpcode::jalr)
            .i(XRegister::ra, Funct3::jalr, XRegister::ra, 0);
        // Set the return value (0).
        this->instructions[instructionsLength++] = Instruction(BaseOpcode::op)
            .r(XRegister::a0, Funct3::_and, XRegister::zero, XRegister::zero);

        // Epilogue.
        this->instructions[instructionsLength++] = Instruction(BaseOpcode::load)
            .i(XRegister::s0, Funct3::lw, XRegister::sp, stackSize - 4);
        this->instructions[instructionsLength++] = Instruction(BaseOpcode::load)
            .i(XRegister::ra, Funct3::lw, XRegister::sp, stackSize - 8);
        this->instructions[instructionsLength++] = Instruction(BaseOpcode::opImm)
                .i(XRegister::sp, Funct3::addi, XRegister::sp, stackSize);
        this->instructions[instructionsLength++] = Instruction(BaseOpcode::jalr)
            .i(XRegister::zero, Funct3::jalr, XRegister::ra, 0);
    }

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

    void RiscVVisitor::visit(source::variable *variable)
    {
        const auto freeRegister = this->registerInUse ? XRegister::a0 : XRegister::t0;

        ++this->instructionsLength;
        this->instructions = reinterpret_cast<Instruction *>(
                realloc(this->instructions, this->instructionsLength * sizeof(Instruction)));
        this->instructions[this->instructionsLength - 1] =
            Instruction(BaseOpcode::opImm) // movl $x, %eax; where $x is a number.
                .i(freeRegister, Funct3::addi, XRegister::zero, constants[variable->name()]);
    }

    void RiscVVisitor::visit(source::integer_literal *number)
    {
        const auto freeRegister = this->registerInUse ? XRegister::a0 : XRegister::t0;

        ++this->instructionsLength;
        this->instructions = reinterpret_cast<Instruction *>(
                realloc(this->instructions, this->instructionsLength * sizeof(Instruction)));
        this->instructions[this->instructionsLength - 1] =
            Instruction(BaseOpcode::opImm) // movl $x, %eax; where $x is a number.
                .i(freeRegister, Funct3::addi, XRegister::zero, number->number());
    }

    void RiscVVisitor::visit(source::binary_expression *expression)
    {
        const auto lhs_register = this->registerInUse ? XRegister::a0 : XRegister::t0;

        this->registerInUse = true;
        expression->lhs().accept(this);

        ++this->instructionsLength;
        this->instructions = reinterpret_cast<Instruction *>(
                realloc(this->instructions, this->instructionsLength * sizeof(Instruction)));
        this->instructions[instructionsLength - 1] = // movl %eax, -x(%rbp); where x is a number.
            Instruction(BaseOpcode::store)
                .s(static_cast<std::uint32_t>(this->variableCounter * 4), Funct3::sw, XRegister::sp, XRegister::a0);
        auto lhs_stack_position = ++this->variableCounter;

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

        this->instructionsLength += 2;
        this->instructions = reinterpret_cast<Instruction *>(
                realloc(this->instructions, this->instructionsLength * sizeof(Instruction)));

         this->instructions[instructionsLength - 2] = Instruction(BaseOpcode::load)
            .i(XRegister::a0, Funct3::lw, XRegister::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[instructionsLength - 1] = Instruction(BaseOpcode::op)
                    .r(lhs_register, Funct3::add, XRegister::a0, XRegister::t0);
                break;
            case source::binary_operator::subtraction:
                this->instructions[instructionsLength - 1] = Instruction(BaseOpcode::op)
                    .r(lhs_register, Funct3::sub, XRegister::a0, XRegister::t0, Funct7::sub);
                break;
            case source::binary_operator::multiplication:
                this->instructions[instructionsLength - 1] = Instruction(BaseOpcode::op)
                    .r(lhs_register, Funct3::mul, XRegister::a0, XRegister::t0, Funct7::muldiv);
                break;
        }
    }

    Symbol writeNext(source::block *ast)
    {
        auto visitor = std::make_unique<RiscVVisitor>();
        visitor->visit(ast);

        Symbol program{ "main" };

        for (std::size_t i = 0; i < 3; ++i)
        {
            program.symbols[i] = visitor->references[i];
        }
        for (std::size_t i = 0; i < visitor->instructionsLength; ++i)
        {
            program.text.insert(program.text.cend(),
                    visitor->instructions[i].cbegin(), visitor->instructions[i].cend());
        }
        return program;
    }
}