#pragma once

#include <vector>
#include <string>
#include <cstdint>

#include <DynamicOutput/DynamicOutput.hpp>

#include <Unreal/Hooks.hpp>
#include <Unreal/UObjectGlobals.hpp>
#include <Unreal/UObject.hpp>
#include <Unreal/UClass.hpp>
#include <Unreal/UFunction.hpp>

namespace PalUIExtension::Static
{
    using namespace RC;
    using namespace RC::Unreal;

    // Registry + deferred installer.
    //
    // Register() resolves the target UFunction immediately if it is already loaded. If it isn't, the request is parked in s_pending and a single
    // ProcessEvent pre-callback watches for a call on a matching <class name, function name> pair, then installs the real hook.
    //
    // NOT THREAD SAFE. 

    class EventHookRegistry
    {
    public:
        using StringType = std::basic_string<TCHAR>;
        using HookId = uint32_t;

        static constexpr HookId InvalidHook = 0;

    private:
        struct HookRequest
        {
            HookId       Id{};
            StringType   Name;           // owned, not a borrowed const TCHAR*
            StringType   FunctionPath;   // owned
            FName        ClassName{};    // e.g. WBP_PalStatus_C
            FName        FunctionName{}; // e.g. Construct

            UnrealScriptFunctionCallable Pre{};
            UnrealScriptFunctionCallable Post{};
            void* CustomData{};
        };

        struct InstalledHook
        {
            HookId     Id{};
            StringType Name;
            UFunction* Function{};
            int        PreId{};
            int        PostId{};
        };

        inline static std::vector<HookRequest>   s_pending{};
        inline static std::vector<InstalledHook> s_installed{};
        inline static HookId                     s_next_id{ 1 };
        inline static bool                       s_watcher_installed{ false };
        inline static bool                       s_shutdown{ false };

    public:
        static void Initialize(size_t expected_hooks = 16)
        {
            s_shutdown = false;
            s_pending.reserve(expected_hooks);
            s_installed.reserve(expected_hooks);
        }

        static HookId Register(
            const TCHAR* HookName,
            const TCHAR* FunctionPath,
            UnrealScriptFunctionCallable Pre = [](UnrealScriptFunctionCallableContext&, void*) {},
            UnrealScriptFunctionCallable Post = [](UnrealScriptFunctionCallableContext&, void*) {},
            void* CustomData = nullptr)
        {
            if (!HookName || !FunctionPath)
            {
                Output::send<LogLevel::Error>(STR("[EventHookRegistry] null HookName or FunctionPath\n"));
                return InvalidHook;
            }

            HookRequest Req{};
            Req.Id = s_next_id++;
            Req.Name = HookName;
            Req.FunctionPath = FunctionPath;
            Req.Pre = Pre;
            Req.Post = Post;
            Req.CustomData = CustomData;

            StringType ClassPart, FuncPart;
            if (!SplitFunctionPath(Req.FunctionPath, ClassPart, FuncPart))
            {
                Output::send<LogLevel::Error>(
                    STR("[EventHookRegistry] '{}': cannot parse '{}' -- expected ")
                    STR("/Path/To/Package.ClassName:FunctionName\n"), HookName, FunctionPath);
                return InvalidHook;
            }
            Req.ClassName = FName(ClassPart.c_str(), FNAME_Add);
            Req.FunctionName = FName(FuncPart.c_str(), FNAME_Add);

            // Fast path: already loaded, install now, never touch ProcessEvent.
            if (UFunction* Fn = UObjectGlobals::StaticFindObject<UFunction*>(nullptr, nullptr, Req.FunctionPath.c_str()))
            {
                Install(Req, Fn);
                return Req.Id;
            }

            Output::send<LogLevel::Verbose>(
                STR("[EventHookRegistry] '{}': '{}' not loaded yet, deferring\n"), HookName, FunctionPath);

            s_pending.push_back(std::move(Req));
            EnsureWatcher();
            return s_next_id - 1;
        }

        static bool Deregister(const TCHAR* HookName)
        {
            if (!HookName) return false;
            const StringType Name{ HookName };

            for (size_t i = 0; i < s_installed.size(); ++i)
            {
                if (s_installed[i].Name != Name) continue;
                Uninstall(s_installed[i]);
                s_installed.erase(s_installed.begin() + i);
                return true;
            }

            for (size_t i = 0; i < s_pending.size(); ++i)
            {
                if (s_pending[i].Name != Name) continue;
                s_pending.erase(s_pending.begin() + i);
                return true;
            }

            Output::send<LogLevel::Warning>(
                STR("[EventHookRegistry] Deregister: no hook named '{}'\n"), HookName);
            return false;
        }

