#include "Memory.hpp";
#include "Maths.hpp";
#include "ObfuscateString.h"
#include <string>
#include <spdlog/spdlog.h>
#include <spdlog/sinks/rotating_file_sink.h>
#include <filesystem>

// Constants
const std::string PLUGIN_NAME = "WuchangFF";
const std::string PLUGIN_LOG = PLUGIN_NAME + ".log";
const std::string gameExecutable = "Project_Plague-Win64-Shipping.exe";
const float defaultAspect = 1.777778; // Default aspect that must not be compensated because FOV is already HOR+

// Logger
std::shared_ptr<spdlog::logger> logger;

// Screen informations
static int screenWidth = GetSystemMetrics(SM_CXSCREEN);
static int screenHeight = GetSystemMetrics(SM_CYSCREEN);
static float aspectRatio = (float)screenWidth / screenHeight;
static float baseAspect = 2.2222223; // Base game aspect ratio

// Plugin states
static bool AOBScanDone = false;
static bool g_fix_enabled = false;
static bool g_fov_fix_enabled = false;
static bool g_aspect_ratio_fix_enabled = false;
static bool g_DOF_fix_enabled = false;
static bool g_Vignetting_fix_enabled = false;
static bool g_Fog_fix_enabled = false;
static int g_AdditionalFOVValue = 0;

// Shared values
static float g_FOV_In = 0;
static float g_Compensated_FOV = 0;
static float g_FOV_Out = 0;

// AOB Scan pointers
static uint8_t* FOVaddress = nullptr;
static uint8_t* Aspectaddress_1 = nullptr;
static uint8_t* Aspectaddress_2 = nullptr;
static uint8_t* DOFaddress = nullptr;
static uint8_t* Vignettingaddress = nullptr;
static uint8_t* Fogaddress = nullptr;

// FOV Addresses
uintptr_t t_fov_address_to_patch = 0;
uintptr_t t_aspect_address_to_patch = 0;

// VEH debugger handles
PVOID vehFOVHandle = nullptr;
PVOID vehApectHandle_1 = nullptr;
PVOID vehApectHandle_2 = nullptr;

// Prototypes 
static void FOVFixEnabled(bool fix_enabled);
static void AspectFixEnabled(bool fix_enabled);
static void DOFFixEnabled(bool fix_enabled);
static void VignettingFixEnabled(bool fix_enabled);
static void FogFixEnabled(bool fix_enabled);

