#include "elna/backend/target.hpp"
#include "elna/backend/riscv.hpp"
#include <cstring>

namespace elna::riscv
{
    elfio_writer::iterator::reference elfio_writer::iterator::operator*() const noexcept
    {
        return payload;
    }

    elfio_writer::iterator::pointer elfio_writer::iterator::operator->() const noexcept
    {
        return &payload;
    }

    elfio_writer::iterator& elfio_writer::iterator::operator++()
    {
        this->payload.data += *this->sizes;
        this->payload.label = *(++this->labels);
        this->payload.size = *(++this->sizes);

        return *this;
    }

    elfio_writer::iterator& elfio_writer::iterator::operator++(int)
    {
        auto tmp = *this;
        ++(*this);
        return *this;
    }

    bool elfio_writer::iterator::operator==(const iterator& that) const
    {
        return this->labels == that.labels;
    }

    bool elfio_writer::iterator::operator!=(const iterator& that) const
    {
        return !(*this == that);
    }

    elfio_writer::elfio_writer(ELFIO::section *text)
        : text(text), labels(std::make_shared<std::vector<std::string>>()),
        sizes(std::make_shared<std::vector<std::size_t>>())
    {
    }

    void elfio_writer::operator()(const std::string& label, const std::byte *data, std::size_t size)
    {
        labels->push_back(label + '\0');
        sizes->push_back(size);
        text->append_data(reinterpret_cast<const char *>(data), size);
    }

    std::string_view elfio_writer::operator()(const std::byte *data, std::size_t size)
    {
        auto found = std::find_if(begin(), end(),
                [data, size](elfio_writer::entry entry) {
                    return size == entry.size && std::memcmp(entry.data, data, size) == 0;
                });
        if (found == end())
        {
            (*this)(".CL" + std::to_string(labels->size()), data, size);
            return labels->back();
        }
        return found->label;
    }

    elfio_writer::iterator elfio_writer::begin() const
    {
        return elfio_writer::iterator(labels->cbegin(), sizes->cbegin(),
                reinterpret_cast<const std::byte *>(text->get_data()));
    }

    elfio_writer::iterator elfio_writer::end() const
    {
        return elfio_writer::iterator(labels->cend(), sizes->cend());
    }

    std::ptrdiff_t elfio_writer::lookup(const std::string& label)
    {
        auto found = std::find(labels->cbegin(), labels->cend(), label);

        if (found == labels->cend())
        {
            return -1;
        }
        return std::distance(labels->cbegin(), found);
    }

    void riscv32_elf(source::program *ast, const std::filesystem::path& out_file)
    {
        ELFIO::elfio writer;

        writer.create(ELFIO::ELFCLASS32, ELFIO::ELFDATA2LSB);

        writer.set_os_abi(ELFIO::ELFOSABI_NONE);
        writer.set_type(ELFIO::ET_REL);
        writer.set_machine(ELFIO::EM_RISCV);

        // Create code section
        ELFIO::section* text_sec = writer.sections.add(".text");
        text_sec->set_type(ELFIO::SHT_PROGBITS);
        text_sec->set_flags(ELFIO::SHF_ALLOC | ELFIO::SHF_EXECINSTR);
        text_sec->set_addr_align(0x1);

        // Create string table section
        ELFIO::section* str_sec = writer.sections.add(".strtab");
        str_sec->set_type(ELFIO::SHT_STRTAB);

        // Create string table writer
        ELFIO::string_section_accessor stra(str_sec);

        // Create symbol table section
        ELFIO::section* sym_sec = writer.sections.add(".symtab");
        sym_sec->set_type(ELFIO::SHT_SYMTAB);
        sym_sec->set_info(2);
        sym_sec->set_addr_align(0x4);
        sym_sec->set_entry_size(writer.get_default_entry_size(ELFIO::SHT_SYMTAB));
        sym_sec->set_link(str_sec->get_index());

        // Create relocation table section
        ELFIO::section* rel_sec = writer.sections.add(".rel.text");
        rel_sec->set_type(ELFIO::SHT_REL);
        rel_sec->set_info(text_sec->get_index());
        rel_sec->set_addr_align(0x4);
        rel_sec->set_entry_size(writer.get_default_entry_size(ELFIO::SHT_REL));
        rel_sec->set_link(sym_sec->get_index());
        rel_sec->set_flags(ELFIO::SHF_ALLOC);

        // Create read only data section
        ELFIO::section* ro_sec = writer.sections.add(".rodata");
        ro_sec->set_type(ELFIO::SHT_PROGBITS);
        ro_sec->set_flags(ELFIO::SHF_ALLOC);
        ro_sec->set_addr_align(0x4);

        elfio_writer text_writer{ text_sec };
        elfio_writer read_only_writer{ ro_sec };
        visitor _visitor{ text_writer, read_only_writer };
        _visitor.visit(ast);

        // Create symbol relocation table writers
        ELFIO::symbol_section_accessor syma(writer, sym_sec);
        ELFIO::relocation_section_accessor rela(writer, rel_sec);
        ELFIO::Elf_Word digit_symbol;

        for (auto symbol : read_only_writer)
        {
            syma.add_symbol(stra, symbol.label.data(), 0x00000000,
                symbol.label.size(), ELFIO::STB_LOCAL, ELFIO::STT_NOTYPE, 0, ro_sec->get_index());
        }
        ELFIO::Elf_Word printf_symbol = syma.add_symbol(stra, "printf", 0x00000000, 0,
                ELFIO::STB_GLOBAL, ELFIO::STT_NOTYPE, 0, ELFIO::SHN_UNDEF);

        for (auto symbol : text_writer)
        {
            syma.add_symbol(stra, symbol.label.data(), 0x00000000, symbol.label.size(),
                    ELFIO::STB_GLOBAL, ELFIO::STT_FUNC, 0, text_sec->get_index());
        }
        for (auto& reference : _visitor.references)
        {
            switch (reference.target)
            {
                case address_t::high20:
                    digit_symbol = read_only_writer.lookup(reference.name) + 1;
                    rela.add_entry(reference.offset, digit_symbol, 26 /* ELFIO::R_RISCV_HI20 */);
                    rela.add_entry(reference.offset, digit_symbol, 51 /* ELFIO::R_RISCV_RELAX */);
                    break;
                case address_t::lower12i:
                    digit_symbol = read_only_writer.lookup(reference.name) + 1;
                    rela.add_entry(reference.offset, digit_symbol, 27 /* ELFIO::R_RISCV_LO12_I */);
                    rela.add_entry(reference.offset, digit_symbol, 51 /* ELFIO::R_RISCV_RELAX */);
                    break;
                case address_t::text:
                    rela.add_entry(reference.offset, printf_symbol, 18 /* ELFIO::R_RISCV_CALL */);
                    rela.add_entry(reference.offset, printf_symbol, 51 /* ELFIO::R_RISCV_RELAX */);
                    break;
            }
        }

        // Create ELF object file
        writer.save(out_file);
    }
}