#include "DismembermentGraph/DismembermentGraphEditor.h"
#include "DismembermentGraph/DismembermentGraphAsset.h"
#include "DismembermentGraph/DismembermentGraph.h"
#include "DismembermentGraph/DismembermentGraphNode.h"
#include "DismembermentGraph/DismembermentGraphNodeCut.h"
#include "DismembermentGraph/DismembermentGraphNodeBoneSelect.h"
#include "DismembermentGraph/DismembermentGraphNodeBloodEffect.h"
#include "AssetToolsModule.h"
#include "IAssetTools.h"
#include "EdGraphUtilities.h"
#include "Modules/ModuleManager.h"
#include "GraphEditorActions.h"
#include "AssetTypeActions_Base.h"

#define LOCTEXT_NAMESPACE "DismembermentGraphEditor"

/**
 * Dismemberment graph schema actions factory
 */
class FDismembermentGraphSchemaActionMenuBuilder : public FGraphSchemaActionMenuBuilder
{
public:
    // Constructor
    FDismembermentGraphSchemaActionMenuBuilder(const UEdGraphSchema* InSchema)
        : FGraphSchemaActionMenuBuilder(InSchema)
    {
    }
};

/**
 * Dismemberment graph editor module
 */
class FDismembermentGraphEditorModule : public IModuleInterface
{
public:
    // IModuleInterface interface
    virtual void StartupModule() override
    {
        // Register asset type actions
        IAssetTools& AssetTools = FModuleManager::LoadModuleChecked<FAssetToolsModule>("AssetTools").Get();
        
        // Register asset category
        DismembermentGraphAssetCategoryBit = AssetTools.RegisterAdvancedAssetCategory(FName(TEXT("Dismemberment")), LOCTEXT("DismembermentGraphAssetCategory", "Dismemberment"));
        
        // Register asset type actions
        TSharedPtr<IAssetTypeActions> Action = MakeShareable(new FDismembermentGraphAssetTypeActions(DismembermentGraphAssetCategoryBit));
        AssetTools.RegisterAssetTypeActions(Action);
        RegisteredAssetTypeActions.Add(Action);
        
        // Register graph editor factory
        TSharedPtr<FGraphPanelNodeFactory> NodeFactory = MakeShareable(new FDismembermentGraphNodeFactory);
        FEdGraphUtilities::RegisterVisualNodeFactory(NodeFactory);
        
        // Register node factories
        RegisterNodeFactory(UDismembermentGraphNodeCut::StaticClass(), LOCTEXT("CutNode", "Cut Node"), LOCTEXT("CutNodeTooltip", "Performs a cut operation"));
        RegisterNodeFactory(UDismembermentGraphNodeBoneSelect::StaticClass(), LOCTEXT("BoneSelectNode", "Bone Select Node"), LOCTEXT("BoneSelectNodeTooltip", "Selects bones for operations"));
        RegisterNodeFactory(UDismembermentGraphNodeBloodEffect::StaticClass(), LOCTEXT("BloodEffectNode", "Blood Effect Node"), LOCTEXT("BloodEffectNodeTooltip", "Creates blood effects"));
    }
    
    virtual void ShutdownModule() override
    {
        // Unregister asset type actions
        if (FModuleManager::Get().IsModuleLoaded("AssetTools"))
        {
            IAssetTools& AssetTools = FModuleManager::GetModuleChecked<FAssetToolsModule>("AssetTools").Get();
            for (TSharedPtr<IAssetTypeActions> Action : RegisteredAssetTypeActions)
            {
                AssetTools.UnregisterAssetTypeActions(Action.ToSharedRef());
            }
        }
        
        RegisteredAssetTypeActions.Empty();
        
        // Unregister node factories
        for (const auto& Factory : RegisteredNodeFactories)
        {
            FEdGraphUtilities::UnregisterVisualNodeFactory(Factory);
        }
        
        RegisteredNodeFactories.Empty();
    }
    // End of IModuleInterface interface
    
private:
    // Asset category bit
    EAssetTypeCategories::Type DismembermentGraphAssetCategoryBit;
    
    // Registered asset type actions
    TArray<TSharedPtr<IAssetTypeActions>> RegisteredAssetTypeActions;
    
    // Registered node factories
    TArray<TSharedPtr<FGraphPanelNodeFactory>> RegisteredNodeFactories;
    
    // Register a node factory
    void RegisterNodeFactory(UClass* NodeClass, const FText& DisplayName, const FText& Tooltip)
    {
        TSharedPtr<FDismembermentGraphNodeFactory> NodeFactory = MakeShareable(new FDismembermentGraphNodeFactory(NodeClass, DisplayName, Tooltip));
        FEdGraphUtilities::RegisterVisualNodeFactory(NodeFactory);
        RegisteredNodeFactories.Add(NodeFactory);
    }
};

/**
 * Dismemberment graph asset type actions
 */
class FDismembermentGraphAssetTypeActions : public FAssetTypeActions_Base
{
public:
    // Constructor
    FDismembermentGraphAssetTypeActions(EAssetTypeCategories::Type InAssetCategory)
        : AssetCategory(InAssetCategory)
    {
    }
    
    // FAssetTypeActions_Base interface
    virtual FText GetName() const override { return LOCTEXT("DismembermentGraphAssetName", "Dismemberment Graph"); }
    virtual FColor GetTypeColor() const override { return FColor(255, 64, 64); }
    virtual UClass* GetSupportedClass() const override { return UDismembermentGraphAsset::StaticClass(); }
    virtual uint32 GetCategories() override { return AssetCategory; }
    virtual bool HasActions(const TArray<UObject*>& InObjects) const override { return false; }
    virtual void OpenAssetEditor(const TArray<UObject*>& InObjects, TSharedPtr<IToolkitHost> EditWithinLevelEditor) override
    {
        for (UObject* Object : InObjects)
        {
            if (UDismembermentGraphAsset* GraphAsset = Cast<UDismembermentGraphAsset>(Object))
            {
                TSharedRef<FDismembermentGraphEditor> NewEditor = MakeShareable(new FDismembermentGraphEditor());
                NewEditor->InitDismembermentGraphEditor(EToolkitMode::Standalone, EditWithinLevelEditor, GraphAsset);
            }
        }
    }
    // End of FAssetTypeActions_Base interface
    
private:
    // Asset category
    EAssetTypeCategories::Type AssetCategory;
};

/**
 * Dismemberment graph node factory
 */
class FDismembermentGraphNodeFactory : public FGraphPanelNodeFactory
{
public:
    // Constructor
    FDismembermentGraphNodeFactory()
        : NodeClass(nullptr)
    {
    }
    
    FDismembermentGraphNodeFactory(UClass* InNodeClass, const FText& InDisplayName, const FText& InTooltip)
        : NodeClass(InNodeClass)
        , DisplayName(InDisplayName)
        , Tooltip(InTooltip)
    {
    }
    
    // FGraphPanelNodeFactory interface
    virtual TSharedPtr<SGraphNode> CreateNode(UEdGraphNode* Node) const override
    {
        if (NodeClass && Node->IsA(NodeClass))
        {
            return SNew(SDismembermentGraphNode, Node);
        }
        
        return nullptr;
    }
    // End of FGraphPanelNodeFactory interface
    
private:
    // Node class
    UClass* NodeClass;
    
    // Display name
    FText DisplayName;
    
    // Tooltip
    FText Tooltip;
};

IMPLEMENT_MODULE(FDismembermentGraphEditorModule, FLESHEditor)

#undef LOCTEXT_NAMESPACE