#include "loader/elf.hpp"
#include <fstream>
#include <algorithm>

static const uint64_t EI_NIDENT = 16;
static const uint32_t PT_LOAD = 1;
static const uint32_t PT_DYNAMIC = 2;
static const uint32_t DT_NULL = 0;
static const uint32_t DT_NEEDED = 1;
static const uint32_t DT_INIT = 12;
static const uint32_t DT_FINI = 13;
static const uint32_t DT_INIT_ARRAY = 25;
static const uint32_t DT_FINI_ARRAY = 26;
static const uint32_t DT_INIT_ARRAYSZ = 27;
static const uint32_t DT_FINI_ARRAYSZ = 28;
static const uint32_t DT_STRTAB = 5;
static const uint32_t DT_SYMTAB = 6;
static const uint32_t DT_STRSZ = 10;
static const uint32_t SHT_SYMTAB = 2;
static const uint32_t SHT_DYNSYM = 11;
static const uint32_t SHT_RELA = 4;
static const uint32_t SHT_REL = 9;

ELFLoader::ELFLoader()
    : m_arch(Arch::X64), m_entry(0),
      m_init_array(0), m_init_array_size(0),
      m_fini_array(0), m_fini_array_size(0),
      m_init(0), m_fini(0) {}

ELFLoader::~ELFLoader() {}

bool ELFLoader::load(const string& path)
{
    m_path = path;
    ifstream f(path, ios::binary);
    if (!f) return false;

    f.seekg(0, ios::end);
    size_t sz = f.tellg();
    f.seekg(0, ios::beg);

    m_data.resize(sz);
    f.read((char*)m_data.data(), sz);
    f.close();

    if (m_data.size() < EI_NIDENT) return false;
    if (m_data[0] != 0x7F || m_data[1] != 'E' || m_data[2] != 'L' || m_data[3] != 'F')
        return false;

    return parseHeader();
}

static uint16_t r16(const vector<uint8_t>& d, size_t off)
{
    if (off + 2 > d.size()) return 0;
    return d[off] | (d[off+1] << 8);
}

static uint32_t r32(const vector<uint8_t>& d, size_t off)
{
    if (off + 4 > d.size()) return 0;
    return d[off] | (d[off+1] << 8) | (d[off+2] << 16) | (d[off+3] << 24);
}

static uint64_t r64(const vector<uint8_t>& d, size_t off)
{
    if (off + 8 > d.size()) return 0;
    return (uint64_t)d[off] | ((uint64_t)d[off+1] << 8) |
           ((uint64_t)d[off+2] << 16) | ((uint64_t)d[off+3] << 24) |
           ((uint64_t)d[off+4] << 32) | ((uint64_t)d[off+5] << 40) |
           ((uint64_t)d[off+6] << 48) | ((uint64_t)d[off+7] << 56);
}

struct Elf64_Ehdr
{
    uint8_t ident[16];
    uint16_t type;
    uint16_t machine;
    uint32_t version;
    uint64_t entry;
    uint64_t phoff;
    uint64_t shoff;
    uint32_t flags;
    uint16_t ehsize;
    uint16_t phentsize;
    uint16_t phnum;
    uint16_t shentsize;
    uint16_t shnum;
    uint16_t shstrndx;
};

struct Elf64_Phdr
{
    uint32_t type;
    uint32_t flags;
    uint64_t offset;
    uint64_t vaddr;
    uint64_t paddr;
    uint64_t filesz;
    uint64_t memsz;
    uint64_t align;
};

struct Elf64_Shdr
{
    uint32_t name;
    uint32_t type;
    uint64_t flags;
    uint64_t addr;
    uint64_t offset;
    uint64_t size;
    uint32_t link;
    uint32_t info;
    uint64_t addralign;
    uint64_t entsize;
};

struct Elf64_Dyn
{
    int64_t d_tag;
    uint64_t d_val;
};

