#include "CommonHeaders.h"
#include "UEngine.hpp"
#include "Logger.hpp"

// Constants
const std::string PLUGIN_NAME = "DeathStranding2";
const std::string PLUGIN_LOG = PLUGIN_NAME + ".log";

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

// Screen informations
static int screenWidth = GetSystemMetrics(SM_CXSCREEN);
static int screenHeight = GetSystemMetrics(SM_CYSCREEN);
static float g_AspectRatio = (float)screenWidth / screenHeight;

// Plugin states
static bool AOBScanDone = false;
static std::atomic<bool> g_fix_enabled = false;
static std::atomic<bool> g_fov_fix_enabled = false;
static std::atomic<bool> g_ultrawide_fix_enabled = false;
static std::atomic<bool> g_HUD_fix_enabled = false;
static std::atomic<bool> g_camera_fix_enabled = false;
static std::atomic<bool> g_DOF_fix_enabled = false;
static std::atomic<bool> g_vignetting_fix_enabled = false;
static std::atomic<bool> g_sharpening_fix_enabled = false;
static int g_AdditionalFOVValue = 0;
static float g_cameraDistanceMultiplier = 1.f;
static float g_sharpeningValue = 1.f;
static float g_HUDOffset = 0.f;

// Shared values
static float g_FOV_In = 80.f;
static float g_FOV_Out = 80.f;
static float g_Camera_In = 3.f;
static float g_Camera_Out = 3.f;

// AOB Scan pointers
static uint8_t* FOVaddress = nullptr;
static uint8_t* Ultrawideaddress = nullptr;
static uint8_t* Cutscenesaddress = nullptr;
static uint8_t* HUD1address = nullptr;
static uint8_t* HUD2address = nullptr;
static uint8_t* Cameraaddress = nullptr;
static uint8_t* DOFaddress = nullptr;
static uint8_t* Vignetteaddress = nullptr;
static uint8_t* Sharpeningaddress = nullptr;
static uint8_t* PSNCheckaddress = nullptr;

// Hooking
static SafetyHookMid FOVHook{};
static SafetyHookMid CameraHook{};
static SafetyHookMid UWHook{};
static SafetyHookMid HUDHook{};
static SafetyHookMid VignetteHook{};
static SafetyHookMid SharpeningHook{};

// Prototypes
static void FOVFixEnabled();
static void UltraWideFixEnabled();
static void HUDFixEnabled();
static void CameraFixEnabled();
static void DOFFixEnabled();
static void VignetteFixEnabled();
static void SharpeningFixEnabled();
static void PSNCheckRemoval();

