#include "loader/pe.hpp"
#include <fstream>
#include <cstring>

PELoader::PELoader() : m_arch(Arch::X64), m_entry(0), m_image_base(0) {}
PELoader::~PELoader() {}

bool PELoader::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();

    return parseHeaders();
}

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

bool PELoader::parseHeaders()
{
    if (m_data.size() < 64) return false;
    if (r16(m_data, 0) != 0x5A4D) return false;

    uint32_t pe_off = r32(m_data, 0x3C);
    if (pe_off + 4 > m_data.size()) return false;
    if (r32(m_data, pe_off) != 0x00004550) return false;

    uint16_t machine = r16(m_data, pe_off + 4);
    uint16_t num_sections = r16(m_data, pe_off + 6);

    if (machine == 0x8664) m_arch = Arch::X64;
    else if (machine == 0x14C) m_arch = Arch::X86;
    else return false;

    uint16_t opt_hdr_size = r16(m_data, pe_off + 20);
    uint32_t opt_hdr_off = pe_off + 24;

    if (opt_hdr_off + 2 > m_data.size()) return false;
    uint16_t magic = r16(m_data, opt_hdr_off);

    if (magic == 0x10B)
    {
        m_arch = Arch::X86;
        m_image_base = r32(m_data, opt_hdr_off + 28);
        m_entry = r32(m_data, opt_hdr_off + 16) + m_image_base;
    }
    else if (magic == 0x20B)
    {
        m_arch = Arch::X64;
        m_image_base = r64(m_data, opt_hdr_off + 24);
        m_entry = r32(m_data, opt_hdr_off + 16) + m_image_base;
    }
    else return false;

    uint16_t section_hdr_size = 40;
    uint32_t sections_off = opt_hdr_off + opt_hdr_size;

    struct Section
    {
        string name;
        addr_t vaddr;
        uint32_t vsize;
        uint32_t raw_ptr;
        uint32_t raw_size;
        uint32_t characteristics;
    };
    vector<Section> sections;

    for (uint16_t i = 0; i < num_sections; i++)
    {
        uint32_t off = sections_off + i * section_hdr_size;
        if (off + section_hdr_size > m_data.size()) break;

        char name_buf[9] = {};
        memcpy(name_buf, &m_data[off], 8);
        Section s;
        s.name = name_buf;
        s.vsize = r32(m_data, off + 8);
        s.vaddr = r32(m_data, off + 12) + m_image_base;
        s.raw_size = r32(m_data, off + 16);
        s.raw_ptr = r32(m_data, off + 20);
        s.characteristics = r32(m_data, off + 36);
        sections.push_back(s);
    }

    for (const auto& s : sections)
    {
        if (s.characteristics & 0x20000000)
        {
            addr_t start = s.vaddr;
            addr_t end = start + (s.vsize ? s.vsize : s.raw_size);
            m_exec_ranges.push_back({start, end});
        }
    }

    // Collect data section ranges (non-executable readable)
    for (const auto& s : sections)
    {
        if (s.raw_ptr == 0) continue;
        bool is_exec = (s.characteristics & 0x20000000) != 0;
        bool is_read = (s.characteristics & 0x40000000) != 0;
        if (!is_exec && is_read)
        {
            addr_t start = s.vaddr;
            addr_t end = start + (s.vsize ? s.vsize : s.raw_size);
            m_data_ranges.push_back({start, end});
        }
    }

    // Parse .pdata (exception handler table) for precise function boundaries
    uint32_t pdata_rva = 0;
    uint32_t pdata_size = 0;
    uint32_t data_dir_off = (magic == 0x10B) ? opt_hdr_off + 96 : opt_hdr_off + 112;
    if (magic == 0x20B)
    {
        pdata_rva = r32(m_data, data_dir_off + 24);
        pdata_size = r32(m_data, data_dir_off + 28);
    }
    if (pdata_rva && pdata_size >= 12)
    {
        uint32_t pdata_off = 0;
        for (const auto& s : sections)
        {
            addr_t va = pdata_rva + m_image_base;
            if (va >= s.vaddr && va < s.vaddr + (s.vsize ? s.vsize : s.raw_size))
            {
                pdata_off = s.raw_ptr + static_cast<uint32_t>(va - s.vaddr);
                break;
            }
        }
        if (pdata_off)
        {
            size_t nentries = pdata_size / 12;
            for (size_t i = 0; i < nentries; i++)
            {
                uint32_t off = pdata_off + i * 12;
                if (off + 12 > m_data.size()) break;
                uint32_t begin_rva = r32(m_data, off);
                uint32_t end_rva = r32(m_data, off + 4);
                if (begin_rva && end_rva && end_rva > begin_rva)
                    m_pdata.push_back({begin_rva + m_image_base, end_rva + m_image_base});
            }
        }
    }

    bool is_64 = (m_arch == Arch::X64);
    uint32_t import_rva = 0;
    import_rva = r32(m_data, data_dir_off + 8);

    auto rvaToOffset = [&](uint32_t rva) -> uint32_t {
        for (const auto& s : sections)
        {
            addr_t va = rva + m_image_base;
            if (va >= s.vaddr && va < s.vaddr + (s.vsize ? s.vsize : s.raw_size))
                return s.raw_ptr + static_cast<uint32_t>(va - s.vaddr);
        }
        return 0;
    };

    auto readString = [&](uint32_t off) -> string {
        string s;
        for (size_t c = off; c < m_data.size() && m_data[c]; c++)
            s += (char)m_data[c];
        return s;
    };

    if (import_rva)
    {
        uint32_t desc_off = rvaToOffset(import_rva);

        for (uint32_t i = 0; ; i++)
        {
            uint32_t desc = desc_off + i * 20;
            if (desc + 20 > m_data.size()) break;

            uint32_t oft_rva = r32(m_data, desc);
            if (oft_rva == 0) break;

            uint32_t name_rva = r32(m_data, desc + 12);
            string dll_name;
            if (name_rva)
            {
                uint32_t name_off = rvaToOffset(name_rva);
                if (name_off) dll_name = readString(name_off);
            }
            if (dll_name.empty()) continue;

            uint32_t thunk_rva = r32(m_data, desc + 16);
            if (thunk_rva == 0) thunk_rva = oft_rva;

            uint32_t thunk_off = rvaToOffset(thunk_rva);
            if (!thunk_off) continue;

            size_t entry_size = is_64 ? 8 : 4;
            uint64_t ordinal_flag = is_64 ? 0x8000000000000000ULL : 0x80000000;

            for (uint32_t j = 0; ; j++)
            {
                uint32_t entry = static_cast<uint32_t>(thunk_off + j * entry_size);
                if (entry + entry_size > m_data.size()) break;

                uint64_t val = is_64 ? r64(m_data, entry) : r32(m_data, entry);
                if (val == 0) break;

                string func_name;
                if (val & ordinal_flag)
                {
                    func_name = "ord_" + to_string(val & 0xFFFF);
                }
                else
                {
                    addr_t hint_va = val + m_image_base;
                    uint32_t hint_off = rvaToOffset(static_cast<uint32_t>(val));
                    if (hint_off && hint_off + 2 < m_data.size())
                        func_name = readString(hint_off + 2);
                }

                addr_t iat_addr = thunk_rva + m_image_base + j * entry_size;
                ImportEntry e;
                e.address = iat_addr;
                e.name = func_name;
                e.library = dll_name;
                m_imports.push_back(e);
            }
        }
    }

    // Parse export directory
    uint32_t export_rva = r32(m_data, data_dir_off);
    if (export_rva)
    {
        uint32_t export_off = rvaToOffset(export_rva);
        if (export_off && export_off + 40 <= m_data.size())
        {
            uint32_t export_base = r32(m_data, export_off + 16);
            uint32_t num_funcs = r32(m_data, export_off + 20);
            uint32_t num_names = r32(m_data, export_off + 24);
            uint32_t funcs_rva = r32(m_data, export_off + 28);
            uint32_t names_rva = r32(m_data, export_off + 32);
            uint32_t ords_rva = r32(m_data, export_off + 36);

            uint32_t funcs_off = rvaToOffset(funcs_rva);
            uint32_t names_off = rvaToOffset(names_rva);
            uint32_t ords_off = rvaToOffset(ords_rva);

            if (funcs_off && names_off && ords_off)
            {
                for (uint32_t i = 0; i < num_names; i++)
                {
                    uint32_t name_rva = r32(m_data, names_off + i * 4);
                    uint32_t name_off = rvaToOffset(name_rva);
                    if (!name_off) continue;

                    string exp_name = readString(name_off);
                    if (exp_name.empty()) continue;

                    uint16_t ord = r16(m_data, ords_off + i * 2);
                    if (ord >= num_funcs) continue;

                    uint32_t func_rva = r32(m_data, funcs_off + ord * 4);
                    if (func_rva == 0) continue;

                    addr_t func_va = func_rva + m_image_base;
                    m_exports.push_back({func_va, exp_name});
                }
            }
        }
    }

    return true;
}