extern "C" __declspec(dllexport) void SetFixEnabled(bool enabled)
{
    g_fix_enabled = enabled;
    if (g_fix_enabled && !AOBScanDone) {
        logger->info("--------------- AOB scan started ---------------");
        if (FOVaddress == nullptr) {            
            constexpr auto FOVStringObfuscated = make_obfuscated<0x4A>("EB ?? F3 0F ?? ?? ?? ?? ?? ?? F3 0F ?? ?? ?? 0F ?? ?? 8B");
            FOVaddress = Memory::AOBScan(gameExecutable, FOVStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            //"Project_Plague-Win64-Shipping.exe" + 3E1037A - EB 08                   - jmp "Project_Plague-Win64-Shipping.exe" + 3E10384
            //"Project_Plague-Win64-Shipping.exe" + 3E1037C - F3 0F 10 83 E0 02 00 00 - movss xmm0, [rbx + 000002E0]
            //"Project_Plague-Win64-Shipping.exe" + 3E10384 - F3 0F 11 47 30          - movss[rdi + 30], xmm0
            //"Project_Plague-Win64-Shipping.exe" + 3E10389 - 0F 57 C0                - xorps xmm0, xmm0
            //"Project_Plague-Win64-Shipping.exe" + 3E1038C - 8B 83 F0 02 00 00       - mov eax, [rbx + 000002F0]

            if (!FOVaddress)
                logger->warn("FOV signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                logger->info("FOV signature found at address: 0x{:X}.", reinterpret_cast<uintptr_t>(FOVaddress));
                FOVaddress += 0xa; // Offset for the target opcode
            }
        }

        if (Aspectaddress_1 == nullptr) {
            constexpr auto AspectStringObfuscated = make_obfuscated<0x4A>("C3 CC CC CC CC CC CC CC CC 48 89 ?? ?? ?? 57 48 ?? ?? ?? 0F ?? ?? 48 ?? ?? 48 ?? ?? 0F ?? ?? ?? 0F ?? ?? ?? F2 0F");
            Aspectaddress_1 = Memory::AOBScan(gameExecutable, AspectStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            //"Project_Plague-Win64-Shipping.exe" + 10A7443 - 89 41 44 - mov[rcx + 44], eax
            //"Project_Plague-Win64-Shipping.exe" + 10A7446 - 8B 42 48 - mov eax, [rdx + 48]
            //"Project_Plague-Win64-Shipping.exe" + 10A7449 - 89 41 48 - mov[rcx + 48], eax
            //"Project_Plague-Win64-Shipping.exe" + 10A744C - 8B 49 4C - mov ecx, [rcx + 4C]
            //"Project_Plague-Win64-Shipping.exe" + 10A744F - 33 4A 4C - xor ecx, [rdx + 4C]

            if (!Aspectaddress_1)
                logger->warn("Aspect ratio #1 signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                logger->info("Aspect ratio #1 signature found at address: 0x{:X}.", reinterpret_cast<uintptr_t>(Aspectaddress_1));
                Aspectaddress_1 += 0x62;
            }
        }

        if (Aspectaddress_2 == nullptr) {
            constexpr auto AspectStringObfuscated = make_obfuscated<0x4A>("C3 CC CC CC CC CC 48 89 ?? ?? ?? 57 48 ?? ?? ?? 0F ?? ?? 48 ?? ?? 48 ?? ?? 0F ?? ?? ?? 0F ?? ?? ?? F2 0F");
            Aspectaddress_2 = Memory::AOBScan(gameExecutable, AspectStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            //"Project_Plague-Win64-Shipping.exe" + 1914483 - 89 41 44 - mov[rcx + 44], eax
            //"Project_Plague-Win64-Shipping.exe" + 1914486 - 8B 42 48 - mov eax, [rdx + 48]
            //"Project_Plague-Win64-Shipping.exe" + 1914489 - 89 41 48 - mov[rcx + 48], eax
            //"Project_Plague-Win64-Shipping.exe" + 191448C - 8B 49 4C - mov ecx, [rcx + 4C]
            //"Project_Plague-Win64-Shipping.exe" + 191448F - 33 4A 4C - xor ecx, [rdx + 4C]

            if (!Aspectaddress_2)
                logger->warn("Aspect ratio #2 signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                logger->info("Aspect ratio #2 signature found at address: 0x{:X}.", reinterpret_cast<uintptr_t>(Aspectaddress_2));
                Aspectaddress_2 += 0x5F;
            }
        }

        if (DOFaddress == nullptr) {
            constexpr auto DOFStringObfuscated = make_obfuscated<0x4A>("48 ?? ?? 8B ?? ?? E8 ?? ?? ?? ?? 48 ?? ?? 48 6B ?? ?? 48 8D");
            DOFaddress = Memory::AOBScan(gameExecutable, DOFStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            //"Project_Plague-Win64-Shipping.exe" + 2BD06A4 - 48 8B 05 1D 15 9D 05 - mov rax, ["Project_Plague-Win64-Shipping.exe" + 85A1BC8]
            //"Project_Plague-Win64-Shipping.exe" + 2BD06AB - 48 8B CB             - mov rcx, rbx
            //"Project_Plague-Win64-Shipping.exe" + 2BD06AE - 8B 78 04             - mov edi, [rax + 04]
            //"Project_Plague-Win64-Shipping.exe" + 2BD06B1 - E8 AA 5F 8A 01       - call "Project_Plague-Win64-Shipping.exe" + 4476660
            //"Project_Plague-Win64-Shipping.exe" + 2BD06B6 - 48 63 C8             - movsxd  rcx, eax

            if (!DOFaddress)
                logger->warn("DOF signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                logger->info("DOF signature found at address: 0x{:X}.", reinterpret_cast<uintptr_t>(DOFaddress));
                DOFaddress += 0x3;
            }
        }

        if (Vignettingaddress == nullptr) {
            constexpr auto VignettingStringObfuscated = make_obfuscated<0x4A>("8B ?? 83 ?? ?? 7D ?? 89 B3 ?? ?? ?? ?? EB");
            Vignettingaddress = Memory::AOBScan(gameExecutable, VignettingStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            //"Project_Plague-Win64-Shipping.exe" + 4474088 - C7 83 B0 1F 00 00 00 00 80 3F - mov[rbx + 00001FB0], 3F800000
            //"Project_Plague-Win64-Shipping.exe" + 4474092 - 48 8B 05 87 FE 1F 04          - mov rax, ["Project_Plague-Win64-Shipping.exe" + 8673F20]
            //"Project_Plague-Win64-Shipping.exe" + 4474099 - 8B 08                         - mov ecx, [rax]
            //"Project_Plague-Win64-Shipping.exe" + 447409B - 83 F9 02                      - cmp ecx, 02
            //"Project_Plague-Win64-Shipping.exe" + 447409E - 7D 08                         - jnl "Project_Plague-Win64-Shipping.exe" + 44740A8

            if (!Vignettingaddress)
                logger->warn("Vignetting signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                logger->info("Vignetting signature found at address: 0x{:X}.", reinterpret_cast<uintptr_t>(Vignettingaddress));
            }
        }

        if (Fogaddress == nullptr) {
            constexpr auto FogStringObfuscated = make_obfuscated<0x4A>("75 ?? B3 ?? EB ?? 32 ?? 48 8B ?? ?? ?? 48 ?? ?? 74 ?? E8 ?? ?? ?? ?? 0F ?? ?? 48 8B");
            Fogaddress = Memory::AOBScan(gameExecutable, FogStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            //"Project_Plague-Win64-Shipping.exe" + 28E27FD - 48 8B 05 54 9A C6 05 - mov rax, ["Project_Plague-Win64-Shipping.exe" + 854C258]
            //"Project_Plague-Win64-Shipping.exe" + 28E2804 - 83 78 04 01          - cmp dword ptr[rax + 04], 01
            //"Project_Plague-Win64-Shipping.exe" + 28E2808 - 75 04                - jne "Project_Plague-Win64-Shipping.exe" + 28E280E
            //"Project_Plague-Win64-Shipping.exe" + 28E280A - B3 01                - mov bl, 01
            //"Project_Plague-Win64-Shipping.exe" + 28E280C - EB 02                - jmp "Project_Plague-Win64-Shipping.exe" + 28E2810

            if (!Fogaddress)
                logger->warn("Fog signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                logger->info("Fog signature found at address: 0x{:X}.", reinterpret_cast<uintptr_t>(Fogaddress));
            }            
        }

        if (FOVaddress && Aspectaddress_1 && Aspectaddress_2 && DOFaddress && Vignettingaddress && Fogaddress) {
            logger->info("All AOB signatures found. Ready to patch...");
            AOBScanDone = true;
        }

        logger->info("--------------- AOB scan finished ---------------");
    }

    if (g_fix_enabled) {
        if (FOVaddress) FOVFixEnabled(g_fov_fix_enabled || g_aspect_ratio_fix_enabled);
        if (Aspectaddress_1 && Aspectaddress_2) AspectFixEnabled(g_aspect_ratio_fix_enabled);
        if (DOFaddress) DOFFixEnabled(g_DOF_fix_enabled);
        if (Vignettingaddress) VignettingFixEnabled(g_Vignetting_fix_enabled);
        if (Fogaddress) FogFixEnabled(g_Fog_fix_enabled);
    }
    else {
        if (FOVaddress) FOVFixEnabled(false);
        if (Aspectaddress_1 && Aspectaddress_2) AspectFixEnabled(false);
        if (DOFaddress) DOFFixEnabled(false);
        if (Vignettingaddress) VignettingFixEnabled(false);
        if (Fogaddress) FogFixEnabled(false);
        logger->info("All fixes disabled.");
    }
}

extern "C" __declspec(dllexport) void SetFOVFixEnabled(bool enabled, bool init)
{
    g_fov_fix_enabled = enabled;
    if (!init) FOVFixEnabled(g_fov_fix_enabled || g_aspect_ratio_fix_enabled); // FOV fix must be enabled when aspect ratio is too to compensate FOV
}

// Setters for Reshade addon call
extern "C" __declspec(dllexport) void SetAspectRatioFixEnabled(bool enabled, bool init)
{
    g_aspect_ratio_fix_enabled = enabled;
    if (!init) AspectFixEnabled(g_aspect_ratio_fix_enabled);
}

extern "C" __declspec(dllexport) void SetDOFFixEnabled(bool enabled, bool init)
{
    g_DOF_fix_enabled = enabled;
    if (!init) DOFFixEnabled(g_DOF_fix_enabled);
}

extern "C" __declspec(dllexport) void SetVignettingFixEnabled(bool enabled, bool init)
{
    g_Vignetting_fix_enabled = enabled;
    if (!init) VignettingFixEnabled(g_Vignetting_fix_enabled);
}

extern "C" __declspec(dllexport) void SetFogFixEnabled(bool enabled, bool init)
{
    g_Fog_fix_enabled = enabled;
    if (!init) FogFixEnabled(g_Fog_fix_enabled);
}

extern "C" __declspec(dllexport) void SetFOV(int fov)
{
    g_AdditionalFOVValue = fov;
}

// Getters for Reshade addon call
extern "C" __declspec(dllexport) float GetFOVIn() {
    return g_FOV_In;
}

extern "C" __declspec(dllexport) float GetCompensatedFOV() {
    return g_Compensated_FOV;
}

extern "C" __declspec(dllexport) float GetFOVOut() {
    return g_FOV_Out;
}

// VEH debugger breakpoint fixes
LONG WINAPI FOVVehHandler(PEXCEPTION_POINTERS ep) {
    if (!ep || !ep->ContextRecord || !ep->ExceptionRecord) return EXCEPTION_CONTINUE_SEARCH;
    if (ep->ExceptionRecord->ExceptionCode != EXCEPTION_SINGLE_STEP) return EXCEPTION_CONTINUE_SEARCH;

    CONTEXT* ctx = ep->ContextRecord; // Check if instruction breakpoint is the one targeted
    if (ctx->Rip == reinterpret_cast<uintptr_t>(FOVaddress)) {
        t_fov_address_to_patch = static_cast<uintptr_t>(ctx->Rdi) + 0x30;

        float* xmm0 = reinterpret_cast<float*>(&ctx->Xmm0);
        g_FOV_In = *xmm0;
        // Calculate compensated FOV if aspect fix is checked
        if (g_aspect_ratio_fix_enabled && baseAspect > defaultAspect)
            g_Compensated_FOV = *xmm0 = Maths::CompensateHorizontalFOV(g_FOV_In, baseAspect, aspectRatio);
        else
            g_Compensated_FOV = *xmm0;

        *xmm0 += (g_fov_fix_enabled ? g_AdditionalFOVValue : 0);
        g_FOV_Out = *xmm0;

        *reinterpret_cast<float*>(t_fov_address_to_patch) = g_FOV_Out; // Patch new FOV value in memory

        // Skip original opcode size in bytes otherwise the game will crash
        ctx->Rip += 5;

        return EXCEPTION_CONTINUE_EXECUTION;
    }

    return EXCEPTION_CONTINUE_SEARCH;
}

LONG WINAPI AspectVehHandler_1(PEXCEPTION_POINTERS ep) {
    if (!ep || !ep->ContextRecord || !ep->ExceptionRecord) return EXCEPTION_CONTINUE_SEARCH;
    if (ep->ExceptionRecord->ExceptionCode != EXCEPTION_SINGLE_STEP) return EXCEPTION_CONTINUE_SEARCH;

    CONTEXT* ctx = ep->ContextRecord; // Check if instruction breakpoint is the one targeted
    if (ctx->Rip == reinterpret_cast<uintptr_t>(Aspectaddress_1)) {
        // Retrieves the addresse where aspect ratio will be stored
        t_aspect_address_to_patch = static_cast<uintptr_t>(ctx->Rcx) + 0x48;
        // Retrieves where the base aspect ratio is actually stored
        baseAspect = *reinterpret_cast<float*>(&ctx->Rax);
        *reinterpret_cast<float*>(t_aspect_address_to_patch) = baseAspect; // Patch base aspect ratio

        // Skip original opcode size in bytes otherwise the game will crash
        ctx->Rip += 3;

        return EXCEPTION_CONTINUE_EXECUTION;
    }

    return EXCEPTION_CONTINUE_SEARCH;
}

LONG WINAPI AspectVehHandler_2(PEXCEPTION_POINTERS ep) {
    if (!ep || !ep->ContextRecord || !ep->ExceptionRecord) return EXCEPTION_CONTINUE_SEARCH;
    if (ep->ExceptionRecord->ExceptionCode != EXCEPTION_SINGLE_STEP) return EXCEPTION_CONTINUE_SEARCH;

    CONTEXT* ctx = ep->ContextRecord; // Check if instruction breakpoint is the one targeted
    if (ctx->Rip == reinterpret_cast<uintptr_t>(Aspectaddress_2)) {
        // Retrieves the addresse where aspect ratio will be stored
        t_aspect_address_to_patch = static_cast<uintptr_t>(ctx->Rcx) + 0x48;
        *reinterpret_cast<float*>(t_aspect_address_to_patch) = aspectRatio; // Patch new aspect ratio value in memory

        // Skip original opcode size in bytes otherwise the game will crash
        ctx->Rip += 3;

        return EXCEPTION_CONTINUE_EXECUTION;
    }

    return EXCEPTION_CONTINUE_SEARCH;
}

// Code injection functions
static void FOVFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && FOVaddress) {
        vehFOVHandle = nullptr;
        vehFOVHandle = Memory::SetupOrClearHardwareBreakPointForAllThreads(reinterpret_cast<uintptr_t>(FOVaddress), vehFOVHandle ,true, FOVVehHandler, 0);
        if (vehFOVHandle)
            logger->info("FOV fix enabled and hardware breakpoint armed");
    }
    if (!fix_enabled && FOVaddress) {
        if (vehFOVHandle)
            vehFOVHandle = Memory::SetupOrClearHardwareBreakPointForAllThreads(reinterpret_cast<uintptr_t>(FOVaddress), vehFOVHandle, false, nullptr, 0);
        if (!vehFOVHandle)
            logger->info("FOV fix disabled and hardware breakpoint disarmed");
    }
}

static void AspectFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && Aspectaddress_1 && Aspectaddress_2) {
        vehApectHandle_1 = nullptr;
        vehApectHandle_2 = nullptr;
        vehApectHandle_1 = Memory::SetupOrClearHardwareBreakPointForAllThreads(reinterpret_cast<uintptr_t>(Aspectaddress_1), vehApectHandle_1, true, AspectVehHandler_1, 1);
        vehApectHandle_2 = Memory::SetupOrClearHardwareBreakPointForAllThreads(reinterpret_cast<uintptr_t>(Aspectaddress_2), vehApectHandle_2, true, AspectVehHandler_2, 2);

        if (vehApectHandle_1 && vehApectHandle_2)
            logger->info("Aspect ratio fix enabled and hardware breakpoint armed");

        if (FOVaddress) FOVFixEnabled(true); // Compensate FOV if aspect fix is checked
    }
    if (!fix_enabled && Aspectaddress_1) {
        if (vehApectHandle_1)
            vehApectHandle_1 = Memory::SetupOrClearHardwareBreakPointForAllThreads(reinterpret_cast<uintptr_t>(Aspectaddress_1), vehApectHandle_1, false, nullptr, 1);

        if (vehApectHandle_2)
            vehApectHandle_2 = Memory::SetupOrClearHardwareBreakPointForAllThreads(reinterpret_cast<uintptr_t>(Aspectaddress_2), vehApectHandle_2, false, nullptr, 2);
        if (!vehApectHandle_1 && !vehApectHandle_2)
            logger->info("Aspect ratio fix disabled and hardware breakpoint disarmed");

        // When aspect fix is unchecked no more FOV compensation is needed
        if (FOVaddress && !g_fov_fix_enabled) FOVFixEnabled(false);
    }
}

// Memory patch fixes
static void DOFFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && DOFaddress) {
        Memory::PatchBytes(DOFaddress, "\x31\xFF\x90", 3); // xor edi,edi r.DepthOfFieldQuality = 0
        logger->info("Depth of field fix enabled");
    }
    if (!fix_enabled && DOFaddress) {
        Memory::RestoreBytes(DOFaddress);
        logger->info("Depth of field fix disabled");
    }
}