extern "C" __declspec(dllexport) void SetFixEnabled(bool enabled, bool init) {
    g_fix_enabled = enabled;
    if (!AOBScanDone) { // Decima Engine
        logger = InitializeLogger("Death Stranding 2: On The Beach", PLUGIN_LOG);
        logger->info("Plugin {} loaded.", PLUGIN_NAME);
        logger->info("--------------- AOB scan started ---------------");
        constexpr auto FOVStringObfuscated = make_obfuscated<0xFA>("C5 F8 ?? ?? ?? ?? C5 F8 ?? ?? ?? ?? C5 78 ?? ?? ?? ?? C5 FA ?? ?? C5 F2 59 ?? ?? ?? ?? ?? 48 83 ?? ?? C3"); // +0x1e
        constexpr auto AspectStringObfuscated = make_obfuscated<0x8B>("C5 FA 10 ?? ?? C5 F2 ?? ?? C5 CA ?? ?? E9 ?? ?? ?? ?? C5 FA 10 ?? ?? ?? ?? ?? C5");
        constexpr auto AspectCutscenesStringObfuscated = make_obfuscated<0x9D>("0F 84 ?? ?? ?? ?? 49 8B 9D ?? ?? ?? ?? 49 63 85 ?? ?? ?? ?? 48 6B ?? ?? 48 ?? ?? 48 ?? ?? 0F 84");
        constexpr auto HUD1StringObfuscated = make_obfuscated<0x3D>("C5 E8 ?? ?? ?? ?? ?? ?? ?? C5 F9 ?? ?? 25 ?? ?? ?? ?? 3D ?? ?? ?? ?? 74 ?? C5 F8 11 ?? ?? ?? ?? ?? 48 ?? ?? 48 ?? ?? FF ?? ?? 48 83");
        constexpr auto HUD2StringObfuscated = make_obfuscated<0x3D>("7A ?? 75 ?? C5 F8 ?? ?? ?? ?? ?? ?? 7A ?? 75 ?? 39 83 ?? ?? ?? ?? 0F 84");
        constexpr auto CameraStringObfuscated = make_obfuscated<0x84>("C5 FA 10 ?? ?? ?? ?? ?? 40 38 ?? ?? ?? ?? ?? 0F 84 ?? ?? ?? ?? C5 FA 10");
        constexpr auto DOFStringObfuscated = make_obfuscated<0x84>("48 83 ?? ?? C4 C1 ?? ?? ?? C5 FA ?? ?? C5 F8 ?? ?? C5 F8 ?? ?? ?? ?? C5 F8 ?? ?? ?? C5 F8 ?? ?? 76");
        constexpr auto VignetteStringObfuscated = make_obfuscated<0x83>("C5 FA ?? ?? ?? 48 83 ?? ?? 0F 85 ?? ?? ?? ?? C5 F8 ?? ?? ?? ?? 48 8B ?? ?? ?? ?? ?? ?? C6 43");
        constexpr auto SharpeningStringObfuscated = make_obfuscated<0x78>("48 8D ?? ?? ?? ?? ?? C5 ?? ?? E8 ?? ?? ?? ?? C5 78 28 ?? ?? ?? ?? ?? ?? C5 F8 28");
        constexpr auto PSNCheckStringObfuscated = make_obfuscated<0x8A>("E8 ?? ?? ?? ?? 48 89 ?? ?? ?? ?? ?? ?? 48 81 ?? ?? ?? ?? ?? E9 ?? ?? ?? ?? 0F 8A ?? ?? ?? ?? EB");

        using AOBScan::Make;
        using OffsetScan::Make;
        // Prepare all data for scanning
        std::vector<AOBScanEntry> signatures = {
            Make(&FOVaddress, FOVStringObfuscated, "FOV"),
            Make(&Ultrawideaddress, AspectStringObfuscated, "Ultrawide"),
            Make(&Cutscenesaddress, AspectCutscenesStringObfuscated, "Cutscenes"),
            Make(&HUD1address, HUD1StringObfuscated, "HUD"),
            Make(&HUD2address, HUD2StringObfuscated, "HUD"),
            Make(&Cameraaddress, CameraStringObfuscated, "Camera"),
            Make(&DOFaddress, DOFStringObfuscated, "Depth of field"),
            Make(&Vignetteaddress, VignetteStringObfuscated, "Vignette"),
            Make(&Sharpeningaddress, SharpeningStringObfuscated, "Sharpening"),
            #ifdef RELEASE_INTERNAL
                    Make(&PSNCheckaddress, PSNCheckStringObfuscated, "PSN check")
            #endif
        };
        // Scan all signature in a batch
        Memory::AOBScanBatch(signatures, logger);
        if (FOVaddress && Ultrawideaddress && HUD1address && HUD2address && Cutscenesaddress && Cameraaddress && DOFaddress && Vignetteaddress && Sharpeningaddress)
            logger->info("All AOB signatures found. Ready to patch...");

        logger->info("-------------- Fixes initialisation -------------");
        AOBScanDone = true;
    }

    PSNCheckRemoval();
    if (init) return;
    FOVFixEnabled();
    UltraWideFixEnabled();
    HUDFixEnabled();
    CameraFixEnabled();
    DOFFixEnabled();
    VignetteFixEnabled();
    SharpeningFixEnabled();
}