bool PELoader::parseImportTable()
{
    return true;
}

bool PELoader::parseExportTable()
{
    return true;
}

vector<uint8_t> PELoader::readBytes(addr_t vaddr, size_t size) const
{
    vector<uint8_t> result;
    uint32_t off = 0;
    uint32_t raw_off = 0;
    uint32_t raw_end = 0;

    uint16_t num_sections = r16(m_data, r32(m_data, 0x3C) + 6);
    uint32_t pe_off = r32(m_data, 0x3C);
    uint16_t opt_sz = r16(m_data, pe_off + 20);
    uint32_t sections_off = pe_off + 24 + opt_sz;

    for (uint16_t i = 0; i < num_sections; i++)
    {
        uint32_t so = sections_off + i * 40;
        if (so + 40 > m_data.size()) break;
        addr_t sva = r32(m_data, so + 12) + m_image_base;
        uint32_t svs = r32(m_data, so + 8);
        uint32_t srs = r32(m_data, so + 16);
        uint32_t srp = r32(m_data, so + 20);
        addr_t end = sva + (svs ? svs : srs);

        if (vaddr >= sva && vaddr < end)
        {
            off = srp + static_cast<uint32_t>(vaddr - sva);
            raw_off = srp;
            raw_end = srp + srs;
            break;
        }
    }

    if (!off) return result;

    size_t available = (raw_end > off) ? (raw_end - off) : 0;
    size_t to_read = min(size, available);
    if (off + to_read > m_data.size())
        to_read = m_data.size() - off;

    if (to_read > 0)
        result.assign(m_data.begin() + off, m_data.begin() + off + to_read);

    return result;
}

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

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

vector<pair<addr_t, addr_t>> PELoader::getRuntimeFunctions() const
{
    return m_pdata;
}

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

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