/***
    Copyright (C) 2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: UtilProcessMap.NT.cpp
    Date: 2022-1-24
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuProcessMap.NT.hpp"
#include "AuProcessMap.hpp"

namespace Aurora::Process
{
    static AuThreadPrimitives::Mutex gMutex;
    static AuBST<AuUInt, AuString> gPathCache;
    static AuBST<AuUInt, AuString> gModNameCache;

    AuString ModuleToSomething(HMODULE handle, bool path)
    {
        try
        {
            std::wstring file;

            if (!handle)
            {
                return {};
            }

            if (gMutex)
            {
                AU_LOCK_GUARD(gMutex);
                if (path)
                {
                    auto itr = gPathCache.find(reinterpret_cast<AuUInt>(handle));
                    if (itr != gPathCache.end()) return itr->second;
                }
                else
                {
                    auto itr = gModNameCache.find(reinterpret_cast<AuUInt>(handle));
                    if (itr != gModNameCache.end()) return itr->second;
                }
            }

            if (!AuTryResize(file, 16 * 1024))
            {
                return {};
            }

            auto length = GetModuleFileNameW(handle, &file[0], DWORD(file.size()));
            if (!length)
            {
                return {};
            }

            file.resize(length);

            if (!path)
            {
                auto idx = file.find_last_of('\\');
                if (idx != std::wstring::npos)
                {
                    file = file.substr(idx + 1);
                }
            }

            auto ret = Locale::ConvertFromWChar(file.c_str(), file.length());

            if (gMutex)
            {
                AU_LOCK_GUARD(gMutex);
                if (path)
                {
                    AuTryInsert(gPathCache, AuMakePair(reinterpret_cast<AuUInt>(handle), ret));
                }
                else
                {
                    AuTryInsert(gModNameCache, AuMakePair(reinterpret_cast<AuUInt>(handle), ret));
                }
            }

            return ret;
        }
        catch (...)
        {
            return "";
        }
    }

    AuString ModuleToName(HMODULE handle)
    {
        return ModuleToSomething(handle, false);
    }

    AuString ModuleToPath(HMODULE handle)
    {
        return ModuleToSomething(handle, true);
    }

    static AuUInt GetModuleBaseAddressFromPathGrossYesThisIsCached(HMODULE mod, const AuString &path)
    {
        if (path.empty())
        {
            return 0;
        }

        auto dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(mod);
        auto nt = &(reinterpret_cast<PIMAGE_NT_HEADERS>(reinterpret_cast<AuUInt>(mod) + dosHeader->e_lfanew)->OptionalHeader.ImageBase);

        auto offset = static_cast<AuUInt32>(reinterpret_cast<const AuUInt8 *>(nt) - reinterpret_cast<const AuUInt8 *>(mod));
        using Value_t = AuRemovePointer_t<decltype(nt)>;
        Value_t value {};
        DWORD dwBytesRead {};

        auto hHandle = CreateFileW(Locale::ConvertFromUTF8(path).c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, NULL, NULL);
        if (hHandle == INVALID_HANDLE_VALUE)
        {
            return 0;
        }
        
        if (!SetFilePointer(hHandle, offset, NULL, 0))
        {
            AuWin32CloseHandle(hHandle);
            return {};
        }

        if (!ReadFile(hHandle, &value, sizeof(Value_t), &dwBytesRead, NULL))
        {
            AuWin32CloseHandle(hHandle);
            return {};
        }

        if (dwBytesRead != sizeof(Value_t))
        {
            AuWin32CloseHandle(hHandle);
            return {};
        }

        AuWin32CloseHandle(hHandle);
        return value;
    }

    static void FetchModuleSections(Sections &sections, AuUInt offset, HMODULE mod)
    {
        auto dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(mod);
        auto nt = reinterpret_cast<PIMAGE_NT_HEADERS>(reinterpret_cast<AuUInt>(mod) + dosHeader->e_lfanew);

        auto pSection = IMAGE_FIRST_SECTION(nt);
        auto pSectionCur = pSection;

        auto imageBase = nt->OptionalHeader.ImageBase;

        for (auto i = 0; i < nt->FileHeader.NumberOfSections; i++)
        {
            auto cur = pSectionCur++;
            Section seg;

            seg.origVa = cur->VirtualAddress + offset;
            seg.baseVa = cur->VirtualAddress + imageBase;
            seg.size   = cur->SizeOfRawData;
            seg.fsOff  = cur->PointerToRawData;

            seg.pt.readable = cur->Characteristics & (IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_EXECUTE);
            seg.pt.NX       = !(cur->Characteristics & IMAGE_SCN_MEM_EXECUTE);
            seg.pt.writable = cur->Characteristics & IMAGE_SCN_MEM_WRITE;

            seg.name = AuString(cur->Name, cur->Name + strnlen(reinterpret_cast<const char *>(cur->Name), sizeof(cur->Name)));

            AuTryInsert(sections, seg);
        }
    }

    static AuSPtr<PublicModule> HandleToPublicModule(HMODULE h)
    {
        auto pub = AuMakeShared<PublicModule>();
        if (!pub) return {};
        pub->moduleMeta = AuMakeShared<ModuleMeta>();
        if (!pub->moduleMeta) return {};

        pub->moduleMeta->moduleBase = AuUInt(h);
        pub->moduleMeta->moduleName = ModuleToName(h);
        pub->moduleMeta->modulePath = ModuleToPath(h);
        pub->moduleMeta->origVa = GetModuleBaseAddressFromPathGrossYesThisIsCached(h, pub->moduleMeta->modulePath);
        FetchModuleSections(pub->sections, pub->moduleMeta->origVa, h);
        return pub;
    }

    void InvaildateModule(HMODULE hmod)
    {
        if (!gMutex)
        {
            return;
        }

        try
        {
            // acquire gmutex
            {
                AU_LOCK_GUARD(gMutex);
                auto itr1 = gPathCache.find(reinterpret_cast<AuUInt>(hmod));
                if (itr1 != gPathCache.end()) gPathCache.erase(itr1);

                auto itr2 = gModNameCache.find(reinterpret_cast<AuUInt>(hmod));
                if (itr2 != gModNameCache.end()) gModNameCache.erase(itr2);
            }

            // release lock
            {
                RemoveModuleCache({"", reinterpret_cast<AuUInt>(hmod)});
            }
        }
        catch (...)
        {

        }
    }

    bool MakeAwarePtr(AuUInt pointer)
    {
        HMODULE handle;
        if (!GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, reinterpret_cast<LPCWSTR>(pointer), &handle))
        {
            return false;
        }

        MakeAware(handle);
        return true;
    }

    AuOptional<Section> LookupArbitrarySection(AuUInt address)
    {
        MEMORY_BASIC_INFORMATION info;
        Section section;

        if (!VirtualQuery((LPCVOID)address, &info, sizeof(info)))
        {
            return {};
        }

        section.baseVa = reinterpret_cast<AuUInt>(info.BaseAddress);
        section.size = info.RegionSize;
        section.pt.NX = (info.Protect & (PAGE_EXECUTE | PAGE_EXECUTE_READ)) == 0;
        section.pt.writable = (info.Protect & (PAGE_READWRITE | PAGE_WRITECOPY | PAGE_EXECUTE_WRITECOPY)) != 0;
        section.pt.readable = (!section.pt.NX) && ((info.Protect & (PAGE_NOACCESS)) != 0);
        section.pt.acSanity = info.Protect == PAGE_EXECUTE_WRITECOPY;

        if (info.Type & (MEM_MAPPED | MEM_IMAGE))
        {
            section.name = kSectionNameFile;
        }
        else if (info.Type & (MEM_PRIVATE))
        {
            section.name = kSectionNameHeap;
        }

        return section;
    }

    void MakeAware(HMODULE hmod)
    {
        ModuleBasePair handle {"", reinterpret_cast<AuUInt>(hmod)};

        if (!gMutex)
        {
            return;
        }

        if (hmod == INVALID_HANDLE_VALUE)
        {
            SysPushErrorArg();
            return;
        }

        if (!hmod)
        {
            SysPushErrorArg();
            return;
        }

        try
        {
            if (IsInModuleCache(handle))
            {
                return;
            }

            auto ptr = HandleToPublicModule(hmod);
            if (!ptr)
            {
                return;
            }

            InsertModuleCache(handle, ptr);
        }
        catch (...)
        {

        }
    }

    PublicModule GetExecutableRoot()
    {
        try
        {
            auto handle = GetModuleHandleW(NULL);
            MakeAware(handle);
            return GetFromModuleCache(reinterpret_cast<AuUInt>(handle));
        }
        catch (...)
        {
            return {};
        }
    }
    
    void InitProcessMapNt()
    {
    }

    void DeinitProcessMapNt()
    {
    }
}