#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";

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

// Plugin states
static bool AOBScanDone = false;
static bool g_fix_enabled = false;
static bool g_aspect_ratio_axis_constrain_fix_enabled = false;
static bool g_DOF_fix_enabled = false;
static bool g_Vignetting_fix_enabled = false;
static bool g_Fog_fix_enabled = false;

// AOB Scan pointers
static uint8_t* HORPLUSaddress = nullptr;
static uint8_t* DOFaddress = nullptr;
static uint8_t* Vignettingaddress = nullptr;
static uint8_t* Fogaddress = nullptr;

// Prototypes 
static void HORPlusFixEnabled(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 (HORPLUSaddress == nullptr) {
            constexpr auto HORPLUSStringObfuscated = make_obfuscated<0x4A>("41 0F ?? ?? ?? ?? ?? ?? 48 8D ?? ?? ?? ?? ?? 4C ?? ?? 4D ?? ?? E8");
            HORPLUSaddress = Memory::aob_scan(gameExecutable, HORPLUSStringObfuscated.decrypt(), PAGE_EXECUTE_READ);

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

        if (DOFaddress == nullptr) {
            constexpr auto DOFStringObfuscated = make_obfuscated<0x4A>("48 ?? ?? 8B ?? ?? E8 ?? ?? ?? ?? 48 ?? ?? 48 6B ?? ?? 48 8D");
            DOFaddress = Memory::aob_scan(gameExecutable, DOFStringObfuscated.decrypt(), PAGE_EXECUTE_READ);

            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::aob_scan(gameExecutable, VignettingStringObfuscated.decrypt(), PAGE_EXECUTE_READ);

            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::aob_scan(gameExecutable, FogStringObfuscated.decrypt(), PAGE_EXECUTE_READ);

            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 (HORPLUSaddress && DOFaddress && Vignettingaddress && Fogaddress) {
            logger->info("All AOB signatures found. Ready to patch...");
            AOBScanDone = true;
        }

        logger->info("--------------- AOB scan finished ---------------");
    }
    if (g_fix_enabled) {
        if (HORPLUSaddress) HORPlusFixEnabled(g_aspect_ratio_axis_constrain_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 (HORPLUSaddress) HORPlusFixEnabled(false);
        if (DOFaddress) DOFFixEnabled(false);
        if (Vignettingaddress) VignettingFixEnabled(false);
        if (Fogaddress) FogFixEnabled(false);
        logger->info("All fixes disabled.");
    }
}

// Setters for Reshade addon call
extern "C" __declspec(dllexport) void SetARAxisConstrainFixEnabled(bool enabled, bool init)
{
    g_aspect_ratio_axis_constrain_fix_enabled = enabled;
    if (!init) HORPlusFixEnabled(g_aspect_ratio_axis_constrain_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);
}

// Memory patch fixes
static void HORPlusFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && HORPLUSaddress) {
        Memory::PatchBytes(HORPLUSaddress, "\x31\xD2\x90\x90\x90\x90\x90\x90", 8); // xor edx,edx AspectRatioAxisConstraint=AspectRatio_MaintainYFOV
        logger->info("HOR+ fix enabled");
    }
    if (!fix_enabled && HORPLUSaddress) {
        Memory::RestoreBytes(HORPLUSaddress);
        logger->info("HOR+ fix disabled");
    }
}

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
        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;
}