#pragma once

#include <string>
#include <vector>

#include <Unreal/NameTypes.hpp>

#include "UnrealWrapper.hpp"
#include "Widget.hpp"
#include "PanelWidget.hpp"

namespace PalUIExtension::Api
{

    class WidgetTree : public UnrealWrapper
    {
    public:
        static auto ClassPath() -> const TCHAR* { return STR("/Script/UMG.WidgetTree"); }
        static auto DebugTag()  -> const TCHAR* { return STR("WidgetTree"); }

        WidgetTree() = default;
        explicit WidgetTree(UObject* Obj) : UnrealWrapper(Obj, ClassOf<WidgetTree>(), DebugTag()) {}
      
        static auto FromUserWidget(UObject* UserWidget) -> WidgetTree
        {
            if (!UserWidget) return WidgetTree{};

            UObject* Tree = RawObject{ UserWidget, STR("UserWidget") }.TryObjectProperty(STR("WidgetTree"));
            return Tree ? WidgetTree{ Tree } : WidgetTree{};
        }

        // The tree a widget LIVES IN.
        static auto ContainingTree(UObject* AnyWidget) -> WidgetTree
        {
            if (!AnyWidget) return WidgetTree{};

            UClass* TreeClass = ClassOf<WidgetTree>();
            if (!TreeClass) return WidgetTree{};   // ClassOf already logged

            if (UObject* Tree = AnyWidget->GetTypedOuter(TreeClass))
            {
                return WidgetTree{ Tree };
            }
           
            UObject* Outer = AnyWidget->GetOuterPrivate();
            for (int32 Hops = 0; Outer && Hops < kMaxOuterHops; ++Hops)
            {
                if (Outer->IsA(TreeClass)) return WidgetTree{ Outer };
                Outer = Outer->GetOuterPrivate();
            }
            return WidgetTree{};
        }

        static auto FromWidget(UObject* AnyWidget) -> WidgetTree
        {
            if (WidgetTree Owned = FromUserWidget(AnyWidget)) return Owned;
            return ContainingTree(AnyWidget);
        }

        // ---- root ---------------------------------------------------------

        auto GetRootWidgetObject() const -> UObject*
        {
            auto Ptr = PropertyPtr<UObject*>(STR("RootWidget"));
            return Ptr ? *Ptr : nullptr;
        }

        auto Root() const -> Widget
        {
            UObject* R = GetRootWidgetObject();
            return R ? Widget{ R } : Widget{};
        }

        // ==================================================================
        // TRAVERSAL
        // ==================================================================

        enum class EScope : uint8
        {
            ThisTreeOnly,
            IncludeNestedUserWidgets,
        };

        // What a visitor wants to happen next.
        enum class EVisit : uint8
        {
            Continue,      // descend into this node's children
            SkipChildren,  // this node is fine, its subtree is not interesting
            Stop,          // end the walk entirely
        };

        static constexpr int32 kDefaultNodeBudget = 8192;
        static constexpr int32 kMaxDepth = 64;
        static constexpr int32 kMaxOuterHops = 32;

        template <typename Fn>
        auto ForEach(Fn&& Visit, EScope Scope = EScope::ThisTreeOnly, int32 NodeBudget = kDefaultNodeBudget) -> bool
        {
            WalkState State{ Scope, NodeBudget, false, false };
            WalkNode(GetRootWidgetObject(), 0, State, Visit);

            if (State.bExhausted)
            {
                Output::send<LogLevel::Warning>(
                    STR("[WidgetTree] walk aborted -- exceeded {} nodes or depth {}. ")
                    STR("Results are incomplete.\n"), NodeBudget, kMaxDepth);
            }
            return !State.bExhausted;
        }

        // ==================================================================
        // SEARCH
        // ==================================================================

        auto FindObjectByName(FName Name, EScope Scope = EScope::ThisTreeOnly) -> UObject*
        {
            if (Name.IsNone()) return nullptr;

            UObject* Found = nullptr;
            ForEach([&](UObject* Node, int32) -> EVisit {
                    if (Node->GetNamePrivate() == Name)
                    {
                        Found = Node;
                        return EVisit::Stop;
                    }
                    return EVisit::Continue;
                }, Scope);

            return Found;
        }

        auto FindObjectByName(const TCHAR* Name, EScope Scope = EScope::ThisTreeOnly) -> UObject*
        {
            if (!Name) return nullptr;
            return FindObjectByName(FName(Name, FNAME_Find), Scope);
        }

        // Same search, wrapped. Invalid Widget on a miss.
        auto Find(const TCHAR* Name, EScope Scope = EScope::ThisTreeOnly) -> Widget
        {
            UObject* O = FindObjectByName(Name, Scope);
            return O ? Widget{ O } : Widget{};
        }

        template <typename T>
        auto FindAs(const TCHAR* Name, EScope Scope = EScope::ThisTreeOnly) -> T
        {
            UObject* O = FindObjectByName(Name, Scope);
            if (!O)
            {
                Output::send<LogLevel::Warning>(
                    STR("[WidgetTree] no widget named '{}' in this tree\n"), Name);
                return T{};
            }

            if (!ObjectIsA<T>(O))
            {
                Output::send<LogLevel::Warning>(
                    STR("[WidgetTree] '{}' resolved to {}, which is not a {}\n"),
                    Name, O->GetFullName(), T::DebugTag());
                return T{};
            }
            return T{ O };
        }
        
        template <typename T>
        auto FindFirstOfType(EScope Scope = EScope::ThisTreeOnly) -> T
        {
            UClass* Cls = ClassOf<T>();
            if (!Cls) return T{};

            UObject* Found = nullptr;
            ForEach([&](UObject* Node, int32) -> EVisit
                {
                    if (Node->IsA(Cls)) { Found = Node; return EVisit::Stop; }
                    return EVisit::Continue;
                }, Scope);

            return Found ? T{ Found } : T{};
        }

