// system/DllInjector.h
#pragma once

#ifdef CURRENT_MODULE
#undef CURRENT_MODULE
#endif
#define CURRENT_MODULE "DllInjector"

#include "core/FilePath.h"
#include "core/Logger.h"
#include <windows.h>
#include <string>
#include <iostream>

namespace DllInjector
{
    struct HandleCloser { void operator()(HANDLE h) const noexcept { if (h && h != INVALID_HANDLE_VALUE) CloseHandle(h); } };
    using UniqueHandle = std::unique_ptr<void, HandleCloser>;

    /// @brief Inject a specific DLL to a given process identified by PID using LoadLibraryA
    /// @param processId  Target process PID
    /// @param dllPath    Full DLL path in ANSI
    /// @return true = Injected successfully with a non-zero handle
    inline bool Inject(DWORD processId, const std::string& dllPath)
    {
        UniqueHandle hProcess{ OpenProcess(
            PROCESS_QUERY_INFORMATION | PROCESS_CREATE_THREAD |
            PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ,
            FALSE, processId) };

        if (!hProcess)
        {
            LOG_E("OpenProcess failed, PID=" + std::to_string(processId) + ", Error=" + std::to_string(GetLastError()));
            return false;
        }

        SIZE_T pathBytes = dllPath.size() + 1;
        LPVOID remoteMem = VirtualAllocEx(hProcess.get(), nullptr, pathBytes, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
        if (!remoteMem)
        {
            LOG_E("VirtualAllocEx failed, Error=" + std::to_string(GetLastError()));
            return false;
        }

        if (!WriteProcessMemory(hProcess.get(), remoteMem, dllPath.c_str(), pathBytes, nullptr))
        {
            LOG_E("WriteProcessMemory failed, Error=" + std::to_string(GetLastError()));
            VirtualFreeEx(hProcess.get(), remoteMem, 0, MEM_RELEASE);
            return false;
        }

        HMODULE hKernel32 = GetModuleHandleW(L"kernel32.dll");
        if (!hKernel32)
        {
            VirtualFreeEx(hProcess.get(), remoteMem, 0, MEM_RELEASE);
            LOG_E("GetModuleHandleW(kernel32.dll) failed");
            return false;
        }

        auto pLoadLibraryA = reinterpret_cast<LPTHREAD_START_ROUTINE>(GetProcAddress(hKernel32, "LoadLibraryA"));

        UniqueHandle hThread{ CreateRemoteThread(hProcess.get(), nullptr, 0, pLoadLibraryA, remoteMem, 0, nullptr) };
        if (!hThread)
        {
            LOG_E("CreateRemoteThread failed, Error=" + std::to_string(GetLastError()));
            VirtualFreeEx(hProcess.get(), remoteMem, 0, MEM_RELEASE);
            return false;
        }

        WaitForSingleObject(hThread.get(), 8000);

        DWORD exitCode = 0;
        GetExitCodeThread(hThread.get(), &exitCode);
        VirtualFreeEx(hProcess.get(), remoteMem, 0, MEM_RELEASE);

        if (exitCode == 0)
        {
            LOG_E("LoadLibraryA returned NULL in target process");
            return false;
        }

        LOG_I("Injected successfully! hModule=0x" + std::to_string(exitCode));
        return true;
    }

	/// @brief Injects a DLL next to the current process's executable into a given process identified by PID
    /// @param processId     Target process PID
    /// @param dllFileName   Just the file name e.g. "iOperator.dll"
    inline bool InjectLocalDll(DWORD processId, const std::string& dllFileName)
    {
        std::string fullPath = FilePath::GetExecutableDirectory() + dllFileName;
        return Inject(processId, fullPath);
    }
}