bool ELFLoader::parseHeader()
{
    if (m_data.size() < 64) return false;

    uint8_t cls = m_data[4];
    bool is_64 = (cls == 2);

    if (is_64)
    {
        m_arch = Arch::X64;
        if (m_data.size() < 64) return false;
        uint16_t type = r16(m_data, 16);
        m_entry = r64(m_data, 24);
        uint64_t phoff = r64(m_data, 32);
        uint64_t shoff = r64(m_data, 40);
        uint16_t phnum = r16(m_data, 56);
        uint16_t shnum = r16(m_data, 60);
        uint16_t shstrndx = r16(m_data, 62);

        for (uint16_t i = 0; i < phnum; i++)
        {
            size_t off = phoff + i * 56;
            if (off + 56 > m_data.size()) break;
            uint32_t ptype = r32(m_data, off);
            uint32_t pflags = r32(m_data, off + 4);
            uint64_t poffset = r64(m_data, off + 8);
            uint64_t pvaddr = r64(m_data, off + 16);
            uint64_t pfilesz = r64(m_data, off + 32);
            uint64_t pmemsz = r64(m_data, off + 40);

            if (ptype == PT_LOAD)
            {
                m_file_map.push_back({poffset, pvaddr, pmemsz});
                if (pflags & 1)
                    m_exec_ranges.push_back({pvaddr, pvaddr + pmemsz});
            }

            if (ptype == PT_DYNAMIC)
            {
                addr_t dyn_va = pvaddr;
                size_t dyn_sz = pmemsz;
                size_t dyn_off = poffset;

                size_t end = dyn_off + dyn_sz;
                for (size_t j = dyn_off; j + 16 <= end; j += 16)
                {
                    int64_t tag = (int64_t)r64(m_data, j);
                    uint64_t val = r64(m_data, j + 8);

                    if (tag == DT_NULL) break;
                    if (tag == DT_NEEDED && val)
                    {
                        uint64_t strtab = 0;
                        for (size_t k = dyn_off; k + 16 <= end; k += 16)
                        {
                            int64_t t2 = (int64_t)r64(m_data, k);
                            uint64_t v2 = r64(m_data, k + 8);
                            if (t2 == DT_STRTAB) { strtab = v2; break; }
                        }
                        if (strtab)
                        {
                            size_t name_off = strtab + val;
                            if (name_off < m_data.size())
                            {
                                string lib;
                                for (size_t c = name_off; c < m_data.size() && m_data[c]; c++)
                                    lib += (char)m_data[c];
                                if (!lib.empty()) m_needed_libs.push_back(lib);
                            }
                        }
                    }
                    if (tag == DT_INIT) m_init = val;
                    if (tag == DT_FINI) m_fini = val;
                    if (tag == DT_INIT_ARRAY) m_init_array = val;
                    if (tag == DT_INIT_ARRAYSZ) m_init_array_size = val;
                    if (tag == DT_FINI_ARRAY) m_fini_array = val;
                    if (tag == DT_FINI_ARRAYSZ) m_fini_array_size = val;
                }
            }
        }

        struct SectionInfo {
            uint32_t type;
            uint64_t addr;
            uint64_t offset;
            uint64_t size;
            uint32_t link;
            uint64_t entsize;
            uint32_t name_idx;
        };
        vector<SectionInfo> sections(shnum);
        uint64_t shstrtab_off = 0;
        for (uint16_t i = 0; i < shnum; i++)
        {
            size_t off = shoff + i * 64;
            if (off + 64 > m_data.size()) break;
            sections[i].type = r32(m_data, off + 4);
            sections[i].addr = r64(m_data, off + 16);
            sections[i].offset = r64(m_data, off + 24);
            sections[i].size = r64(m_data, off + 32);
            sections[i].link = r32(m_data, off + 40);
            sections[i].entsize = r64(m_data, off + 56);
            sections[i].name_idx = r32(m_data, off);
            if (i == shstrndx) shstrtab_off = sections[i].offset;
        }

        auto sectionName = [&](uint32_t idx) -> string {
            if (!shstrtab_off) return "";
            size_t pos = shstrtab_off + idx;
            string s;
            for (; pos < m_data.size() && m_data[pos]; pos++)
                s += (char)m_data[pos];
            return s;
        };

        vector<uint64_t> strtab_offsets(shnum, 0);
        for (uint16_t i = 0; i < shnum; i++)
        {
            if (sections[i].type == SHT_SYMTAB || sections[i].type == SHT_DYNSYM)
            {
                uint32_t link = sections[i].link;
                if (link < shnum)
                    strtab_offsets[i] = sections[link].offset;
            }
        }

        for (uint16_t i = 0; i < shnum; i++)
        {
            if (sections[i].type != SHT_SYMTAB && sections[i].type != SHT_DYNSYM)
                continue;

            uint64_t sym_off = sections[i].offset;
            uint64_t sym_sz = sections[i].size;
            uint64_t entsize = sections[i].entsize ? sections[i].entsize : 24;
            uint64_t str_off = strtab_offsets[i];

            size_t nsym = sym_sz / entsize;
            for (size_t j = 0; j < nsym; j++)
            {
                size_t ent = sym_off + j * entsize;
                if (ent + 24 > m_data.size()) break;
                uint32_t sym_name_off = r32(m_data, ent);
                uint8_t sym_info = m_data[ent + 4];
                uint8_t sym_other = m_data[ent + 5];
                uint16_t sym_shndx = r16(m_data, ent + 6);
                uint64_t sym_val = r64(m_data, ent + 8);
                uint64_t sym_size = r64(m_data, ent + 16);

                uint8_t sym_type = sym_info & 0xF;
                uint8_t sym_bind = sym_info >> 4;

                if (sym_type != 2) continue;
                if (!sym_name_off || !str_off) continue;

                size_t name_pos = str_off + sym_name_off;
                if (name_pos >= m_data.size()) continue;

                string sym_name;
                for (size_t c = name_pos; c < m_data.size() && m_data[c]; c++)
                    sym_name += (char)m_data[c];
                if (sym_name.empty()) continue;

                if (sym_shndx != 0 && sym_val)
                {
                    m_symbols[sym_val] = sym_name;
                    if (sym_bind == 1)
                        m_exports.push_back({sym_val, sym_name});
                }
            }
        }

        // Collect data section ranges (allocatable, non-executable)
        for (uint16_t i = 0; i < shnum; i++)
        {
            uint64_t sh_flags = 0;
            size_t so = shoff + i * 64;
            if (so + 64 > m_data.size()) continue;
            sh_flags = r64(m_data, so + 8);
            bool is_exec = (sh_flags & 0x4) != 0;
            bool is_alloc = (sh_flags & 0x2) != 0;
            if (is_alloc && !is_exec && sections[i].addr && sections[i].size)
                m_data_ranges.push_back({sections[i].addr, sections[i].addr + sections[i].size});
        }

        for (uint16_t i = 0; i < shnum; i++)
        {
            string sname = sectionName(sections[i].name_idx);
            if (sections[i].type != SHT_RELA && sname != ".rela.plt" && sname != ".rela.dyn")
                continue;

            uint64_t rela_off = sections[i].offset;
            uint64_t rela_sz = sections[i].size;
            size_t nrela = rela_sz / 24;
            if (nrela == 0) continue;

            uint32_t symtab_link = sections[i].link;
            if (symtab_link >= shnum) continue;

            uint64_t sym_off = sections[symtab_link].offset;
            uint64_t sym_sz = sections[symtab_link].size;
            uint64_t str_off2 = 0;
            uint32_t str_link = sections[symtab_link].link;
            if (str_link < shnum) str_off2 = sections[str_link].offset;
            uint64_t sym_entsize = sections[symtab_link].entsize ? sections[symtab_link].entsize : 24;

            string lib_name;
            for (size_t li = 0; li < m_needed_libs.size(); li++)
            {
                if (li) lib_name += ", ";
                lib_name += m_needed_libs[li];
            }

            for (size_t j = 0; j < nrela; j++)
            {
                size_t ent = rela_off + j * 24;
                if (ent + 24 > m_data.size()) break;
                uint64_t r_offset = r64(m_data, ent);
                uint64_t r_info = r64(m_data, ent + 8);
                uint32_t sym_idx = (uint32_t)(r_info >> 32);

                if (!sym_idx) continue;

                size_t sym_ent = sym_off + sym_idx * sym_entsize;
                if (sym_ent + 24 > m_data.size()) continue;
                uint32_t func_name_off = r32(m_data, sym_ent);
                if (!func_name_off || !str_off2) continue;

                size_t fn_pos = str_off2 + func_name_off;
                if (fn_pos >= m_data.size()) continue;

                string func_name;
                for (size_t c = fn_pos; c < m_data.size() && m_data[c]; c++)
                    func_name += (char)m_data[c];
                if (func_name.empty()) continue;

                bool dup = false;
                for (const auto& imp : m_imports)
                    if (imp.address == r_offset) { dup = true; break; }
                if (dup) continue;

                ImportEntry e;
                e.address = r_offset;
                e.name = func_name;
                e.library = lib_name;
                m_imports.push_back(e);
            }
        }
    }
    else
    {
        m_arch = Arch::X86;
        if (m_data.size() < 52) return false;

        m_entry = r32(m_data, 24);
        uint32_t phoff = r32(m_data, 28);
        uint32_t shoff = r32(m_data, 32);
        uint16_t phnum = r16(m_data, 44);
        uint16_t shnum = r16(m_data, 48);
        uint16_t shstrndx = r16(m_data, 50);

        for (uint16_t i = 0; i < phnum; i++)
        {
            size_t off = phoff + i * 32;
            if (off + 32 > m_data.size()) break;
            uint32_t ptype = r32(m_data, off);
            uint32_t poffset = r32(m_data, off + 4);
            uint32_t pvaddr = r32(m_data, off + 8);
            uint32_t pfilesz = r32(m_data, off + 16);
            uint32_t pmemsz = r32(m_data, off + 20);
            uint32_t pflags = r32(m_data, off + 24);

            if (ptype == PT_LOAD)
            {
                m_file_map.push_back({poffset, pvaddr, pmemsz});
                if (pflags & 1)
                    m_exec_ranges.push_back({pvaddr, pvaddr + pmemsz});
            }

            if (ptype == PT_DYNAMIC)
            {
                size_t dyn_off = poffset;
                size_t dyn_sz = pmemsz;
                size_t end = dyn_off + dyn_sz;

                for (size_t j = dyn_off; j + 8 <= end; j += 8)
                {
                    int32_t tag = (int32_t)r32(m_data, j);
                    uint32_t val = r32(m_data, j + 4);

                    if (tag == DT_NULL) break;
                    if (tag == DT_NEEDED && val)
                    {
                        uint32_t strtab = 0;
                        for (size_t k = dyn_off; k + 8 <= end; k += 8)
                        {
                            int32_t t2 = (int32_t)r32(m_data, k);
                            uint32_t v2 = r32(m_data, k + 4);
                            if (t2 == DT_STRTAB) { strtab = v2; break; }
                        }
                        if (strtab)
                        {
                            size_t name_off = strtab + val;
                            if (name_off < m_data.size())
                            {
                                string lib;
                                for (size_t c = name_off; c < m_data.size() && m_data[c]; c++)
                                    lib += (char)m_data[c];
                                if (!lib.empty()) m_needed_libs.push_back(lib);
                            }
                        }
                    }
                    if (tag == DT_INIT) m_init = val;
                    if (tag == DT_FINI) m_fini = val;
                    if (tag == DT_INIT_ARRAY) m_init_array = val;
                    if (tag == DT_INIT_ARRAYSZ) m_init_array_size = val;
                    if (tag == DT_FINI_ARRAY) m_fini_array = val;
                    if (tag == DT_FINI_ARRAYSZ) m_fini_array_size = val;
                }
            }
        }

        struct SectionInfo32 {
            uint32_t type;
            uint32_t addr;
            uint32_t offset;
            uint32_t size;
            uint32_t link;
            uint32_t entsize;
            uint32_t name_idx;
        };
        vector<SectionInfo32> sections(shnum);
        uint32_t shstrtab_off = 0;
        for (uint16_t i = 0; i < shnum; i++)
        {
            size_t off = shoff + i * 40;
            if (off + 40 > m_data.size()) break;
            sections[i].type = r32(m_data, off + 4);
            sections[i].addr = r32(m_data, off + 12);
            sections[i].offset = r32(m_data, off + 16);
            sections[i].size = r32(m_data, off + 20);
            sections[i].link = r32(m_data, off + 28);
            sections[i].entsize = r32(m_data, off + 36);
            sections[i].name_idx = r32(m_data, off);
            if (i == shstrndx) shstrtab_off = sections[i].offset;
        }

        auto sectionName = [&](uint32_t idx) -> string {
            if (!shstrtab_off) return "";
            size_t pos = shstrtab_off + idx;
            string s;
            for (; pos < m_data.size() && m_data[pos]; pos++)
                s += (char)m_data[pos];
            return s;
        };

        vector<uint32_t> strtab_offsets(shnum, 0);
        for (uint16_t i = 0; i < shnum; i++)
        {
            if (sections[i].type == SHT_SYMTAB || sections[i].type == SHT_DYNSYM)
            {
                uint32_t link = sections[i].link;
                if (link < shnum)
                    strtab_offsets[i] = sections[link].offset;
            }
        }

        for (uint16_t i = 0; i < shnum; i++)
        {
            if (sections[i].type != SHT_SYMTAB && sections[i].type != SHT_DYNSYM)
                continue;

            uint32_t sym_off = sections[i].offset;
            uint32_t sym_sz = sections[i].size;
            uint32_t entsize = sections[i].entsize ? sections[i].entsize : 16;
            uint32_t str_off = strtab_offsets[i];

            size_t nsym = sym_sz / entsize;
            for (size_t j = 0; j < nsym; j++)
            {
                size_t ent = sym_off + j * entsize;
                if (ent + 16 > m_data.size()) break;
                uint32_t sym_name_off = r32(m_data, ent);
                uint8_t sym_info = m_data[ent + 4];
                uint8_t sym_other = m_data[ent + 5];
                uint16_t sym_shndx = r16(m_data, ent + 6);
                uint32_t sym_val = r32(m_data, ent + 8);
                uint32_t sym_size = r32(m_data, ent + 12);

                uint8_t sym_type = sym_info & 0xF;
                uint8_t sym_bind = sym_info >> 4;

                if (sym_type != 2) continue;
                if (!sym_name_off || !str_off) continue;

                size_t name_pos = str_off + sym_name_off;
                if (name_pos >= m_data.size()) continue;

                string sym_name;
                for (size_t c = name_pos; c < m_data.size() && m_data[c]; c++)
                    sym_name += (char)m_data[c];
                if (sym_name.empty()) continue;

                if (sym_shndx != 0 && sym_val)
                {
                    m_symbols[sym_val] = sym_name;
                    if (sym_bind == 1)
                        m_exports.push_back({sym_val, sym_name});
                }
            }
        }

        for (uint16_t i = 0; i < shnum; i++)
        {
            uint32_t sh_flags = 0;
            size_t so = shoff + i * 40;
            if (so + 40 > m_data.size()) continue;
            sh_flags = r32(m_data, so + 8);
            bool is_exec = (sh_flags & 0x4) != 0;
            bool is_alloc = (sh_flags & 0x2) != 0;
            if (is_alloc && !is_exec && sections[i].addr && sections[i].size)
                m_data_ranges.push_back({sections[i].addr, sections[i].addr + sections[i].size});
        }

        for (uint16_t i = 0; i < shnum; i++)
        {
            string sname = sectionName(sections[i].name_idx);
            if (sections[i].type != SHT_RELA && sname != ".rela.plt" && sname != ".rela.dyn")
                continue;

            uint32_t rela_off = sections[i].offset;
            uint32_t rela_sz = sections[i].size;
            size_t nrela = rela_sz / 12;
            if (nrela == 0) continue;

            uint32_t symtab_link = sections[i].link;
            if (symtab_link >= shnum) continue;

            uint32_t sym_off = sections[symtab_link].offset;
            uint32_t sym_sz = sections[symtab_link].size;
            uint32_t str_off2 = 0;
            uint32_t str_link = sections[symtab_link].link;
            if (str_link < shnum) str_off2 = sections[str_link].offset;
            uint32_t sym_entsize = sections[symtab_link].entsize ? sections[symtab_link].entsize : 16;

            string lib_name;
            for (size_t li = 0; li < m_needed_libs.size(); li++)
            {
                if (li) lib_name += ", ";
                lib_name += m_needed_libs[li];
            }

            for (size_t j = 0; j < nrela; j++)
            {
                size_t ent = rela_off + j * 12;
                if (ent + 12 > m_data.size()) break;
                uint32_t r_offset = r32(m_data, ent);
                uint32_t r_info = r32(m_data, ent + 4);
                uint32_t sym_idx = r_info >> 8;

                if (!sym_idx) continue;

                size_t sym_ent = sym_off + sym_idx * sym_entsize;
                if (sym_ent + 16 > m_data.size()) continue;
                uint32_t func_name_off = r32(m_data, sym_ent);
                if (!func_name_off || !str_off2) continue;

                size_t fn_pos = str_off2 + func_name_off;
                if (fn_pos >= m_data.size()) continue;

                string func_name;
                for (size_t c = fn_pos; c < m_data.size() && m_data[c]; c++)
                    func_name += (char)m_data[c];
                if (func_name.empty()) continue;

                bool dup = false;
                for (const auto& imp : m_imports)
                    if (imp.address == r_offset) { dup = true; break; }
                if (dup) continue;

                ImportEntry e;
                e.address = r_offset;
                e.name = func_name;
                e.library = lib_name;
                m_imports.push_back(e);
            }
        }
    }

    return true;
}

