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

// Constants
const std::string PLUGIN_NAME = "UnchartedLOTC";
const std::string PLUGIN_LOG = PLUGIN_NAME + ".log";
std::string previousGameExecutable = "";
std::string gameExecutable = "";

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

// Plugin states
static bool AOBScanDone = false;
static bool g_fix_enabled = false;
static bool g_fov_fix_enabled = false;
static bool g_photo_fix_enabled = false;
static bool g_cinematics_fix_enabled = false;
static bool g_aspect_ratio_fix_enabled = false;
static bool g_DOF_fix_enabled = false;
static bool g_CA_fix_enabled = false; // Chromatic aberrations
static bool g_BD_fix_enabled = false; // Barell distortion
static int g_AdditionalValue = 0;
static int g_AdditionalPhotoValue = 0;
static float g_AspectValue = 0;

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

// AOB Scan pointers
static uint8_t* FOVaddress = nullptr;
static uint8_t* PHOTOaddress = nullptr;
static uint8_t* CINEMATICSASPECTaddress = nullptr;
static uint8_t* Aspectaddress = nullptr;
static uint8_t* DOFaddress = nullptr;
static uint8_t* EFFECTSaddress = nullptr; // Base effects
static uint8_t* CAaddress = nullptr; // Base effects
static uint8_t* BDaddress_1 = nullptr; // Barrel distortions
static uint8_t* BDaddress_2 = nullptr; // Barrel distortions

// Hooking
static SafetyHookMid FOVHook{};
static SafetyHookMid PhotoFOVHook{};
static SafetyHookMid AspectRatioHook{};
static SafetyHookMid DOFHook{};

// Prototypes 
static void FOVFixEnabled(bool fix_enabled);
static void PhotoFOVFixEnabled(bool fix_enabled);
static void CinematicsARFixEnabled(bool fix_enabled);
static void AspectRatioFixEnabled(bool fix_enabled);
static void DOFFixEnabled(bool fix_enabled);
static void CAFixEnabled(bool fix_enabled);
static void BDFixEnabled(bool fix_enabled);

// Original byte
static std::vector<BYTE> originalBytes = std::vector<BYTE>();

static std::string GetExecutableFileNameOnly() {
    char path[MAX_PATH];
    GetModuleFileNameA(nullptr, path, MAX_PATH);
    return std::filesystem::path(path).filename().string();
}

