#include <stdio.h>
#include <stdlib.h>
#include <windows.h>
#include <tlhelp32.h>
#include <psapi.h>
#include <new>
#include <cstdarg>

// Structure to hold section information
struct SectionInfo {
    unsigned long long start;
    unsigned long long end;
    unsigned long long size;
    unsigned long long original_start;  // Start address as seen in IDA
};

// Function to get process ID by name
DWORD get_process_id_by_name(const wchar_t* process_name) {
    PROCESSENTRY32W pe32;
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (hSnapshot == INVALID_HANDLE_VALUE) {
        return 0;
    }

    pe32.dwSize = sizeof(PROCESSENTRY32W);
    if (!Process32FirstW(hSnapshot, &pe32)) {
        CloseHandle(hSnapshot);
        return 0;
    }

    do {
        if (_wcsicmp(pe32.szExeFile, process_name) == 0) {
            CloseHandle(hSnapshot);
            return pe32.th32ProcessID;
        }
    } while (Process32NextW(hSnapshot, &pe32));

    CloseHandle(hSnapshot);
    return 0;
}

// Function to get the base address of a module in a given process
unsigned long long get_base_address(DWORD pid, const wchar_t* module_name) {
    MODULEENTRY32W me32;
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32, pid);

    if (hSnapshot == INVALID_HANDLE_VALUE) {
        return 0;
    }

    me32.dwSize = sizeof(MODULEENTRY32W);
    if (!Module32FirstW(hSnapshot, &me32)) {
        CloseHandle(hSnapshot);
        return 0;
    }

    unsigned long long base_address = 0;

    do {
        if (_wcsicmp(me32.szModule, module_name) == 0) {
            base_address = (unsigned long long)me32.modBaseAddr;
            break;
        }
    } while (Module32NextW(hSnapshot, &me32));

    CloseHandle(hSnapshot);
    return base_address;
}

// Function to get section information
bool get_section_info(HANDLE hProcess, unsigned long long base_addr, const char* section_name, SectionInfo& section_info) {
    IMAGE_DOS_HEADER dos_header;
    if (!ReadProcessMemory(hProcess, (LPCVOID)base_addr, &dos_header, sizeof(dos_header), NULL)) {
        return false;
    }

    IMAGE_NT_HEADERS64 nt_headers;
    if (!ReadProcessMemory(hProcess, (LPCVOID)(base_addr + dos_header.e_lfanew), &nt_headers, sizeof(nt_headers), NULL)) {
        return false;
    }

    unsigned long long section_offset = base_addr + dos_header.e_lfanew + sizeof(IMAGE_NT_HEADERS64);
    DWORD section_count = nt_headers.FileHeader.NumberOfSections;
    IMAGE_SECTION_HEADER* section_headers = new (std::nothrow) IMAGE_SECTION_HEADER[section_count];
    if (!section_headers) {
        return false;
    }

    if (!ReadProcessMemory(hProcess, (LPCVOID)section_offset, section_headers, sizeof(IMAGE_SECTION_HEADER) * section_count, NULL)) {
        delete[] section_headers;
        return false;
    }

    bool found = false;
    for (DWORD i = 0; i < section_count; i++) {
        if (strncmp((char*)section_headers[i].Name, section_name, IMAGE_SIZEOF_SHORT_NAME) == 0) {
            // Calculate precise section start including alignment
            section_info.start = base_addr + section_headers[i].VirtualAddress;

            // Calculate aligned size
            DWORD aligned_size = section_headers[i].Misc.VirtualSize;
            if (aligned_size % nt_headers.OptionalHeader.SectionAlignment != 0) {
                aligned_size = (aligned_size + nt_headers.OptionalHeader.SectionAlignment - 1)
                    & ~(nt_headers.OptionalHeader.SectionAlignment - 1);
            }

            section_info.end = section_info.start + aligned_size;
            section_info.size = aligned_size;

            // Calculate original address using precise ImageBase
            section_info.original_start = nt_headers.OptionalHeader.ImageBase + section_headers[i].VirtualAddress;

            found = true;
            break;
        }
    }

    delete[] section_headers;
    return found;
}

