#pragma once

#include <string>
#include <unordered_map>
#include <cstdint>

#include <DynamicOutput/DynamicOutput.hpp>
#include <Unreal/UObject.hpp>

#include <Api/Widget.hpp>

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

    class WidgetRegistry
    {
    public:
        using Generation = uint32_t;
        using KeyType = std::basic_string<TCHAR>;

        // ---- generation control --------------------------------------

        // Advances the generation and drops every existing entry.

        static auto BeginGeneration() -> Generation
        {
            ++s_generation;
            s_entries.clear();
            Output::send<LogLevel::Verbose>(STR("[WidgetRegistry] generation advanced to {}\n"), s_generation);
            return s_generation;
        }

        // Calls Widget::RemoveFromParent on every entry still in the
        // current generation, THEN advances the generation. 
        // 
        // If an entry's pointer is already stale, this will call RemoveFromParent on whatever
        // object now lives at that address. That is the one real cost of the gap this class cannot close. Keep your generations short.
        static auto TeardownGeneration() -> void
        {
            for (auto& [Key, E] : s_entries)
            {
                if (E.Gen != s_generation || !E.Obj) continue;
                Widget W{ E.Obj };
                if (W) W.RemoveFromParent();
            }
            BeginGeneration();
        }

        static auto CurrentGeneration() -> Generation { return s_generation; }

        // ---- registration -----------------------------------------------

        static auto Register(const KeyType& Key, UObject* Obj) -> bool
        {
            if (!Obj){
                Output::send<LogLevel::Warning>(STR("[WidgetRegistry] Register('{}'): refusing a null object\n"), Key.c_str());
                return false;
            }

            if (s_entries.find(Key) != s_entries.end())
            {
                Output::send<LogLevel::Verbose>(STR("[WidgetRegistry] Register('{}'): overwriting existing entry\n"), Key.c_str());
            }

            s_entries[Key] = Entry{ Obj, s_generation };
            return true;
        }

        static auto Register(const TCHAR* Key, UObject* Obj) -> bool
        {
            return Register(KeyType{ Key }, Obj);
        }

        static auto Unregister(const KeyType& Key) -> bool
        {
            return s_entries.erase(Key) > 0;
        }

        // ---- lookup -----------------------------------------------------

        // Returns nullptr if the key is missing OR if the stored entry predates the current generation. 
        // A stale generation is treated identically to "never registered," and the dead entry is erased
        // on the way out so it doesn't accumulate.
        static auto GetObject(const KeyType& Key) -> UObject*
        {
            auto It = s_entries.find(Key);
            if (It == s_entries.end()) return nullptr;

            if (It->second.Gen != s_generation)
            {
                Output::send<LogLevel::Verbose>(STR("[WidgetRegistry] GetObject('{}'): stale generation ({} vs current {}), dropping\n"), Key.c_str(), It->second.Gen, s_generation);
                s_entries.erase(It);
                return nullptr;
            }

            return It->second.Obj;
        }

        static auto GetObject(const TCHAR* Key) -> UObject*
        {
            return GetObject(KeyType{ Key });
        }

        // Typed lookup. T must be one of the wrapper types in this tree (Widget, TextBlock, CanvasPanel, PanelWidget, ...) exposing
        // ClassOf<T>() -- the same requirement UnrealWrapper's derived types already have. Constructing T{Obj} runs T's own IsA()
        // check, so a cross-class stale pointer comes back as an invalid wrapper (IsValid() == false) instead of a wrong-typed one.
        template <typename T>
        static auto Get(const KeyType& Key) -> T
        {
            UObject* Obj = GetObject(Key);
            return Obj ? T{ Obj } : T{};
        }

        template <typename T>
        static auto Get(const TCHAR* Key) -> T
        {
            return Get<T>(KeyType{ Key });
        }

        static auto Contains(const KeyType& Key) -> bool
        {
            return GetObject(Key) != nullptr;
        }

        static auto Count() -> size_t { return s_entries.size(); }

        // Drops everything without touching the Slate tree. Same caveat as BeginGeneration(): this does not remove widgets from their parents, it just makes the registry stop tracking them.
        static auto Clear() -> void { s_entries.clear(); }

    private:
        struct Entry { UObject* Obj{}; Generation Gen{}; };

        inline static std::unordered_map<KeyType, Entry> s_entries{};
        inline static Generation                          s_generation{ 1 };
    };
}