extern "C" __declspec(dllexport) void SetFixEnabled(bool enabled, bool init)
{
    g_fix_enabled = enabled;
    gameExecutable = GetExecutableFileNameOnly();

    if (g_fix_enabled && !AOBScanDone) {
        logger->info("Module {} detected", gameExecutable);
        logger->info("--------------- AOB scan started ---------------");
        if (FOVaddress == nullptr) {
            if (gameExecutable == "u4.exe") {
                auto FOVStringObfuscated = make_obfuscated<0x4A>("C5 FA ?? ?? ?? ?? ?? ?? C5 F2 ?? ?? ?? ?? ?? ?? C5 7A ?? ?? ?? ?? ?? ?? C5 7A");
                FOVaddress = Memory::AOBScan(gameExecutable, FOVStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            }
            else {
                auto FOVStringObfuscated = make_obfuscated<0x4A>("C5 FA ?? ?? ?? ?? ?? ?? C5 F2 ?? ?? ?? ?? ?? ?? E8 ?? ?? ?? ?? C5 92");
                FOVaddress = Memory::AOBScan(gameExecutable, FOVStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            }

            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));
            }
        }
        if (PHOTOaddress == nullptr) {
            auto PhotoFOVStringObfuscated = make_obfuscated<0x4A>("C5 FA ?? ?? ?? ?? ?? ?? C5 FA ?? ?? ?? ?? ?? ?? 41 ?? ?? ?? ?? ?? 48 8B ?? ?? ?? ?? ?? 48 83");
            PHOTOaddress = Memory::AOBScan(gameExecutable, PhotoFOVStringObfuscated.decrypt(), PAGE_EXECUTE_READ);

            if (!PHOTOaddress)
                logger->warn("Photo FOV signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                logger->info("Photo FOV signature found at address: 0x{:X}.", reinterpret_cast<uintptr_t>(PHOTOaddress));
                PHOTOaddress += 0x8;
            }
        }        
        if (CINEMATICSASPECTaddress == nullptr) {
            constexpr auto CinematicsAspectStringObfuscated = make_obfuscated<0x4A>("48 8B ?? ?? ?? ?? ?? 4C 8B ?? ?? ?? ?? ?? ?? 4C 8B ?? 49 C1");
            CINEMATICSASPECTaddress = Memory::AOBScan(gameExecutable, CinematicsAspectStringObfuscated.decrypt(), PAGE_EXECUTE_READ);

            if (!CINEMATICSASPECTaddress)
                logger->warn("Cinematics aspect ratio signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                logger->info("Cinematics aspect ratio signature found at address: 0x{:X}.", reinterpret_cast<uintptr_t>(CINEMATICSASPECTaddress));
            }
        }
        if (Aspectaddress == nullptr) {
            constexpr auto AspectStringObfuscated = make_obfuscated<0x4A>("44 8B ?? ?? ?? ?? ?? 48 ?? ?? 48 89 ?? ?? ?? ?? ?? E8 ?? ?? ?? ?? C5");
            Aspectaddress = Memory::AOBScan(gameExecutable, AspectStringObfuscated.decrypt(), PAGE_EXECUTE_READ);

            if (!Aspectaddress)
                logger->warn("Aspect ratio signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                logger->info("Aspect ratio signature found at address: 0x{:X}.", reinterpret_cast<uintptr_t>(Aspectaddress));
                Aspectaddress += 0x16;
            }
        }
        if (DOFaddress == nullptr) { 
            constexpr auto DOFStringObfuscated = make_obfuscated<0x4A>("41 8B 87 ?? ?? ?? ?? 41 89 86 ?? ?? ?? ?? 41 C7 86");
            DOFaddress = Memory::AOBScan(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));
            }
        }
        if (EFFECTSaddress == nullptr) {
            if (gameExecutable == "u4.exe") {
                auto EffectsStringObfuscated = make_obfuscated<0x4A>("41 8B 86 ?? ?? ?? ?? 89 86 ?? ?? ?? ?? 41 8B 86 ?? ?? ?? ?? 89 86 ?? ?? ?? ?? 41 8B 86 ?? ?? ?? ?? 89 86 ?? ?? ?? ?? 41 8B 87");
                EFFECTSaddress = Memory::AOBScan(gameExecutable, EffectsStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            }
            else {
                auto EffectsStringObfuscated = make_obfuscated<0x4A>("41 8B 87 ?? ?? ?? ?? 89 86 ?? ?? ?? ?? 41 8B 87 ?? ?? ?? ?? 89 86 ?? ?? ?? ?? 41 8B 87 ?? ?? ?? ?? 89 86 ?? ?? ?? ?? 41 8B 86");
                EFFECTSaddress = Memory::AOBScan(gameExecutable, EffectsStringObfuscated.decrypt(), PAGE_EXECUTE_READ);
            }

            if (!EFFECTSaddress)
                logger->warn("Effects signature not found. Maybe your game has been updated and is no more compatible with this plugin.");
            else {
                BDaddress_1 = EFFECTSaddress;
                BDaddress_2 = EFFECTSaddress + 0xd;
                CAaddress = EFFECTSaddress + 0x1a;
                logger->info("Effects signature found at address: 0x{:X}, 0x{:X}, 0x{:X}.", reinterpret_cast<uintptr_t>(EFFECTSaddress), reinterpret_cast<uintptr_t>(BDaddress_2), reinterpret_cast<uintptr_t>(CAaddress));
            }
            if (FOVaddress && PHOTOaddress && CINEMATICSASPECTaddress && Aspectaddress && DOFaddress && EFFECTSaddress)
                logger->info("All AOB signatures found. Ready to patch...");
        }
        if (FOVaddress && PHOTOaddress && CINEMATICSASPECTaddress && Aspectaddress && DOFaddress && EFFECTSaddress) AOBScanDone = true;
        logger->info("--------------- AOB scan finished ---------------");
    }
    if (g_fix_enabled) {
        if (FOVaddress) FOVFixEnabled(g_fov_fix_enabled);
        if (PHOTOaddress) PhotoFOVFixEnabled(g_photo_fix_enabled);
        if (CINEMATICSASPECTaddress) CinematicsARFixEnabled(g_cinematics_fix_enabled);
        if (Aspectaddress) AspectRatioFixEnabled(g_aspect_ratio_fix_enabled);
        if (DOFaddress) DOFFixEnabled(g_DOF_fix_enabled);
        if (CAaddress) CAFixEnabled(g_CA_fix_enabled);
        if (BDaddress_1 && BDaddress_2) BDFixEnabled(g_BD_fix_enabled);
    }
    else {
        if (FOVaddress) FOVFixEnabled(false);
        if (PHOTOaddress) PhotoFOVFixEnabled(false);
        if (CINEMATICSASPECTaddress) CinematicsARFixEnabled(false);
        if (Aspectaddress) AspectRatioFixEnabled(false);
        if (DOFaddress) DOFFixEnabled(false);
        if (CAaddress) CAFixEnabled(false);
        if (BDaddress_1 && BDaddress_2) BDFixEnabled(false);
        logger->info("All fixes disabled.");
    }
}