bool ELFLoader::parseDynamic()
{
    return true;
}

bool ELFLoader::parseSections()
{
    return true;
}

bool ELFLoader::parseSymbols()
{
    return true;
}

vector<uint8_t> ELFLoader::readBytes(addr_t vaddr, size_t size) const
{
    for (const auto& fmap : m_file_map)
    {
        addr_t seg_vaddr = get<1>(fmap);
        uint64_t seg_size = get<2>(fmap);
        if (vaddr >= seg_vaddr && vaddr < seg_vaddr + seg_size)
        {
            uint64_t file_off = get<0>(fmap) + (vaddr - seg_vaddr);
            size_t avail = min(size, (size_t)(seg_vaddr + seg_size - vaddr));
            if (file_off >= m_data.size()) avail = 0;
            else if (file_off + avail > m_data.size())
                avail = m_data.size() - file_off;

            vector<uint8_t> result(m_data.begin() + file_off, m_data.begin() + file_off + avail);
            return result;
        }
    }
    return {};
}

vector<pair<addr_t, addr_t>> ELFLoader::getDataRanges() const
{
    return m_data_ranges;
}

vector<pair<addr_t, addr_t>> ELFLoader::getExecutableRanges() const
{
    return m_exec_ranges;
}

vector<pair<addr_t, string>> ELFLoader::getExportedFunctions() const
{
    return m_exports;
}

vector<ImportEntry> ELFLoader::getImportedFunctions() const
{
    return m_imports;
}