        // Every widget of a given wrapper type, in traversal order.
        template <typename T>
        auto FindAllOfType(EScope Scope = EScope::ThisTreeOnly) -> std::vector<T>
        {
            std::vector<T> Out;

            UClass* Cls = ClassOf<T>();
            if (!Cls) return Out;

            ForEach([&](UObject* Node, int32) -> EVisit
                {
                    if (Node->IsA(Cls)) Out.emplace_back(Node);
                    return EVisit::Continue;
                }, Scope);

            return Out;
        }

        // Flat list of every widget object the walk reaches.
        auto AllWidgetObjects(EScope Scope = EScope::ThisTreeOnly) -> std::vector<UObject*>
        {
            std::vector<UObject*> Out;
            ForEach([&](UObject* Node, int32) -> EVisit
                {
                    Out.push_back(Node);
                    return EVisit::Continue;
                }, Scope);
            return Out;
        }

        // ==================================================================
        // DIAGNOSTICS
        // ==================================================================

        auto Dump(EScope Scope = EScope::IncludeNestedUserWidgets) -> void
        {
            if (!m_obj)
            {
                Output::send<LogLevel::Warning>(STR("[WidgetTree] Dump on an invalid tree\n"));
                return;
            }

            Output::send<LogLevel::Warning>(
                STR("[WidgetTree] dump of {}\n"), m_obj->GetFullName());

            int32 Count = 0;
            ForEach([&](UObject* Node, int32 Depth) -> EVisit
                {
                    ++Count;

                    std::basic_string<TCHAR> Indent;
                    Indent.reserve(static_cast<size_t>(Depth) * 2);
                    for (int32 i = 0; i < Depth; ++i) Indent += STR("  ");

                    UClass* Cls = Node->GetClassPrivate();
                    std::basic_string<TCHAR> ClassName =
                        Cls ? Cls->GetName() : std::basic_string<TCHAR>(STR("<no class>"));

                    if (ObjectIsA<PanelWidget>(Node))
                    {
                        Output::send<LogLevel::Warning>(
                            STR("  {}{} [{}] ({} children)\n"),
                            Indent.c_str(), Node->GetName().c_str(), ClassName.c_str(),
                            static_cast<int32>(PanelWidget{ Node }.GetChildObjects().size()));
                    }
                    else
                    {
                        Output::send<LogLevel::Warning>(
                            STR("  {}{} [{}]\n"),
                            Indent.c_str(), Node->GetName().c_str(), ClassName.c_str());
                    }
                    return EVisit::Continue;
                }, Scope);

            Output::send<LogLevel::Warning>(STR("[WidgetTree] {} widgets\n"), Count);
        }

        // ==================================================================
        // CONSTRUCTION
        // ==================================================================

        auto ConstructWidget(UClass* WidgetClass, const TCHAR* InstanceName) -> UObject*
        {
            if (!m_obj)
            {
                Output::send<LogLevel::Warning>(
                    STR("[WidgetTree] ConstructWidget on an invalid tree\n"));
                return nullptr;
            }
            if (!WidgetClass)
            {
                Output::send<LogLevel::Warning>(
                    STR("[WidgetTree] ConstructWidget: null class\n"));
                return nullptr;
            }

            FStaticConstructObjectParameters Params{ WidgetClass };
            Params.Outer = m_obj;
            if (InstanceName) Params.Name = FName(InstanceName);

            UObject* Obj = UObjectGlobals::StaticConstructObject(Params);
            if (!Obj)
            {
                Output::send<LogLevel::Warning>(
                    STR("[WidgetTree] ConstructWidget: StaticConstructObject returned null for '{}'\n"),
                    InstanceName ? InstanceName : STR("<unnamed>"));
            }
            return Obj;
        }

        template <typename T>
        auto ConstructWidget(const TCHAR* InstanceName) -> T
        {
            UObject* O = ConstructWidget(ClassOf<T>(), InstanceName);
            return O ? T{ O } : T{};
        }

    private:
        struct WalkState
        {
            EScope Scope;
            int32  Budget;
            bool   bStopped;
            bool   bExhausted;
        };
      
        static auto OwnedTreeObject(UObject* Node) -> UObject*
        {
            return RawObject{ Node, STR("Widget") }.TryObjectProperty(STR("WidgetTree"));
        }

        template <typename Fn>
        static auto WalkNode(UObject* Node, int32 Depth, WalkState& State, Fn& Visit) -> void
        {
            if (!Node || State.bStopped) return;

            if (State.Budget-- <= 0 || Depth >= kMaxDepth)
            {
                State.bExhausted = true;
                State.bStopped = true;
                return;
            }

            const EVisit Result = Visit(Node, Depth);
            if (Result == EVisit::Stop) { State.bStopped = true; return; }
            if (Result == EVisit::SkipChildren) return;
         
            if (UObject* NestedTree = OwnedTreeObject(Node))
            {
                if (State.Scope == EScope::IncludeNestedUserWidgets)
                {
                    UObject* NestedRoot =
                        RawObject{ NestedTree, STR("WidgetTree") }.TryObjectProperty(STR("RootWidget"));
                    WalkNode(NestedRoot, Depth + 1, State, Visit);
                }
                return;
            }

            if (!ObjectIsA<PanelWidget>(Node)) return;
          
            for (UObject* Child : PanelWidget{ Node }.GetChildObjects())
            {
                if (State.bStopped) return;
                WalkNode(Child, Depth + 1, State, Visit);
            }
        }
    };
}