#include <windows.h>
#include <tlhelp32.h>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <psapi.h>

#define IMAGE_REL_BASED_DIR64 10
typedef ULONGLONG QWORD;

struct MANUAL_MAPPING_DATA {
    PVOID ImageBase;
    PVOID EntryPoint;
};

// Улучшенное логирование
void Log(const std::string& message) {
    std::ofstream log("C:\\Windows\\Temp\\inj.log", std::ios::app);
    if (log) {
        log << "[" << __TIME__ << "] " << message << std::endl;
    }
}

DWORD GetProcessIdByName(const std::wstring& processName) {
    DWORD processId = 0;
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (snapshot == INVALID_HANDLE_VALUE) {
        Log("Snapshot error: " + std::to_string(GetLastError()));
        return 0;
    }

    PROCESSENTRY32W entry = { sizeof(PROCESSENTRY32W) };
    if (Process32FirstW(snapshot, &entry)) {
        do {
            if (processName == entry.szExeFile) {
                processId = entry.th32ProcessID;
                break;
            }
        } while (Process32NextW(snapshot, &entry));
    }

    CloseHandle(snapshot);
    return processId;
}

void ApplyRelocations(PIMAGE_NT_HEADERS64 pNtHeaders, PVOID base, ULONGLONG delta) {
    if (delta == 0) return;

    auto dir = &pNtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
    if (dir->Size == 0) return;

    auto reloc = (PIMAGE_BASE_RELOCATION)((ULONGLONG)base + dir->VirtualAddress);
    while (reloc->VirtualAddress > 0) {
        ULONGLONG relocAddr = (ULONGLONG)base + reloc->VirtualAddress;
        UINT numEntries = (reloc->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
        auto entries = (PBASE_RELOCATION_ENTRY)(reloc + 1);

        for (UINT i = 0; i < numEntries; i++) {
            if (entries[i].Type == IMAGE_REL_BASED_DIR64) {
                ULONGLONG* patchAddr = (ULONGLONG*)(relocAddr + entries[i].Offset);
                *patchAddr += delta;
            }
        }
        reloc = (PIMAGE_BASE_RELOCATION)((ULONGLONG)reloc + reloc->SizeOfBlock);
    }
}

bool ResolveImports(HANDLE hProcess, PIMAGE_NT_HEADERS64 pNtHeaders, PVOID base) {
    auto dir = &pNtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (dir->Size == 0) return true;

    auto importDesc = (PIMAGE_IMPORT_DESCRIPTOR)((ULONGLONG)base + dir->VirtualAddress);
    while (importDesc->Name) {
        char* moduleName = (char*)((ULONGLONG)base + importDesc->Name);
        HMODULE hModule = GetModuleHandleA(moduleName);
        if (!hModule) {
            hModule = LoadLibraryA(moduleName);
            if (!hModule) return false;
        }

        auto thunk = (PIMAGE_THUNK_DATA64)((ULONGLONG)base + importDesc->FirstThunk);
        while (thunk->u1.AddressOfData) {
            if (IMAGE_SNAP_BY_ORDINAL64(thunk->u1.Ordinal)) {
                auto proc = GetProcAddress(hModule, (LPCSTR)IMAGE_ORDINAL64(thunk->u1.Ordinal));
                if (!proc) return false;
                thunk->u1.Function = (ULONGLONG)proc;
            }
            else {
                auto import = (PIMAGE_IMPORT_BY_NAME)((ULONGLONG)base + thunk->u1.AddressOfData);
                auto proc = GetProcAddress(hModule, import->Name);
                if (!proc) return false;
                thunk->u1.Function = (ULONGLONG)proc;
            }
            thunk++;
        }
        importDesc++;
    }
    return true;
}

bool ManualMap(HANDLE hProcess, const std::string& dllPath) {
    Log("Starting manual mapping...");

    HANDLE hFile = CreateFileA(dllPath.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        Log("CreateFile failed: " + std::to_string(GetLastError()));
        return false;
    }

    DWORD fileSize = GetFileSize(hFile, NULL);
    std::vector<BYTE> dllData(fileSize);
    DWORD bytesRead;

    if (!ReadFile(hFile, dllData.data(), fileSize, &bytesRead, NULL)) {
        Log("ReadFile failed: " + std::to_string(GetLastError()));
        CloseHandle(hFile);
        return false;
    }
    CloseHandle(hFile);

    auto pDosHeader = (PIMAGE_DOS_HEADER)dllData.data();
    if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
        Log("Invalid DOS header");
        return false;
    }

    auto pNtHeaders = (PIMAGE_NT_HEADERS64)((ULONGLONG)dllData.data() + pDosHeader->e_lfanew);
    if (pNtHeaders->Signature != IMAGE_NT_SIGNATURE ||
        pNtHeaders->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
        Log("Invalid NT header");
        return false;
    }

    PVOID remoteBase = VirtualAllocEx(hProcess, (PVOID)pNtHeaders->OptionalHeader.ImageBase,
        pNtHeaders->OptionalHeader.SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);

    if (!remoteBase) {
        remoteBase = VirtualAllocEx(hProcess, NULL, pNtHeaders->OptionalHeader.SizeOfImage,
            MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
        if (!remoteBase) {
            Log("Memory allocation failed: " + std::to_string(GetLastError()));
            return false;
        }
    }

    WriteProcessMemory(hProcess, remoteBase, dllData.data(), pNtHeaders->OptionalHeader.SizeOfHeaders, NULL);

    auto pSection = IMAGE_FIRST_SECTION(pNtHeaders);
    for (int i = 0; i < pNtHeaders->FileHeader.NumberOfSections; i++, pSection++) {
        PVOID secDest = (BYTE*)remoteBase + pSection->VirtualAddress;
        WriteProcessMemory(hProcess, secDest, dllData.data() + pSection->PointerToRawData, pSection->SizeOfRawData, NULL);
    }

    ULONGLONG delta = (ULONGLONG)remoteBase - pNtHeaders->OptionalHeader.ImageBase;
    ApplyRelocations(pNtHeaders, remoteBase, delta);

    if (!ResolveImports(hProcess, pNtHeaders, remoteBase)) {
        Log("Import resolution failed");
        VirtualFreeEx(hProcess, remoteBase, 0, MEM_RELEASE);
        return false;
    }

    MANUAL_MAPPING_DATA mappingData{};
    mappingData.ImageBase = remoteBase;
    mappingData.EntryPoint = (PVOID)((ULONGLONG)remoteBase + pNtHeaders->OptionalHeader.AddressOfEntryPoint);

    HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0,
        (LPTHREAD_START_ROUTINE)mappingData.EntryPoint,
        mappingData.ImageBase, 0, NULL);

    if (!hThread) {
        Log("Thread creation failed: " + std::to_string(GetLastError()));
        VirtualFreeEx(hProcess, remoteBase, 0, MEM_RELEASE);
        return false;
    }

    WaitForSingleObject(hThread, INFINITE);
    CloseHandle(hThread);
    Log("Manual mapping completed");
    return true;
}