static void VignettingFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && Vignettingaddress) {
        Memory::PatchBytes(Vignettingaddress, "\x31\xC9", 2); // xor ecx,ecx r.Tonemapper.Quality=0
        logger->info("Vignetting fix enabled");
    }
    if (!fix_enabled && Vignettingaddress) {
        Memory::RestoreBytes(Vignettingaddress);
        logger->info("Vignetting fix disabled");
    }
}

static void FogFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && Fogaddress) {
        Memory::PatchBytes(Fogaddress, "\xEB", 1); // jmp "Project_Plague-Win64-Shipping.exe"+28E280E r.Fog = 0
        logger->info("Fog fix enabled");
    }
    if (!fix_enabled && Fogaddress) {
        Memory::RestoreBytes(Fogaddress);
        logger->info("Fog fix disabled");
    }
}

static void InitializeLogger()
{
    try
    {
        std::filesystem::path log_path = std::filesystem::absolute(PLUGIN_LOG);
        if (std::filesystem::exists(log_path))
            std::filesystem::remove(log_path);
        logger = std::make_shared<spdlog::logger>("Wuchang Fallen Feathers", std::make_shared<spdlog::sinks::rotating_file_sink_st>(PLUGIN_LOG, 10 * 1024 * 1024, 1));
        logger->flush_on(spdlog::level::debug); // Flush automatically
    }
    catch (const spdlog::spdlog_ex& ex)
    {
        std::string plugin_error_message = "Could not open " + PLUGIN_LOG;
        MessageBoxA(nullptr, plugin_error_message.c_str(), "Logger Error", MB_ICONERROR | MB_OK);
    }
}


// Standard dll entry
BOOL APIENTRY DllMain(HMODULE hModule, DWORD reason, LPVOID)
{
    if (reason == DLL_PROCESS_ATTACH)
    {
        InitializeLogger();
        logger->info("Plugin {} loaded.", PLUGIN_NAME);
    }
    else if (reason == DLL_PROCESS_DETACH)
    {
        logger->info("Plugin {} unloaded.", PLUGIN_NAME);
        spdlog::drop_all();
    }
    return TRUE;
}