// © 2024 Long Ly. All rights reserved. Any unauthorized use, reproduction, or distribution of this trademark is strictly prohibited and may result in legal action.


#include "GameFeatures/OLSGameFeatureAction_AddInputContextMapping.h"

#include "UserSettings/EnhancedInputUserSettings.h"
#include "EnhancedInputSubsystems.h"
#include "Systems/OLSAssetManager.h"
#include "Engine/LocalPlayer.h"

#if WITH_EDITOR
#include "Misc/DataValidation.h"
#endif

#include "InputMappingContext.h"
#include "OLSLog.h"
#include "Components/GameFrameworkComponentManager.h"
#include "Components/OLSHeroComponent.h"

DEFINE_LOG_CATEGORY(LogOLSGameFA_AddInputContextMapping);

#include UE_INLINE_GENERATED_CPP_BY_NAME(OLSGameFeatureAction_AddInputContextMapping)

#define LOCTEXT_NAMESPACE "GameFeatures"

void UOLSGameFeatureAction_AddInputContextMapping::OnGameFeatureRegistering()
{
	Super::OnGameFeatureRegistering();
	
	RegisterInputMappingContexts();
}

void UOLSGameFeatureAction_AddInputContextMapping::OnGameFeatureActivating(FGameFeatureActivatingContext& context)
{
	FPerContextData& activeData = ContextData.FindOrAdd(context);
	if (!ensure(activeData.ExtensionRequestHandles.IsEmpty()) ||
		!ensure(activeData.ControllersAddedTo.IsEmpty()))
	{
		Reset(activeData);
	}
	
	Super::OnGameFeatureActivating(context);
}