// Setters for Reshade addon call
extern "C" __declspec(dllexport) void SetFOVFixEnabled(bool fix_enabled, bool init)
{
    g_fov_fix_enabled = fix_enabled;
    if (!init) FOVFixEnabled(g_fov_fix_enabled);
}

extern "C" __declspec(dllexport) void SetPhotoFOVFixEnabled(bool fix_enabled, bool init)
{
    g_photo_fix_enabled = fix_enabled;
    if (!init) PhotoFOVFixEnabled(g_photo_fix_enabled); // FOV fix must be enabled when aspect ratio is too to compensate FOV
}

extern "C" __declspec(dllexport) void SetCinematicsARFixEnabled(bool fix_enabled, bool init)
{
    g_cinematics_fix_enabled = fix_enabled;
    if (!init) CinematicsARFixEnabled(g_cinematics_fix_enabled);
}

extern "C" __declspec(dllexport) void SetAspectRatioFixEnabled(bool fix_enabled, bool init)
{
    g_aspect_ratio_fix_enabled = fix_enabled;
    if (!init) AspectRatioFixEnabled(g_aspect_ratio_fix_enabled);
}

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

extern "C" __declspec(dllexport) void SetCAFixEnabled(bool fix_enabled, bool init)
{
    g_CA_fix_enabled = fix_enabled;
    if (!init) CAFixEnabled(g_CA_fix_enabled);
}

extern "C" __declspec(dllexport) void SetBDFixEnabled(bool fix_enabled, bool init)
{
    g_BD_fix_enabled = fix_enabled;
    if (!init) BDFixEnabled(g_BD_fix_enabled);
}

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

extern "C" __declspec(dllexport) void SetPhotoFOV(int fov)
{
    g_AdditionalPhotoValue = fov;
}

extern "C" __declspec(dllexport) void SetAspect(float aspect)
{
    g_AspectValue = aspect;
}

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

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

// Injection functions
static void FOVFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && FOVaddress) {
        if (!FOVHook) { // Hook only once
            FOVHook = safetyhook::create_mid(FOVaddress,
                [](SafetyHookContext& ctx) {
                    g_FOV_In = ctx.xmm0.f32[0];
                    g_FOV_Out = ctx.xmm0.f32[0] += g_AdditionalValue;
                });
        }
        else FOVHook.enable();
        logger->info("FOV fix enabled");
    }
    if (!fix_enabled && FOVaddress && FOVHook) {
        FOVHook.disable();
        logger->info("FOV fix disabled");
    }
}

static void PhotoFOVFixEnabled (bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && PHOTOaddress) {
        if (!PhotoFOVHook) { // Hook only once
            PhotoFOVHook = safetyhook::create_mid(PHOTOaddress,
                [](SafetyHookContext& ctx) {
                    ctx.xmm0.f32[0] += g_AdditionalPhotoValue - (g_fov_fix_enabled ? g_AdditionalValue : 0);
                });
        }
        else {
            PhotoFOVHook.enable();
        }
        logger->info("Photo FOV fix enabled");
    }
    if (!fix_enabled && PHOTOaddress && PhotoFOVHook) {
        PhotoFOVHook.disable();
        logger->info("Photo FOV fix disabled");
    }
}