bool Inject(HANDLE hProcess, const std::string& dllPath) {
    return ManualMap(hProcess, dllPath);
}

void HideConsole() {
    HWND hwnd = GetConsoleWindow();
    if (hwnd) ShowWindow(hwnd, SW_HIDE);
}

int main() {
    HideConsole();
    Log("Injector started");

    const std::string dllPath = "C:\\downloads\\your_dll_name.dll";
    const std::wstring targetProcess = L"RainbowSix_DX11.exe";

    // Вариант 1: Запуск через инжектор
    STARTUPINFOW si = { sizeof(si) };
    PROCESS_INFORMATION pi;
    if (CreateProcessW(L"D:\\Games\\Tom Clancy's Rainbow Six Siege\\RainbowSix_DX11.exe",
                      NULL, NULL, NULL, FALSE, CREATE_SUSPENDED, NULL, NULL, &si, &pi)) {
        Log("Process created in suspended state");
        if (Inject(pi.hProcess, dllPath)) {
            Log("Injection successful, resuming thread");
            ResumeThread(pi.hThread);
            CloseHandle(pi.hThread);
            CloseHandle(pi.hProcess);
            return 0;
        }
    }

    // Вариант 2: Инжекция в уже запущенный процесс
    DWORD pid = GetProcessIdByName(targetProcess);
    if (pid) {
        HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
        if (hProcess) {
            if (Inject(hProcess, dllPath)) {
                Log("Injection into existing process successful");
                CloseHandle(hProcess);
                return 0;
            }
            CloseHandle(hProcess);
        }
    }

    Log("Injection failed");
    return 1;
}