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

namespace ApiResolver {

static HMODULE EnsureModule(std::uint32_t moduleHash, const wchar_t* fallbackName);

HMODULE GetNtdll() {
    auto* m = PeParser::GetModuleByHash(Hashes::H_NTDLL_DLL);
    return m ? m : PeParser::GetModuleByIndex(1);
}

HMODULE GetKernel32() {
    auto* m = PeParser::GetModuleByHash(Hashes::H_KERNEL32_DLL);
    return m ? m : PeParser::GetModuleByIndex(2);
}

HMODULE GetWs2_32() {
    return EnsureModule(Hashes::H_WS2_32_DLL, L"ws2_32.dll");
}

HMODULE GetBcrypt() {
    return EnsureModule(Hashes::H_BCRYPT_DLL, L"bcrypt.dll");
}

FARPROC ResolveExport(HMODULE module, std::uint32_t exportHash) {
    return PeParser::GetExportByHash(module, exportHash);
}

FARPROC ResolveNtdll(std::uint32_t exportHash) {
    return ResolveExport(GetNtdll(), exportHash);
}

FARPROC ResolveWs2(std::uint32_t exportHash) {
    return ResolveExport(GetWs2_32(), exportHash);
}

FARPROC ResolveBcrypt(std::uint32_t exportHash) {
    return ResolveExport(GetBcrypt(), exportHash);
}

static HMODULE EnsureModule(std::uint32_t moduleHash, const wchar_t* fallbackName) {
    HMODULE mod = PeParser::GetModuleByHash(moduleHash);
    if (mod) {
        return mod;
    }

    auto* k32 = GetKernel32();
    if (!k32) {
        return nullptr;
    }

    using LoadLibraryW_t = HMODULE(WINAPI*)(LPCWSTR);
    auto* pLoad = reinterpret_cast<LoadLibraryW_t>(
        ResolveExport(k32, Hashes::H_LoadLibraryW));
    if (!pLoad || !fallbackName) {
        return nullptr;
    }

    mod = pLoad(fallbackName);
    return mod;
}

} // namespace ApiResolver