#include "Extensions/DruidsMCPBridgeSageExtensionRegistry.h"

#include "AssetRegistry/AssetData.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "CoreGlobals.h"
#include "DruidsMCPBridge.h"
#include "Engine/Blueprint.h"
#include "UObject/UObjectGlobals.h"
#include "UObject/UObjectIterator.h"

namespace
{
	UDruidsMCPBridgeSageExtensionRegistry* GDruidsMCPBridgeSageExtensionRegistry = nullptr;
}

UDruidsMCPBridgeSageExtensionRegistry* UDruidsMCPBridgeSageExtensionRegistry::Get()
{
	if (!GDruidsMCPBridgeSageExtensionRegistry)
	{
		GDruidsMCPBridgeSageExtensionRegistry = NewObject<UDruidsMCPBridgeSageExtensionRegistry>(GetTransientPackage());
		GDruidsMCPBridgeSageExtensionRegistry->AddToRoot();
	}
	return GDruidsMCPBridgeSageExtensionRegistry;
}

void UDruidsMCPBridgeSageExtensionRegistry::Shutdown()
{
	UDruidsMCPBridgeSageExtensionRegistry* Registry = GDruidsMCPBridgeSageExtensionRegistry;
	if (!Registry)
	{
		return;
	}

	GDruidsMCPBridgeSageExtensionRegistry = nullptr;

	if (GExitPurge || IsEngineExitRequested())
	{
		return;
	}

	Registry->Reset();
	if (Registry->IsRooted())
	{
		Registry->RemoveFromRoot();
	}
}

void UDruidsMCPBridgeSageExtensionRegistry::Reset()
{
	FScopeLock Lock(&ExtensionsLock);
	RegisteredExtensions.Empty();
}

void UDruidsMCPBridgeSageExtensionRegistry::Refresh()
{
	FScopeLock Lock(&ExtensionsLock);
	RegisteredExtensions.Empty();
	DiscoverNativeExtensions();
	DiscoverBlueprintExtensions();
}

TArray<TWeakObjectPtr<UDruidsMCPBridgeSageExtension>> UDruidsMCPBridgeSageExtensionRegistry::GetActiveExtensions()
{
	Refresh();

	TArray<TWeakObjectPtr<UDruidsMCPBridgeSageExtension>> ActiveExtensions;
	TArray<TObjectPtr<UDruidsMCPBridgeSageExtension>> Extensions;
	RegisteredExtensions.GenerateValueArray(Extensions);
	for (TObjectPtr<UDruidsMCPBridgeSageExtension> Extension : Extensions)
	{
		if (Extension && Extension->IsEnabled())
		{
			ActiveExtensions.Add(Extension.Get());
		}
	}

	return ActiveExtensions;
}

TArray<TSharedPtr<FDruidsMCPBridgeSageExtensionDefinition>> UDruidsMCPBridgeSageExtensionRegistry::GetAllExtensionDefinitions()
{
	TArray<TSharedPtr<FDruidsMCPBridgeSageExtensionDefinition>> Definitions;
	for (TWeakObjectPtr<UDruidsMCPBridgeSageExtension> Extension : GetActiveExtensions())
	{
		if (Extension.IsValid())
		{
			Definitions.Add(Extension->GetExtensionDefinition());
		}
	}
	return Definitions;
}

TWeakObjectPtr<UDruidsMCPBridgeSageExtension> UDruidsMCPBridgeSageExtensionRegistry::GetExtensionForId(const FString& ExtensionId)
{
	Refresh();

	if (const TObjectPtr<UDruidsMCPBridgeSageExtension>* Extension = RegisteredExtensions.Find(ExtensionId))
	{
		return Extension->Get();
	}

	for (const TPair<FString, TObjectPtr<UDruidsMCPBridgeSageExtension>>& Pair : RegisteredExtensions)
	{
		if (Pair.Key.EndsWith(ExtensionId) || Pair.Key.Contains(ExtensionId))
		{
			return Pair.Value.Get();
		}
	}

	return nullptr;
}

void UDruidsMCPBridgeSageExtensionRegistry::DiscoverNativeExtensions()
{
	for (TObjectIterator<UClass> ClassIt; ClassIt; ++ClassIt)
	{
		UClass* Class = *ClassIt;
		if (!Class || Class == UDruidsMCPBridgeSageExtension::StaticClass())
		{
			continue;
		}
		if (!Class->IsChildOf(UDruidsMCPBridgeSageExtension::StaticClass()))
		{
			continue;
		}
		if (Class->HasAnyClassFlags(CLASS_Abstract | CLASS_Deprecated | CLASS_NewerVersionExists))
		{
			continue;
		}

		if (UDruidsMCPBridgeSageExtension* Extension = NewObject<UDruidsMCPBridgeSageExtension>(GetTransientPackage(), Class))
		{
			RegisterExtension(Extension);
		}
	}
}

void UDruidsMCPBridgeSageExtensionRegistry::DiscoverBlueprintExtensions()
{
	FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(TEXT("AssetRegistry"));
	IAssetRegistry& AssetRegistry = AssetRegistryModule.Get();

	if (AssetRegistry.IsLoadingAssets())
	{
		return;
	}

	FARFilter Filter;
	Filter.ClassPaths.Add(FTopLevelAssetPath(TEXT("/Script/Blutility"), TEXT("EditorUtilityBlueprint")));
	Filter.bRecursiveClasses = true;
	Filter.bIncludeOnlyOnDiskAssets = false;

	TArray<FAssetData> AssetDataList;
	AssetRegistry.GetAssets(Filter, AssetDataList);

	for (const FAssetData& AssetData : AssetDataList)
	{
		FString ParentClassPath;
		if (!AssetData.GetTagValue(FName(TEXT("ParentClass")), ParentClassPath))
		{
			continue;
		}

		UObject* ParentClassObject = LoadObject<UObject>(nullptr, *ParentClassPath);
		UClass* ParentClass = Cast<UClass>(ParentClassObject);
		if (!ParentClass || !ParentClass->IsChildOf(UDruidsMCPBridgeSageExtension::StaticClass()))
		{
			continue;
		}

		UBlueprint* Blueprint = Cast<UBlueprint>(AssetData.GetAsset());
		if (!Blueprint || !Blueprint->GeneratedClass)
		{
			continue;
		}
		if (!Blueprint->GeneratedClass->IsChildOf(UDruidsMCPBridgeSageExtension::StaticClass()))
		{
			continue;
		}

		if (UDruidsMCPBridgeSageExtension* Extension = NewObject<UDruidsMCPBridgeSageExtension>(GetTransientPackage(), Blueprint->GeneratedClass))
		{
			RegisterExtension(Extension);
		}
	}
}

void UDruidsMCPBridgeSageExtensionRegistry::RegisterExtension(UDruidsMCPBridgeSageExtension* Extension)
{
	if (!Extension)
	{
		return;
	}

	Extension->DiscoverActionsAndQueries();
	TSharedPtr<FDruidsMCPBridgeSageExtensionDefinition> Definition = Extension->GetExtensionDefinition();
	if (!Definition.IsValid() || Definition->ExtensionId.IsEmpty())
	{
		return;
	}

	RegisteredExtensions.Add(Definition->ExtensionId, Extension);
	UE_LOG(LogDruidsMCPBridge, Verbose, TEXT("Registered Sage extension: %s"), *Definition->ExtensionId);
}