// Setters for Reshade addon call
extern "C" __declspec(dllexport) void SetFixesEnabled(GameFixes fix, bool enabled) { // Set each fix individually
    if (fix == GameFixes::FOV)                  { g_fov_fix_enabled = enabled; FOVFixEnabled(); }
    if (fix == GameFixes::UltraWide)            { g_ultrawide_fix_enabled = enabled; UltraWideFixEnabled(); }
    if (fix == GameFixes::HUD)                  { g_HUD_fix_enabled = enabled; HUDFixEnabled(); }
    if (fix == GameFixes::Camera)               { g_camera_fix_enabled = enabled; CameraFixEnabled(); }
    if (fix == GameFixes::DOF)                  { g_DOF_fix_enabled = enabled; DOFFixEnabled(); }
    if (fix == GameFixes::Vignetting)           { g_vignetting_fix_enabled = enabled; VignetteFixEnabled(); }
    if (fix == GameFixes::Sharpening)           { g_sharpening_fix_enabled = enabled; SharpeningFixEnabled(); }
    if (fix == GameFixes::None) logger->info("------------------ User inputs ------------------");
}

extern "C" __declspec(dllexport) void SetValues(GameSetting setting, float value) {
    if (setting == GameSetting::FOV) g_AdditionalFOVValue = (int)(value);
    if (setting == GameSetting::CameraDistance) g_cameraDistanceMultiplier = value;
    if (setting == GameSetting::HUD) g_HUDOffset = (value * screenWidth) / (float)100;
    if (setting == GameSetting::Sharpening) g_sharpeningValue = 1.f - value / 1.5f;
}

// Getters for Reshade addon call
extern "C" __declspec(dllexport) void GetGameInfos(GameInfos* infos) {
    if (!infos) return;

    infos->FOVIn = g_FOV_In;
    infos->FOVOut = g_FOV_Out;
    infos->cameraIn = g_Camera_In;
    infos->cameraOut = g_Camera_Out;
    infos->screenWidth = screenWidth;
    infos->screenHeight = screenHeight;
    infos->aspectRatio = (float)screenWidth / screenHeight;
}

// Code injection functions
static void FOVFixEnabled() {
	if (!FOVaddress) return;
    if (g_fix_enabled && g_fov_fix_enabled) {
        if (!FOVHook) { // Hook only once
            FOVHook = safetyhook::create_mid(FOVaddress + 0x1e,
                [](SafetyHookContext& ctx) {
                    g_FOV_In = ctx.xmm0.f32[0];
                    ctx.xmm0.f32[0] += (g_fix_enabled && g_fov_fix_enabled ? g_AdditionalFOVValue : 0);
                    g_FOV_Out = ctx.xmm0.f32[0];
                });
        }
		else FOVHook.enable();
    }
    else if (FOVHook) FOVHook.disable();

  logger->info("FOV fix {}", g_fix_enabled && g_fov_fix_enabled ? "enabled" : "disabled");
}

static void UltraWideFixEnabled() {
    if (!Ultrawideaddress || !Cutscenesaddress) return;
    if (g_fix_enabled && g_ultrawide_fix_enabled) {
        if (!UWHook) {
            UWHook = safetyhook::create_mid(Ultrawideaddress,
                [](SafetyHookContext& ctx) {
                    ctx.xmm0.f32[0] = g_AspectRatio;
                });
        }
        else UWHook.enable();
        Memory::PatchBytes(Cutscenesaddress, "\x90\xE9", 2);
    }
    else {
        if (UWHook) UWHook.disable();
        Memory::RestoreBytes(Cutscenesaddress);
    }

    logger->info("Ultrawide fix {}", g_fix_enabled && g_ultrawide_fix_enabled ? "enabled" : "disabled");
}