// Function to find a string in a memory section
unsigned long long find_string_in_section(HANDLE hProcess, const SectionInfo& section_info, const char* string_to_search) {
    SIZE_T string_length = strlen(string_to_search);
    unsigned long long addr = section_info.start;
    const SIZE_T CHUNK_SIZE = 1024 * 1024; // 1MB chunks

    while (addr < section_info.end) {
        MEMORY_BASIC_INFORMATION mbi;
        if (!VirtualQueryEx(hProcess, (LPCVOID)addr, &mbi, sizeof(mbi))) {
            break;
        }

        SIZE_T read_size = min(CHUNK_SIZE, mbi.RegionSize);
        char* buffer = new (std::nothrow) char[read_size];
        if (!buffer) {
            break;
        }

        SIZE_T bytesRead;
        if (ReadProcessMemory(hProcess, (LPCVOID)addr, buffer, read_size, &bytesRead) && bytesRead > 0) {
            for (SIZE_T i = 0; i < bytesRead - string_length; i++) {
                if (memcmp(buffer + i, string_to_search, string_length) == 0) {
                    unsigned long long string_addr = addr + i;
                    delete[] buffer;
                    return string_addr;
                }
            }
        }

        delete[] buffer;
        addr += read_size;
    }

    return 0;
}

// Function to find a reference to an address in the data section
unsigned long long find_reference_in_data_section(HANDLE hProcess, const SectionInfo& section_info, unsigned long long string_addr) {
    unsigned long long addr = section_info.start;
    const SIZE_T CHUNK_SIZE = 1024 * 1024; // 1MB chunks

    while (addr < section_info.end) {
        MEMORY_BASIC_INFORMATION mbi;
        if (!VirtualQueryEx(hProcess, (LPCVOID)addr, &mbi, sizeof(mbi))) {
            break;
        }

        SIZE_T read_size = min(CHUNK_SIZE, mbi.RegionSize);
        read_size = read_size - (read_size % sizeof(unsigned long long)); // Ensure alignment

        unsigned long long* buffer = new (std::nothrow) unsigned long long[read_size / sizeof(unsigned long long)];
        if (!buffer) {
            break;
        }

        SIZE_T bytesRead;
        if (ReadProcessMemory(hProcess, (LPCVOID)addr, buffer, read_size, &bytesRead) && bytesRead > 0) {
            SIZE_T numElements = bytesRead / sizeof(unsigned long long);
            for (SIZE_T i = 0; i < numElements; i++) {
                if (buffer[i] == string_addr) {
                    unsigned long long reference_addr = addr + i * sizeof(unsigned long long);
                    delete[] buffer;
                    return reference_addr;
                }
            }
        }

        delete[] buffer;
        addr += read_size;
    }
    return 0;
}

