#include "pe_parser.hpp"
#include "../common/hashes.hpp"

#include <cstring>

namespace PeParser {

Peb* GetPeb() {
#ifdef _WIN64
    return reinterpret_cast<Peb*>(__readgsqword(0x60));
#else
    return reinterpret_cast<Peb*>(__readfsdword(0x30));
#endif
}

bool WideToLowerAscii(const PeUnicodeString& name, char* out, size_t outCap) {
    if (!name.Buffer || name.Length == 0 || outCap == 0) {
        return false;
    }
    const int wcharCount = name.Length / sizeof(WCHAR);
    int written = WideCharToMultiByte(
        CP_ACP, 0, name.Buffer, wcharCount,
        out, static_cast<int>(outCap - 1), nullptr, nullptr);
    if (written <= 0) {
        return false;
    }
    out[written] = '\0';
    CharLowerA(out);
    return true;
}

HMODULE GetModuleByHash(std::uint32_t moduleHash) {
    auto* peb = GetPeb();
    if (!peb || !peb->Ldr) {
        return nullptr;
    }

    auto* head = &peb->Ldr->InMemoryOrderModuleList;
    for (auto* link = head->Flink; link != head; link = link->Flink) {
        auto* entry = reinterpret_cast<LdrDataTableEntry*>(
            reinterpret_cast<BYTE*>(link) - offsetof(LdrDataTableEntry, InMemoryOrderLinks));

        if (!entry->BaseDllName.Buffer) {
            continue;
        }

        char narrow[MAX_PATH] = {};
        if (!WideToLowerAscii(entry->BaseDllName, narrow, sizeof(narrow))) {
            continue;
        }

        // Match precomputed hash (api_hash constants) or runtime djb2 lowercase.
        if (Hashes::Djb2(narrow) == moduleHash || Hashes::Djb2Lower(narrow) == moduleHash) {
            return reinterpret_cast<HMODULE>(entry->DllBase);
        }
    }
    return nullptr;
}

HMODULE GetModuleByIndex(unsigned index) {
    auto* peb = GetPeb();
    if (!peb || !peb->Ldr) {
        return nullptr;
    }

    auto* head = &peb->Ldr->InMemoryOrderModuleList;
    unsigned i = 0;
    for (auto* link = head->Flink; link != head; link = link->Flink, ++i) {
        if (i == index) {
            auto* entry = reinterpret_cast<LdrDataTableEntry*>(
                reinterpret_cast<BYTE*>(link) - offsetof(LdrDataTableEntry, InMemoryOrderLinks));
            return reinterpret_cast<HMODULE>(entry->DllBase);
        }
    }
    return nullptr;
}

PIMAGE_NT_HEADERS GetNtHeaders(HMODULE module) {
    if (!module) {
        return nullptr;
    }
    auto* base = reinterpret_cast<BYTE*>(module);
    auto* dos = reinterpret_cast<PIMAGE_DOS_HEADER>(base);
    if (dos->e_magic != IMAGE_DOS_SIGNATURE) {
        return nullptr;
    }
    auto* nt = reinterpret_cast<PIMAGE_NT_HEADERS>(base + dos->e_lfanew);
    if (nt->Signature != IMAGE_NT_SIGNATURE) {
        return nullptr;
    }
    return nt;
}

FARPROC GetExportByHash(HMODULE module, std::uint32_t exportHash) {
    if (!module) {
        return nullptr;
    }

    auto* base = reinterpret_cast<BYTE*>(module);
    auto* nt = GetNtHeaders(module);
    if (!nt) {
        return nullptr;
    }

    const auto& dir = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
    if (dir.VirtualAddress == 0 || dir.Size == 0) {
        return nullptr;
    }

    auto* exp = reinterpret_cast<PIMAGE_EXPORT_DIRECTORY>(base + dir.VirtualAddress);
    auto* names = reinterpret_cast<PDWORD>(base + exp->AddressOfNames);
    auto* ords  = reinterpret_cast<PWORD>(base + exp->AddressOfNameOrdinals);
    auto* funcs = reinterpret_cast<PDWORD>(base + exp->AddressOfFunctions);

    for (DWORD i = 0; i < exp->NumberOfNames; ++i) {
        const char* exportName = reinterpret_cast<const char*>(base + names[i]);
        if (Hashes::Djb2(exportName) != exportHash) {
            continue;
        }
        return reinterpret_cast<FARPROC>(base + funcs[ords[i]]);
    }
    return nullptr;
}

} // namespace PeParser