float aspect = (g_AspectRatio > 16.f / 9.f) ? 2.370370388f : 16.f / 9.f;
static float initialHUDOffset = (1920.f - (screenHeight * aspect)) / 2.f;
static std::vector<uint8_t> memoryOffset = {};
static int32_t offset = 0;
static void HUDFixEnabled() {
    if (!HUD1address || !HUD2address) return;
    if (!HUDHook) { // Hook only once
        if (memoryOffset.empty()) {
			memoryOffset = Memory::ReadBytes(HUD1address + 4, 4); // Retrieve the offset used to access the HUD
			offset = *reinterpret_cast<uint32_t*>(memoryOffset.data()); // Convert it to a usable int32 offset
        }
        HUDHook = safetyhook::create_mid(HUD1address,
            [](SafetyHookContext& ctx) {
                *reinterpret_cast<float*>(ctx.rax + offset) = g_fix_enabled && g_HUD_fix_enabled ? initialHUDOffset + g_HUDOffset : initialHUDOffset;
            });
    }
	Memory::PatchBytes(HUD1address + 0x19, "\x90\x90\x90\x90\x90\x90\x90\x90", 8); // Prevent rax+0xD0 from being overridden
	Memory::PatchBytes(HUD2address, "\xEB", 1); // Force the engine to reapply the new HUD position every frame

    logger->info("HUD fix {}", g_fix_enabled && g_HUD_fix_enabled ? "enabled" : "disabled");
}

static void CameraFixEnabled() {
    if (!Cameraaddress) return;
    if (g_fix_enabled && g_camera_fix_enabled) {
        if (!CameraHook) {
            CameraHook = safetyhook::create_mid(Cameraaddress + 0x8,
                [](SafetyHookContext& ctx) {
                    g_Camera_In = ctx.xmm0.f32[0];
                    ctx.xmm0.f32[0] *= g_cameraDistanceMultiplier;
                    g_Camera_Out = ctx.xmm0.f32[0];
                });
        }
        else CameraHook.enable();
    }
    else if (CameraHook) CameraHook.disable();

    logger->info("Camera fix {}", g_fix_enabled && g_camera_fix_enabled ? "enabled" : "disabled");
}

static void VignetteFixEnabled() {
    if (!Vignetteaddress) return;
    if (g_fix_enabled && g_vignetting_fix_enabled) {
        if (!VignetteHook) {
            VignetteHook = safetyhook::create_mid(Vignetteaddress,
                [](SafetyHookContext& ctx) { ctx.xmm3.f32[0] = 0.f; });
        }
        else VignetteHook.enable();
    }
    else if (VignetteHook) VignetteHook.disable();

    logger->info("Vignetting fix {}", g_fix_enabled && g_vignetting_fix_enabled ? "enabled" : "disabled");
}

static void SharpeningFixEnabled() {
    if (!Sharpeningaddress) return;
    if (g_fix_enabled && g_sharpening_fix_enabled) {
        if (!SharpeningHook) {
            SharpeningHook = safetyhook::create_mid(Sharpeningaddress,
                [](SafetyHookContext& ctx) { ctx.xmm6.f32[0] = g_sharpeningValue ; });
        }
        else SharpeningHook.enable();
    }
    else if (SharpeningHook) SharpeningHook.disable();

    logger->info("Sharpening fix {}", g_fix_enabled && g_sharpening_fix_enabled ? "enabled" : "disabled");
}

// Memory patch fixes
static void DOFFixEnabled() {
    if (!DOFaddress) return;
    if (g_fix_enabled && g_DOF_fix_enabled)
        Memory::PatchBytes(DOFaddress, "\xC3", 1);
    else Memory::RestoreBytes(DOFaddress);

    logger->info("Depth of field fix {}", g_fix_enabled && g_DOF_fix_enabled ? "enabled" : "disabled");
}

static void PSNCheckRemoval() {
    if (!PSNCheckaddress) return;
    Memory::PatchBytes(PSNCheckaddress, "\x31\xC0\x90\x90\x90", 5); //
}

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