static void CinematicsARFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && CINEMATICSASPECTaddress) {
        if (originalBytes.empty()) {
            originalBytes.resize(1);
            std::copy_n(CINEMATICSASPECTaddress + 3, 1, originalBytes.begin()); // Copy 1 byte of original offset
            originalBytes[0] -= 0x10; // Decrease it by 0x10
            Memory::PatchBytes(CINEMATICSASPECTaddress + 3, reinterpret_cast<const char*>(originalBytes.data()), 1); // mov rcx,[u4.exe+653E080 - 0x10]
            logger->info("Cinematics aspect ratio fix enabled");
        }
    }
    if (!fix_enabled && CINEMATICSASPECTaddress) {
        originalBytes.clear();
        Memory::RestoreBytes(CINEMATICSASPECTaddress + 3);
        logger->info("Cinematics aspect ratio fix disabled");
    }
}

static void AspectRatioFixEnabled(bool fix_enabled) { // OK ?
    if (g_fix_enabled && fix_enabled && Aspectaddress) {
        if (!AspectRatioHook) {
            AspectRatioHook = safetyhook::create_mid(Aspectaddress,
                [](SafetyHookContext& ctx) {
                    ctx.xmm0.f32[0] = g_AspectValue;
                });
        }
        else {
            AspectRatioHook.enable();
        }
        logger->info("Aspect ratio fix enabled");
    }
    if (!fix_enabled && Aspectaddress && AspectRatioHook) {
        AspectRatioHook.disable();
        logger->info("Aspect ratio fix disabled");
    }
}

static void DOFFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && DOFaddress) {
        Memory::PatchBytes(DOFaddress, "\xB8\x00\x00\x80\x7F\x90\x90", 7); // mov eax,0x7F800000
        logger->info("Depth of field fix enabled");
    }
    if (!fix_enabled && DOFaddress) {
        Memory::RestoreBytes(DOFaddress);
        logger->info("Depth of field fix disabled");
    }
}

static void CAFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && CAaddress) {
        Memory::PatchBytes(CAaddress, "\x31\xC0\x90\x90\x90\x90\x90", 7); // xor eax,eax No chromatic aberrations
        logger->info("Chromatic aberrations fix enabled");
    }
    if (!fix_enabled && CAaddress) {
        Memory::RestoreBytes(CAaddress);
        logger->info("Chromatic aberrations fix disabled");
    }
}

static void BDFixEnabled(bool fix_enabled) {
    if (g_fix_enabled && fix_enabled && BDaddress_1 && BDaddress_2) {
        Memory::PatchBytes(BDaddress_1, "\x31\xC0\x90\x90\x90\x90\x90", 7); // xor eax,eax No barrel distortions
        Memory::PatchBytes(BDaddress_2, "\xB8\x00\x00\x80\x3F\x90\x90", 7); // mov eax,0x3f800000 No barrel distortions
        logger->info("Barrel distortion fix enabled");
    }
    if (!fix_enabled && BDaddress_1) {
        Memory::RestoreBytes(BDaddress_1);
        Memory::RestoreBytes(BDaddress_2);
        logger->info("Barrel distortion fix disabled");
    }
}

// Initialisation de spdlog avec format personnalisé
static void InitializeLogger()
{
    try
    {
        logger = spdlog::basic_logger_mt("Fixlib", PLUGIN_LOG, true);
        spdlog::set_default_logger(logger);
        // Format : [YYYY-MM-DD HH:MM:SS] [INFO] message
        spdlog::set_pattern("[%Y-%m-%d %H:%M:%S] [%^%l%$] %v");
        spdlog::set_level(spdlog::level::debug);
        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)
    {
        //CinematicsDOFFixEnabled(false); // Core is loaded 3 times so i need to restore previous patched bytes
        logger->info("Plugin {} unloaded.", PLUGIN_NAME);
        spdlog::drop_all();
    }
    return TRUE;
}