int main() {
    const wchar_t* process_name = L"Fallout76.exe";
    const wchar_t* module_name = L"Fallout76.exe";

    const char* gameSettings[] = {
        "fVATSMaxChance",
        "fVATSChancePerceptionMultiplier",
        "fVATSMaxEngageDistance",
        "fVATSChanceRaycastVisibilityMininmum",
        "fVATSPingCost",
        "fVATSCriticalChargeBase",
        "fVATSCriticalChargeBonusMult",
        "fVATSCriticalChargeLimit",
        "fVATSCriticalChargeMult",
        "fPlayerAutoReloadDelaySec",
        "bDisablePlayerCollision",
        "fAVDCarryWeightBase",
        "fAbsolutePlayerInventoryWeightLimit",
        "fInvisibilityMaxRefraction",
        "fInvisibilityMinRefraction",
        "fJetpackThrustInitial",
        "fJetpackThrustSustained",
        "fJetpackTimeToSustained",
        "bJetpackInitialThrustOnlyOnTakeoff",
        "fAimStabilityCrouchedScopeMult",
        "fPropModScaleWeapDurability",
        "fLockpickSkillSweetSpotBase",
        "fSweetSpotEasy",
        "fSweetSpotAverage",
        "fSweetSpotHard",
        "fSweetSpotVeryHard",
        "iHackingMinWords",
        "iHackingMaxWords",
        "fCombatMusicNearCombatInnerRadius",
        "fCombatMusicNearCombatOuterRadius",
        "fCombatMusicPlayerTargetedThreatRatio",
        "pkKeywords",
        "fAttackActionPointCost",
        "fspeed",
        "fweight",
        "iDurability",
        "ihealth",
        "bHoldInputToPower",
        "fMinPowerPerShot",
        "poAmmo",
        "paAimModel",
        "ppOverrideProjectile",
        "uNumProjectiles",
        "vaActorValues",
        "ihealth"
    };

    FILE* log_file;
    if (_wfopen_s(&log_file, L"memscan.log", L"w") != 0) {
        wprintf(L"Could not open log file for writing\n");
        return 1;
    }

    DWORD pid = get_process_id_by_name(process_name);
    if (pid == 0) {
        fprintf(log_file, "Process not found\n");
        fclose(log_file);
        return 1;
    }

    HANDLE hProcess = OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, FALSE, pid);
    if (hProcess == NULL) {
        fprintf(log_file, "Failed to open process. Error: %lu\n", GetLastError());
        fclose(log_file);
        return 1;
    }

    unsigned long long base_addr = get_base_address(pid, module_name);
    if (base_addr == 0) {
        fprintf(log_file, "Failed to get base address\n");
        CloseHandle(hProcess);
        fclose(log_file);
        return 1;
    }

    SectionInfo rdata_info, data_info;

    if (!get_section_info(hProcess, base_addr, ".rdata", rdata_info)) {
        fprintf(log_file, "Failed to get .rdata section info\n");
        CloseHandle(hProcess);
        fclose(log_file);
        return 1;
    }

    if (!get_section_info(hProcess, base_addr, ".data", data_info)) {
        fprintf(log_file, "Failed to get .data section info\n");
        CloseHandle(hProcess);
        fclose(log_file);
        return 1;
    }

    fprintf(log_file, "Section Information:\n");
    fprintf(log_file, ".rdata: Runtime: 0x%llx - 0x%llx (size: 0x%llx)\n",
        rdata_info.start, rdata_info.end, rdata_info.size);
    fprintf(log_file, "        Original: 0x%llx\n", rdata_info.original_start);
    fprintf(log_file, ".data:  Runtime: 0x%llx - 0x%llx (size: 0x%llx)\n",
        data_info.start, data_info.end, data_info.size);
    fprintf(log_file, "        Original: 0x%llx\n\n", data_info.original_start);

    fprintf(log_file, "$gameSettings = [\n");

    for (const char* string_to_search : gameSettings) {
        unsigned long long string_addr = find_string_in_section(hProcess, rdata_info, string_to_search);

        if (string_addr != 0) {
            unsigned long long reference_addr = find_reference_in_data_section(hProcess, data_info, string_addr);

            if (reference_addr != 0) {
                // Calculate relative offset from the base address and adjust to match IDA's base
                unsigned long long relative_string_offset = string_addr - base_addr;
                unsigned long long relative_reference_offset = reference_addr - base_addr;

                unsigned long long ida_string_addr = 0x140000000 + relative_string_offset;
                unsigned long long ida_reference_addr = 0x140000000 + relative_reference_offset;

                fprintf(log_file, "    \"%s\" => 0x%llx, // String at 0x%llx\n",
                    string_to_search,
                    ida_reference_addr,
                    ida_string_addr);
            }
        }
    }

    fprintf(log_file, "];\n");

    CloseHandle(hProcess);
    fclose(log_file);
    return 0;
}