#include "CommonHeaders.h"
#include "CommonUEHeaders.h"
#include "Phoenix_classes.hpp"
#include "SDK/UI_BP_PhoenixHUDWidget_classes.hpp"
#include "SDK/UI_BP_AsyncLoadWrapper_classes.hpp"
#include "SDK/UI_BP_BeaconHUDPanel_classes.hpp"

using namespace SDK;

// Constants
const std::string PLUGIN_NAME = "HogwartsLegacy";
const std::string PLUGIN_LOG = PLUGIN_NAME + ".log";
constexpr ULONGLONG DEFAULT_ACTORS_SCAN_BETWEEN_TICKS = 40; // Used for enemies time dilation

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

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

// 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_CA_fix_enabled = false;
static std::atomic <bool> g_Vignetting_fix_enabled = false;
static std::atomic <bool> g_Fog_fix_enabled = false;
static std::atomic <bool> g_TimeDilation_fix_enabled = false;
static std::atomic <bool> g_GodMide_fix_enabled = false;
static std::atomic <bool> g_Stealth_fix_enabled = false;
static int g_AdditionalFOVValue = 0;
static float g_CameraMultiplier = 1.f;
static float g_WorldTimeDilationValue = 1.f;
static float g_AITimeDilationValue = 1.f;
static float g_HUDOffsets = 0.f;

// Shared values
static float g_FOV_In = 67.85f;
static float g_FOV_Out = 67.85f;
static float g_Camera_In = 360.f;
static float g_Camera_Out = 360.f;

// AOB Scan pointers
static uint8_t* FOVaddress = nullptr;
static uint8_t* ConstrainAspectaddress = nullptr;
static uint8_t* Cameraaddress = nullptr;
static uint8_t* WorldTimedilationaddress = nullptr;
static uint8_t* timedilationaddress = nullptr;

// Hooking
static SafetyHookMid FOVHook{};
static SafetyHookMid CameraHook{};
static SafetyHookMid PEHook{};
static SafetyHookMid WorldTimeDilationHook{};
static SafetyHookMid TimeDilationHook{};

// Prototypes
static void FOVFixEnabled();
static void UltraWideFixEnabled();
static void CameraFixEnabled();
static void EnableConsole();
static void HUDFixEnabled();
static void ProcessEvent();
static void EnableCheats(Cheat cheat);
static void LogHUD();

