#include "elna/backend/target.hpp"
#include "elna/backend/riscv.hpp"
#include <cstring>
#include <boost/process/v2.hpp>
#include <boost/endian.hpp>

namespace elna::riscv
{
    elfio_writer::elfio_writer(ELFIO::section *text,
            ELFIO::symbol_section_accessor symbol_accessor, ELFIO::string_section_accessor string_accessor)
        : text(text),
        symbol_accessor(symbol_accessor), string_accessor(string_accessor)
    {
    }

    std::size_t elfio_writer::sink(const std::string& label, const std::byte *data, std::size_t size)
    {
        auto offset = text->get_size();
        text->append_data(reinterpret_cast<const char *>(data), size);

        this->symbol_accessor.add_symbol(this->string_accessor, label.data(), offset, label.size(),
                ELFIO::STB_GLOBAL, ELFIO::STT_FUNC, 0, text->get_index());

        return text->get_size();
    }

    void elfio_writer::sink(const std::string& label)
    {
        this->symbol_accessor.add_symbol(this->string_accessor, "printf", 0x00000000, 0,
                ELFIO::STB_GLOBAL, ELFIO::STT_NOTYPE, 0, ELFIO::SHN_UNDEF);
    }

    std::size_t elfio_writer::size() const
    {
        return this->text->get_size();
    }

    process_writer::process_writer(boost::asio::io_context& context)
        : pipe(context)
    {
    }

    std::size_t process_writer::sink(const std::string& label, const std::byte *data, std::size_t size)
    {
        auto network_order_size = boost::endian::native_to_big(static_cast<std::uint16_t>(size));

        auto written = boost::asio::write(this->pipe, boost::asio::buffer(label.data(), label.size() + 1));
        boost::asio::write(this->pipe, boost::asio::buffer(&network_order_size, 2));
        boost::asio::write(this->pipe, boost::asio::buffer(data, size));
        this->text_size += written;

        return written;
    }

    void process_writer::sink(const std::string& label)
    {
        std::uint16_t data_size{ 0 };
        boost::asio::write(this->pipe, boost::asio::buffer(label.data(), label.size() + 1));
        boost::asio::write(this->pipe, boost::asio::buffer(&data_size, 2));
    }

    std::size_t process_writer::size() const
    {
        return this->text_size;
    }

    std::ptrdiff_t lookup(ELFIO::symbol_section_accessor symbol_accessor, const std::string& label)
    {
        for (ptrdiff_t j = 0; j < symbol_accessor.get_symbols_num(); ++j)
        {
            std::string name;
            ELFIO::Elf64_Addr value;
            ELFIO::Elf_Xword size;
            unsigned char bind;
            unsigned char type;
            ELFIO::Elf_Half section_index;
            unsigned char other;

            symbol_accessor.get_symbol(j, name, value, size, bind, type, section_index, other);
            if (name == label)
            {
                return j;
            }
        }
        return -1;
    }

    void riscv32_elf(source::program *ast, source::intermediate_code_generator intermediate_code_generator,
            std::shared_ptr<source::symbol_table> table, 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_flags(0x4); // EF_RISCV_FLOAT_ABI_DOUBLE
        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 symbol relocation table writers
        ELFIO::symbol_section_accessor syma(writer, sym_sec);
        ELFIO::relocation_section_accessor rela(writer, rel_sec);

        auto writer1 = std::make_shared<elfio_writer>(text_sec, syma, stra);
        boost::asio::io_context ctx;
        auto writer2 = std::make_shared<process_writer>(ctx);
        boost::process::process child(ctx, "./build/hs/elna", {},
                boost::process::process_stdio{ { writer2->pipe }, nullptr, nullptr });

        auto references = generate(intermediate_code_generator, table, writer1);
        generate(intermediate_code_generator, table, writer2);

        syma.arrange_local_symbols();

        for (auto& reference : references)
        {
            ELFIO::Elf_Word relocated_symbol = lookup(syma, reference.name);
            auto reference_offset = boost::endian::native_to_big(static_cast<std::uint32_t>(reference.offset));
            std::uint8_t address_type{ 0 };

            switch (reference.target)
            {
                case address_t::high20:
                    address_type = 26;
                    break;
                case address_t::lower12i:
                    address_type = 27;
                    break;
                case address_t::text:
                    address_type = 18;
                    break;
            }
            rela.add_entry(reference.offset, relocated_symbol, 18 /* ELFIO::R_RISCV_CALL */);
            // rela.add_entry(reference.offset, relocated_symbol, 51 /* ELFIO::R_RISCV_RELAX */);
            writer2->pipe.write_some(boost::asio::buffer(&address_type, 1));
            writer2->pipe.write_some(boost::asio::buffer(reference.name.data(), reference.name.size() + 1));
            writer2->pipe.write_some(boost::asio::buffer(&reference_offset, 4));
        }

        // Create ELF object file
        writer2->pipe.close();
        child.wait();

        writer.save(out_file);
    }
}