void UOLSGameFeatureAction_AddInputContextMapping::OnGameFeatureDeactivating(FGameFeatureDeactivatingContext& context)
{
	Super::OnGameFeatureDeactivating(context);

	FPerContextData* activeData = ContextData.Find(context);
	if (ensure(activeData))
	{
		Reset(*activeData);
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::OnGameFeatureUnregistering()
{
	Super::OnGameFeatureUnregistering();

	UnregisterInputMappingContexts();
}

#if WITH_EDITOR
EDataValidationResult UOLSGameFeatureAction_AddInputContextMapping::IsDataValid(FDataValidationContext& context) const
{
	// Don't call Super since it does not fit in this.
	
	EDataValidationResult result = CombineDataValidationResults(Super::IsDataValid(context), EDataValidationResult::Valid);

	int32 index = 0;

	for (const FOLSInputMappingContextAndPriority& entry : InputMappings)
	{
		if (entry.InputMapping.IsNull())
		{
			result = EDataValidationResult::Invalid;
			context.AddError(FText::Format(LOCTEXT("NullInputMapping", "Null InputMapping at index {0}."), index));
		}
		++index;
	}

	return result;
}
#endif

void UOLSGameFeatureAction_AddInputContextMapping::RegisterInputMappingContexts()
{
	RegisterInputContextMappingsForGameInstanceHandle = FWorldDelegates::OnStartGameInstance.AddUObject(
		this, &UOLSGameFeatureAction_AddInputContextMapping::RegisterInputContextMappingsForGameInstance);

	const TIndirectArray<FWorldContext>& worldContexts = GEngine->GetWorldContexts();
	for (TIndirectArray<FWorldContext>::TConstIterator worldContextIterator = worldContexts.CreateConstIterator();
	     worldContextIterator; ++worldContextIterator)
	{
		RegisterInputContextMappingsForGameInstance(worldContextIterator->OwningGameInstance);
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::RegisterInputContextMappingsForGameInstance(
	UGameInstance* gameInstance)
{
	if (gameInstance != nullptr && !gameInstance->OnLocalPlayerAddedEvent.IsBoundToObject(this))
	{
		gameInstance->OnLocalPlayerAddedEvent.AddUObject(
			this, &UOLSGameFeatureAction_AddInputContextMapping::RegisterInputMappingContextsForLocalPlayer);
		gameInstance->OnLocalPlayerRemovedEvent.AddUObject(
			this, &UOLSGameFeatureAction_AddInputContextMapping::UnregisterInputMappingContextsForLocalPlayer);

		for (TArray<ULocalPlayer*>::TConstIterator localPlayerIterator = gameInstance->GetLocalPlayerIterator();
		     localPlayerIterator; ++localPlayerIterator)
		{
			RegisterInputMappingContextsForLocalPlayer(*localPlayerIterator);
		}
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::RegisterInputMappingContextsForLocalPlayer(ULocalPlayer* localPlayer)
{
	if (ensure(localPlayer))
	{
		UOLSAssetManager& assetManager = UOLSAssetManager::Get();
		
		if (UEnhancedInputLocalPlayerSubsystem* eiSubsystem = ULocalPlayer::GetSubsystem<UEnhancedInputLocalPlayerSubsystem>(localPlayer))
		{
			if (UEnhancedInputUserSettings* settings = eiSubsystem->GetUserSettings())
			{
				for (const FOLSInputMappingContextAndPriority& entry : InputMappings)
				{
					// Skip entries that don't want to be registered
					if (!entry.bShouldRegisterWithSettings)
					{
						continue;
					}

					// Register this IMC with the settings!
					if (UInputMappingContext* imc = assetManager.GetAsset(entry.InputMapping))
					{
						settings->RegisterInputMappingContext(imc);
					}
				}
			}
		}
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::UnregisterInputMappingContexts()
{
	FWorldDelegates::OnStartGameInstance.Remove(RegisterInputContextMappingsForGameInstanceHandle);
	RegisterInputContextMappingsForGameInstanceHandle.Reset();

	const TIndirectArray<FWorldContext>& worldContexts = GEngine->GetWorldContexts();
	for (TIndirectArray<FWorldContext>::TConstIterator worldContextIterator = worldContexts.CreateConstIterator();
	     worldContextIterator; ++worldContextIterator)
	{
		UnregisterInputContextMappingsForGameInstance(worldContextIterator->OwningGameInstance);
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::UnregisterInputContextMappingsForGameInstance(
	UGameInstance* gameInstance)
{
	if (gameInstance)
	{
		gameInstance->OnLocalPlayerAddedEvent.RemoveAll(this);
		gameInstance->OnLocalPlayerRemovedEvent.RemoveAll(this);

		for (TArray<ULocalPlayer*>::TConstIterator localPlayerIterator = gameInstance->GetLocalPlayerIterator();
		     localPlayerIterator; ++localPlayerIterator)
		{
			UnregisterInputMappingContextsForLocalPlayer(*localPlayerIterator);
		}
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::UnregisterInputMappingContextsForLocalPlayer(
	ULocalPlayer* localPlayer)
{
	if (ensure(localPlayer))
	{
		if (UEnhancedInputLocalPlayerSubsystem* eiSubsystem = ULocalPlayer::GetSubsystem<UEnhancedInputLocalPlayerSubsystem>(localPlayer))
		{
			if (UEnhancedInputUserSettings* settings = eiSubsystem->GetUserSettings())
			{
				for (const FOLSInputMappingContextAndPriority& entry : InputMappings)
				{
					// Skip entries that don't want to be registered
					if (!entry.bShouldRegisterWithSettings)
					{
						continue;
					}

					// Register this IMC with the settings!
					if (UInputMappingContext* imc = entry.InputMapping.Get())
					{
						settings->UnregisterInputMappingContext(imc);
					}
				}
			}
		}
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::AddToWorld(const FWorldContext& worldContext,
                                                              const FGameFeatureStateChangeContext& changeContext)
{
	UWorld* world = worldContext.World();
	UGameInstance* gameInstance = worldContext.OwningGameInstance;
	FPerContextData& activeData = ContextData.FindOrAdd(changeContext);

	if (gameInstance && world && world->IsGameWorld())
	{
		if (UGameFrameworkComponentManager* componentManager = UGameInstance::GetSubsystem<
			UGameFrameworkComponentManager>(gameInstance))
		{
			UGameFrameworkComponentManager::FExtensionHandlerDelegate addAbilitiesDelegate =
				UGameFrameworkComponentManager::FExtensionHandlerDelegate::CreateUObject(
					this, &ThisClass::HandleControllerExtension, changeContext);
			TSharedPtr<FComponentRequestHandle> extensionRequestHandle =
				componentManager->AddExtensionHandler(APlayerController::StaticClass(), addAbilitiesDelegate);

			activeData.ExtensionRequestHandles.Add(extensionRequestHandle);
		}
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::Reset(FPerContextData& activeData)
{
	activeData.ExtensionRequestHandles.Empty();

	while (!activeData.ControllersAddedTo.IsEmpty())
	{
		TWeakObjectPtr<APlayerController> controllerPtr = activeData.ControllersAddedTo.Top();
		if (controllerPtr.IsValid())
		{
			RemoveInputMapping(controllerPtr.Get(), activeData);
		}
		else
		{
			activeData.ControllersAddedTo.Pop();
		}
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::HandleControllerExtension(AActor* actor,
                                                                             FName eventName,
                                                                             FGameFeatureStateChangeContext changeContext)
{
	APlayerController* playerController = CastChecked<APlayerController>(actor);
	FPerContextData& activeData = ContextData.FindOrAdd(changeContext);

	// TODO Why does this code mix and match controllers and local players? ControllersAddedTo is never modified
	if ((eventName == UGameFrameworkComponentManager::NAME_ExtensionRemoved) || (eventName == UGameFrameworkComponentManager::NAME_ReceiverRemoved))
	{
		RemoveInputMapping(playerController, activeData);
	}
	else if ((eventName == UGameFrameworkComponentManager::NAME_ExtensionAdded) || (eventName == UOLSHeroComponent::NAME_BindInputsNow))
	{
		AddInputMappingForPlayer(playerController->GetLocalPlayer(), activeData);
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::AddInputMappingForPlayer(UPlayer* player,
                                                                            FPerContextData& activeData)
{
	if (ULocalPlayer* localPlayer = Cast<ULocalPlayer>(player))
	{
		if (UEnhancedInputLocalPlayerSubsystem* inputSystem = localPlayer->GetSubsystem<
			UEnhancedInputLocalPlayerSubsystem>())
		{
			for (const FOLSInputMappingContextAndPriority& entry : InputMappings)
			{
				if (const UInputMappingContext* imc = entry.InputMapping.Get())
				{
					inputSystem->AddMappingContext(imc, entry.Priority);
				}
			}
		}
		else
		{
			OLS_LOG(LogOLSGameFA_AddInputContextMapping, Error,
			        TEXT(
				        "Failed to find `UEnhancedInputLocalPlayerSubsystem` for local player. Input mappings will not be added. Make sure you're set to use the EnhancedInput system via config file."
			        ));
		}
	}
}

void UOLSGameFeatureAction_AddInputContextMapping::RemoveInputMapping(APlayerController* playerController,
                                                                      FPerContextData& activeData)
{
	if (ULocalPlayer* localPlayer = playerController->GetLocalPlayer())
	{
		if (UEnhancedInputLocalPlayerSubsystem* inputSystem = localPlayer->GetSubsystem<UEnhancedInputLocalPlayerSubsystem>())
		{
			for (const FOLSInputMappingContextAndPriority& entry : InputMappings)
			{
				if (const UInputMappingContext* imc = entry.InputMapping.Get())
				{
					inputSystem->RemoveMappingContext(imc);
				}
			}
		}
	}

	activeData.ControllersAddedTo.Remove(playerController);
}

#undef LOCTEXT_NAMESPACE