extern "C" __declspec(dllexport) void SetFixEnabled(bool enabled, bool init) {
    g_fix_enabled = enabled;
    if (init) {
        logger = InitializeLogger("Hogwarts Legacy", PLUGIN_LOG);
        logger->info("Plugin {} loaded.", PLUGIN_NAME);
    }
    if (!AOBScanDone) { // Unreal Engine 4.27.2
        logger->info("--------------- AOB scan started ---------------");
        uint8_t* baseModule = reinterpret_cast<uint8_t*>(GetModuleHandleA(nullptr)); // Get game base address
        constexpr auto FOVStringObfuscated = make_obfuscated<0x4A>("EB ?? F3 0F ?? ?? ?? ?? ?? ?? F3 0F ?? ?? ?? 8B 83");
        constexpr auto CameraStringObfuscated = make_obfuscated<0x4A>("F3 0F ?? ?? C8 ?? ?? ?? E8 ?? ?? ?? ?? 4C 8D ?? ?? ?? ?? ?? 48");
        constexpr auto WorldTimeDilationStringObfuscated = make_obfuscated<0x59>("F3 0F 10 81 ?? ?? ?? ?? F3 0F 59 81 ?? ?? ?? ?? F3 0F 59 81 ?? ?? ?? ?? C3");
        constexpr auto timeDilationObfuscated = make_obfuscated<0xB4>("F3 0F ?? ?? ?? ?? ?? ?? EB ?? F3 0F ?? ?? ?? ?? ?? ?? 48 8B ?? ?? 4C 8B ?? F3 0F ?? ?? 44 8B");
        using AOBScan::Make;
        using OffsetScan::Make;
        // Prepare all data for scanning
        std::vector<AOBScanEntry> signatures = {
            Make(&FOVaddress, FOVStringObfuscated, "FOV"),
            Make(&Cameraaddress, CameraStringObfuscated, "Camera distance"),
            Make(&WorldTimedilationaddress, WorldTimeDilationStringObfuscated, "World time dilation"),
            Make(&timedilationaddress, timeDilationObfuscated, "Time dilation"),
        };
        // Scan all signature in a batch
        Memory::AOBScanBatch(signatures, logger);
        if (FOVaddress) ConstrainAspectaddress = FOVaddress + 0x18;
        if (FOVaddress && ConstrainAspectaddress && Cameraaddress && WorldTimedilationaddress && timedilationaddress)
            logger->info("All AOB signatures found. Ready to patch...");

        if (!GObjectsaddress || !AppendStringaddress || !ProcessEventaddress) {
            logger->info("------------ UEngine offsets search ------------");

            constexpr auto GObjetcsStringObfuscated = make_obfuscated<0x4A>("48 8B ?? ?? ?? ?? ?? 48 8B ?? ?? 48 8D ?? ?? EB ?? 33");
            constexpr auto GWorldStringObfuscated = make_obfuscated<0x4A>("48 8B 05 ?? ?? ?? ?? 48 ?? ?? 0F 84 ?? ?? ?? ?? 45 ?? ?? 48 ?? ?? 74");
            constexpr auto AppendStringStringObfuscated = make_obfuscated<0x4A>("48 89 ?? ?? ?? 48 89 ?? ?? ?? 57 48 ?? ?? ?? 8B ?? 48 ?? ?? 8B ?? 44 ?? ?? ?? C1 ?? ?? 48 ?? ?? 80 3D ?? ?? ?? ?? ?? 89 ?? ?? ?? 44 89 ?? ?? ?? 74 ?? 4C 8D ?? ?? ?? ?? ?? EB ?? 48 8D ?? ?? ?? ?? ?? E8 ?? ?? ?? ?? 4C ?? ??  C6 05 ?? ?? ?? ?? ?? 48 8B ?? ?? ?? 48 ?? ?? 48 ?? ?? ?? 8D ?? ?? 49 03 ?? ?? ?? E8 ?? ?? ?? ?? 83");
            constexpr auto ProcessEventStringObfuscated = make_obfuscated<0x4A>("40 ?? 56 57 41 ?? 41 ?? 41 ?? 41 ?? 48 81 ?? ?? ?? ?? ?? 48 8D ?? ?? ?? 48 89 ?? ?? ?? ?? ?? 48 8B ?? ?? ?? ?? ?? 48 ?? ?? 48 89 ?? ?? ?? ?? ?? 8B");

            // Prepare all data for scanning
            std::vector<OffsetScanEntry> UEoffsetsScans = {
                Make(&GObjectsaddress, GObjetcsStringObfuscated, "GObjects", OffsetCalcType::GetOffsetFromOpcode, &Offsets::GObjects, 0x3),
                Make(&GWorldaddress, GWorldStringObfuscated, "GWorld", OffsetCalcType::GetOffsetFromOpcode, &Offsets::GWorld, 0x3),
                Make(&AppendStringaddress, AppendStringStringObfuscated, "AppendString", OffsetCalcType::UE_CalculateOffset, &Offsets::AppendString),
                Make(&ProcessEventaddress, ProcessEventStringObfuscated, "ProcessEvent", OffsetCalcType::UE_CalculateOffset, &Offsets::ProcessEvent)
            };
            // Retrieve all Unreal Engine offsets in a batch
            Memory::OffsetScanBatch(UEoffsetsScans, baseModule, logger, "");

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

    ProcessEvent();
    if (init) return;
    FOVFixEnabled();
    UltraWideFixEnabled();
    CameraFixEnabled();
    HUDFixEnabled();
    LogHUD();
	SetAllEffectsToBeToggled();
    LogFixToggle(GameFixes::None, g_fix_enabled);
}

// Setters for Reshade addon call
extern "C" __declspec(dllexport) void SetFixesEnabled(GameFixes fix, bool enabled) {
    bool bVisual = g_fix_enabled && enabled;
    if (fix == GameFixes::FOV)                  { g_fov_fix_enabled = enabled; FOVFixEnabled(); }
    if (fix == GameFixes::UltraWide)            { g_ultrawide_fix_enabled = enabled; UltraWideFixEnabled(); }
    if (fix == GameFixes::Camera)               { g_camera_fix_enabled = enabled; CameraFixEnabled(); }
    if (fix == GameFixes::DOF)                  { g_DOF_fix_enabled = enabled; gPendingDOF = true; LogFixToggle(fix, bVisual); }
    if (fix == GameFixes::ChromaticAberrations) { g_CA_fix_enabled = enabled; gPendingCA = true; LogFixToggle(fix, bVisual); }
    if (fix == GameFixes::Vignetting)           { g_Vignetting_fix_enabled = enabled; gPendingVignetting = true; LogFixToggle(fix, bVisual); }
    if (fix == GameFixes::Fog)                  { g_Fog_fix_enabled = enabled; gPendingFog = true; LogFixToggle(fix, bVisual); }
    if (fix == GameFixes::HUD)                  { g_HUD_fix_enabled = enabled; HUDFixEnabled(); LogHUD(); }
    if (fix == GameFixes::TimeDilation)         { g_TimeDilation_fix_enabled = enabled; EnableCheats(Cheat::TimeDilation); }
    if (fix == GameFixes::GodMode)              { g_GodMide_fix_enabled = enabled; EnableCheats(Cheat::GodMode); }
    if (fix == GameFixes::Stealth)              { g_Stealth_fix_enabled = enabled; EnableCheats(Cheat::Stealth); }
}

extern "C" __declspec(dllexport) void SetValues(GameSetting setting, float value) {
    if (setting == GameSetting::FOV) g_AdditionalFOVValue = (int)(value);
    if (setting == GameSetting::CameraDistance) g_CameraMultiplier = value;
    if (setting == GameSetting::HUD) {
        g_HUDOffsets = (value * screenWidth) / 100;
        HUDFixEnabled();
    }
    if (setting == GameSetting::WorldTimeDilation) g_WorldTimeDilationValue = value;
    if (setting == GameSetting::AITimeDilation) g_AITimeDilationValue = value;
}

// 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->consoleEnabled = g_Console_Enabled;
}

// Hooks injections
static UUI_BP_PhoenixHUDWidget_C* gPhoenixHUDWidget = nullptr;
static ULONGLONG lastScanTick = 0; // Last time Actors were scanned for enemies time dilation
static void ProcessEvent() {
    if (!PEHook) {
        PEHook = safetyhook::create_mid(ProcessEventaddress + 0xc,
            [](SafetyHookContext& ctx) {
                UObject* object = (UObject*)ctx.rcx;
                UFunction* function = (UFunction*)ctx.rdx;
                if (!object || !function) return;

                if (object->IsA(UUserWidget::StaticClass())) {
                    std::string funcName = function->GetName();
					std::string objectName = object->Class->GetName();
                    if (objectName == "UI_BP_Title_C" && funcName == "Construct")
						EnableConsole(); // Enable console when title screen is loaded to ensure it will work

                    if (object->IsA(UUI_BP_PhoenixHUDWidget_C::StaticClass())) {
                        if (funcName == "OnShowEnemyInfo") {
                            auto* phoenixHUDWidget = static_cast<UUI_BP_PhoenixHUDWidget_C*>(object);
                            if (phoenixHUDWidget && phoenixHUDWidget->Class) {
                                gPhoenixHUDWidget = phoenixHUDWidget; // Store UUI_BP_PhoenixHUDWidget_C pointer
                                HUDFixEnabled(); // Enable live offsets modifications
                            }
                        }
                        else if (funcName == "Destruct") gPhoenixHUDWidget = nullptr;
                    }
                }
            });
    }
}
// HUD fix
static void HUDFixEnabled() {
    if (!gPhoenixHUDWidget) return;

    auto ApplyCanvas{ [](UWidget* panel, float left, float right = 0.f) {
        if (!panel || !panel->Slot || !panel->Slot->IsA(UCanvasPanelSlot::StaticClass())) return;
        auto* slot = static_cast<UCanvasPanelSlot*>(panel->Slot);

        FMargin offsets = slot->GetOffsets();
        offsets.Left = left;
        offsets.Right = right;
        slot->SetOffsets(offsets);
    }};

    float rootOffset = g_HUD_fix_enabled ? g_HUDOffsets : 0.f;
    ApplyCanvas(gPhoenixHUDWidget->Root, rootOffset, rootOffset); // Position main HUD
    ApplyCanvas(gPhoenixHUDWidget->AsyncReticule, -rootOffset); // Fix Reticle shift
    ApplyCanvas(gPhoenixHUDWidget->BeaconHUDPanel, -rootOffset); // Fix marker shift
}

static void LogHUD() { // Log toggle state
    if (g_HUD_fix_enabled) logger->info("HUD fix enabled");
    else logger->info("HUD fix disabled");
}

static void FOVFixEnabled() {
	if (!FOVaddress) return;
    if (g_fix_enabled && g_fov_fix_enabled) {
        if (!FOVHook) { // Hook only once
            FOVHook = safetyhook::create_mid(FOVaddress + 0xa,
                [](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 CameraFixEnabled() {
	if (!Cameraaddress) return;
    if (g_fix_enabled && g_camera_fix_enabled) {
        if (!CameraHook) { // Hook only once
            CameraHook = safetyhook::create_mid(Cameraaddress + 0x8,
                [](SafetyHookContext& ctx) {
                    g_Camera_In = ctx.xmm1.f32[0];
                    ctx.xmm1.f32[0] *= g_CameraMultiplier;
                    g_Camera_Out = ctx.xmm1.f32[0];
                });
        }
        else CameraHook.enable();
    }
    else if (CameraHook) CameraHook.disable();

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

// Cheats
static void EnableCheats(Cheat cheat) {
    if (WorldTimedilationaddress && !WorldTimeDilationHook) {
        WorldTimeDilationHook = safetyhook::create_mid(WorldTimedilationaddress + 0x10,
            [](SafetyHookContext& ctx) {
                // From AWorldSettings retrieved from world->K2_GetWorldSettings()
                ctx.xmm0.f32[0] *= g_TimeDilation_fix_enabled ? g_WorldTimeDilationValue : 1.f;
                // Apply visual effects only in hook main thread to ensure they will be applied correctly
                if (!g_Console_Enabled) return; // It relies on dev console being reactivated
                if (gPendingFog.exchange(false))
                    ApplyVisualEffect(GameFixes::Fog, g_fix_enabled && g_Fog_fix_enabled);
                if (gPendingDOF.exchange(false))
                    ApplyVisualEffect(GameFixes::DOF, g_fix_enabled && g_DOF_fix_enabled);
                if (gPendingCA.exchange(false))
                    ApplyVisualEffect(GameFixes::ChromaticAberrations, g_fix_enabled && g_CA_fix_enabled);
                if (gPendingVignetting.exchange(false))
                    ApplyVisualEffect(GameFixes::Vignetting, g_fix_enabled && g_Vignetting_fix_enabled);
            });
    }
    if (timedilationaddress && !TimeDilationHook) {
        TimeDilationHook = safetyhook::create_mid(timedilationaddress,
            [](SafetyHookContext& ctx) {
                if (!ctx.rbx) return;
                UObject* object = (UObject*)ctx.rbx;
                if (!object || !object->Class) return;

                ULONGLONG now = GetTickCount64();
                if (now - lastScanTick < DEFAULT_ACTORS_SCAN_BETWEEN_TICKS) return;
                lastScanTick = now;
                
                if (object->IsA(ABiped_Player::StaticClass())) {
                    ABiped_Player* bipedPlayer = static_cast<ABiped_Player*>(object);
                    if (bipedPlayer && bipedPlayer->Controller && bipedPlayer->Controller->IsA(APlayerController::StaticClass()))
                        if (g_GodMide_fix_enabled) bipedPlayer->RestoreHealth();
                }
                if (!object->IsA(AEnemy_Character::StaticClass())) return;

                AEnemy_Character* enemy = static_cast<AEnemy_Character*>(object);
                if (!enemy || !enemy->Class) return;
                enemy->CustomTimeDilation = g_TimeDilation_fix_enabled ? g_AITimeDilationValue : 1.f; // Enemy time dilation
                if (enemy->EnemyAIComponent && enemy->EnemyAIComponent->IsA(UEnemyAIComponent::StaticClass()))
                    enemy->EnemyAIComponent->EnablePerception(!g_Stealth_fix_enabled); // Stealth mode
            });
    }

    if (cheat == Cheat::TimeDilation) logger->info("Time dilation cheat {}", g_TimeDilation_fix_enabled ? "enabled" : "disabled");
    if (cheat == Cheat::GodMode) logger->info("God mode cheat {}", g_GodMide_fix_enabled ? "enabled" : "disabled");
    if (cheat == Cheat::Stealth) logger->info("Stealth cheat {}", g_Stealth_fix_enabled ? "enabled" : "disabled");
}

// Memory patch fixes
static void UltraWideFixEnabled() {
	if (!ConstrainAspectaddress) return;
    if (g_fix_enabled && g_ultrawide_fix_enabled)
        Memory::PatchBytes(ConstrainAspectaddress, "\x31\xC9\x90\x90\x90\x90\x90", 7); // bConstrainAspectRatio = 0
    else Memory::RestoreBytes(ConstrainAspectaddress);

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

// UE Console creation
static void EnableConsole() {
    if (g_Console_Enabled || !GObjectsaddress || !AppendStringaddress || !ProcessEventaddress)
        return;

    logger->info("-------------- Console re-enabling --------------");
    ReactivateDevConsole(logger);
}

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