/**
 * Copyright 2022-2025 Theai, Inc. dba Inworld AI
 */

#include "InworldRuntime/Public/Graph/RuntimeData/InworldGraphRuntimeData_Intent.h"

#include "Graph/InworldGraph.h"

UInworldGraphRuntimeData_Intent::UInworldGraphRuntimeData_Intent()
{
	Key = "Intents";
}

FInworldDataHandle UInworldGraphRuntimeData_Intent::GetData()
{
	FInworldIntentsData Intents;
	Intents.TextEmbeddings = EnabledIntents;
	Intents.TextEmbeddingAssets = EnabledIntentAssets;
	return FInworldData_Struct(Intents);
}

void UInworldGraphRuntimeData_Intent::Initialize_Implementation()
{
	for (const TTuple<FName, TObjectPtr<const UInworldTextEmbeddingAsset>>& Intent : InitialIntents)
	{
		AddIntentTextEmbeddingAsset(Intent.Key, Intent.Value);
	}
	Super::Initialize_Implementation();
}

void UInworldGraphRuntimeData_Intent::Set(const TMap<FName, UInworldTextEmbeddingAsset*>& IntentAssets)
{
	EnabledIntents = {};
	DisabledIntents = {};
	EnabledIntentAssets = {};
	DisabledIntentAssets = {};
	for (const TTuple<FName, UInworldTextEmbeddingAsset*>& Intent : IntentAssets)
	{
		AddIntentTextEmbeddingAsset(Intent.Key, Intent.Value);
	}
}

void UInworldGraphRuntimeData_Intent::AddIntentTextEmbeddingAsset(FName Id, const UInworldTextEmbeddingAsset* Intent, bool Replace)
{
	const bool bIsEnabled{EnabledIntents.Contains(Id) || EnabledIntentAssets.Contains(Id)};
	const bool bIsDisabled{DisabledIntents.Contains(Id) || DisabledIntentAssets.Contains(Id)};
	const bool bHasIntentId{bIsEnabled || bIsDisabled};
	if (bHasIntentId && !Replace)
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Intent::AddIntentSet - "
				 "Failed to add Intent. A Intent with ID %s already exists."),
			*Id.ToString());
		return;
	}
	if (bHasIntentId)
	{
		RemoveIntent(Id);
	}

	if (bIsDisabled)
	{
		DisabledIntentAssets.Add(Id, Intent);
	}
	else
	{
		EnabledIntentAssets.Add(Id, Intent);
	}
}

void UInworldGraphRuntimeData_Intent::AddIntentTextEmbeddings(FName Id, const TArray<FInworldTextEmbedding>& Intent, bool Replace)
{
	const bool bIsEnabled{EnabledIntents.Contains(Id) || EnabledIntentAssets.Contains(Id)};
	const bool bIsDisabled{DisabledIntents.Contains(Id) || DisabledIntentAssets.Contains(Id)};
	const bool bHasIntentId{bIsEnabled || bIsDisabled};
	if (bHasIntentId && !Replace)
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Intent::AddIntentSet - "
				 "Failed to add Intent. A Intent with ID %s already exists."),
			*Id.ToString());
		return;
	}
	if (bHasIntentId)
	{
		RemoveIntent(Id);
	}

	TSharedPtr<const TArray<FInworldTextEmbedding>> SharedTextEmbeddings = MakeShared<const TArray<FInworldTextEmbedding>>(Intent);
	if (bIsDisabled)
	{
		DisabledIntents.Add(Id, SharedTextEmbeddings);
	}
	else
	{
		EnabledIntents.Add(Id, SharedTextEmbeddings);
	}
}

void UInworldGraphRuntimeData_Intent::RemoveIntent(FName Id)
{
	if (EnabledIntents.Contains(Id))
	{
		EnabledIntents.Remove(Id);
	}
	else if (EnabledIntentAssets.Contains(Id))
	{
		EnabledIntentAssets.Remove(Id);
	}
	else if (DisabledIntents.Contains(Id))
	{
		DisabledIntents.Remove(Id);
	}
	else if (DisabledIntentAssets.Contains(Id))
	{
		DisabledIntentAssets.Remove(Id);
	}
	else
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Intent::RemoveIntentSet - "
				 "A Intent with ID %s could not be found."),
			*Id.ToString());
	}
}

void UInworldGraphRuntimeData_Intent::EnableIntent(FName Id)
{
	if (DisabledIntents.Contains(Id))
	{
		EnabledIntents.Add(Id, DisabledIntents[Id]);
		DisabledIntents.Remove(Id);
	}
	else if (DisabledIntentAssets.Contains(Id))
	{
		EnabledIntentAssets.Add(Id, DisabledIntentAssets[Id]);
		DisabledIntentAssets.Remove(Id);
	}
	else
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Intent::EnabledIntentSet - "
				 "A disabled Intent with ID %s could not be found."),
			*Id.ToString());
	}
}

void UInworldGraphRuntimeData_Intent::DisableIntent(FName Id)
{
	if (EnabledIntents.Contains(Id))
	{
		DisabledIntents.Add(Id, EnabledIntents[Id]);
		EnabledIntents.Remove(Id);
	}
	else if (EnabledIntentAssets.Contains(Id))
	{
		DisabledIntentAssets.Add(Id, EnabledIntentAssets[Id]);
		EnabledIntentAssets.Remove(Id);
	}
	else
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Intent::DisableIntentSet - "
				 "An enabled Intent with ID %s could not be found."),
			*Id.ToString());
	}
}

TArray<FName> UInworldGraphRuntimeData_Intent::GetIntentMapKeys() const
{
	TArray<FName> AllKeys;

	TArray<FName> EnabledKeys;
	EnabledIntents.GenerateKeyArray(EnabledKeys);
	AllKeys.Append(EnabledKeys);

	TArray<FName> DisabledKeys;
	DisabledIntents.GenerateKeyArray(DisabledKeys);
	AllKeys.Append(DisabledKeys);

	TArray<FName> EnabledAssetKeys;
	EnabledIntentAssets.GenerateKeyArray(EnabledAssetKeys);
	AllKeys.Append(EnabledAssetKeys);

	TArray<FName> DisabledAssetKeys;
	DisabledIntentAssets.GenerateKeyArray(DisabledAssetKeys);
	AllKeys.Append(DisabledAssetKeys);

	return AllKeys;
}