        static bool Deregister(HookId Id)
        {
            for (size_t i = 0; i < s_installed.size(); ++i)
            {
                if (s_installed[i].Id != Id) continue;
                Uninstall(s_installed[i]);
                s_installed.erase(s_installed.begin() + i);
                return true;
            }
            for (size_t i = 0; i < s_pending.size(); ++i)
            {
                if (s_pending[i].Id != Id) continue;
                s_pending.erase(s_pending.begin() + i);
                return true;
            }
            return false;
        }

        static bool IsInstalled(HookId Id)
        {
            for (const auto& H : s_installed) if (H.Id == Id) return true;
            return false;
        }

        static size_t PendingCount() { return s_pending.size(); }

        // The ProcessEvent callback cannot be unregistered in UE4SS, insteasd it drops all state and sets s_shutdown; the watcher
        // then costs one bool test per ProcessEvent for the rest of the process.
        static void Cleanup()
        {
            for (auto& H : s_installed) Uninstall(H);
            s_installed.clear();
            s_pending.clear();
            s_shutdown = true;
        }

    private:
        static void EnsureWatcher()
        {
            if (s_watcher_installed) return;
            Hook::RegisterProcessEventPreCallback(&EventHookRegistry::OnProcessEvent);
            s_watcher_installed = true;
            Output::send<LogLevel::Verbose>(STR("[EventHookRegistry] ProcessEvent watcher installed\n"));
        }

        // HOT PATH. This runs for effectively every UFunction call in the game.
        // Every line before the loop is there to make the common case a load,
        // a branch, and a return.
        static void OnProcessEvent(UObject* Context, UFunction* Function, void*)
        {
            if (s_shutdown || s_pending.empty()) return;
            if (!Context || !Function) return;

            UClass* Cls = Context->GetClassPrivate();
            if (!Cls) return;
            const FName ClsName = Cls->GetNamePrivate();

            // Backwards so erase() doesn't disturb indices we haven't visited.
            for (size_t i = s_pending.size(); i-- > 0;)
            {
                HookRequest& Req = s_pending[i];
                if (Req.ClassName != ClsName) continue;

                UFunction* Target = UObjectGlobals::StaticFindObject<UFunction*>(nullptr, nullptr, Req.FunctionPath.c_str());

                if (!Target)
                {
                    if (Req.FunctionName != Function->GetNamePrivate()) continue;

                    Output::send<LogLevel::Warning>(
                        STR("[EventHookRegistry] '{}': '{}' still will not resolve by path; ")
                        STR("falling back to the dispatched UFunction {}. If your hook fires on the ")
                        STR("wrong object, your path string is wrong.\n"),
                        Req.Name, Req.FunctionPath, Function->GetFullName());
                    Target = Function;
                }

                HookRequest Owned = std::move(Req);
                s_pending.erase(s_pending.begin() + i);
                Install(Owned, Target);
            }
        }

        static void Install(HookRequest& Req, UFunction* Fn)
        {
            if (!Fn) return;

            for (const auto& H : s_installed)
            {
                if (H.Function != Fn) continue;
                Output::send<LogLevel::Warning>(
                    STR("[EventHookRegistry] '{}': {} is already hooked by '{}'. ")
                    STR("Installing anyway -- both callbacks will fire.\n"),
                    Req.Name, Fn->GetFullName(), H.Name);
                break;
            }

            auto _pre = Req.Pre;
            if (!_pre)  _pre = [](UnrealScriptFunctionCallableContext&, void*) {};
            auto _post = Req.Post;
            if (!_post) _post = [](UnrealScriptFunctionCallableContext&, void*) {};

            auto [PreId, PostId] = UObjectGlobals::RegisterHook(Fn, _pre, _post, Req.CustomData);

            s_installed.push_back(InstalledHook{ Req.Id, Req.Name, Fn, PreId, PostId });

            Output::send<LogLevel::Verbose>(
                STR("[EventHookRegistry] '{}': installed on {} (pre={}, post={})\n"),
                Req.Name, Fn->GetFullName(), PreId, PostId);
        }

        static void Uninstall(InstalledHook& H)
        {
            if (!H.Function) return;
            UObjectGlobals::UnregisterHook(H.Function, std::make_tuple(H.PreId, H.PostId));
            H.Function = nullptr;
        }

        // "/Game/Pal/Blueprint/UI/PalStatus/WBP_PalStatus.WBP_PalStatus_C:Construct"
        //   -> ClassName "WBP_PalStatus_C", FunctionName "Construct"
        static bool SplitFunctionPath(const StringType& Path, StringType& OutClass, StringType& OutFunction)
        {
            const auto Colon = Path.rfind(STR(':'));
            if (Colon == StringType::npos || Colon + 1 >= Path.size()) return false;

            OutFunction = Path.substr(Colon + 1);

            const StringType ObjectPart = Path.substr(0, Colon);
            const auto Dot = ObjectPart.rfind(STR('.'));
            OutClass = (Dot == StringType::npos) ? ObjectPart : ObjectPart.substr(Dot + 1);

            return !OutClass.empty() && !OutFunction.empty();
        }
    };
}