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

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

#include "Graph/InworldGraph.h"

UInworldGraphRuntimeData_Knowledge::UInworldGraphRuntimeData_Knowledge()
{
	Key = "Knowledge";
}

FInworldDataHandle UInworldGraphRuntimeData_Knowledge::GetData()
{
	FInworldKnowledgeData KnowledgeData;
	EnabledKnowledge.GenerateValueArray(KnowledgeData.TextEmbeddings);
	EnabledKnowledgeAssets.GenerateValueArray(KnowledgeData.TextEmbeddingAssets);
	return FInworldData_Struct(KnowledgeData);
}

void UInworldGraphRuntimeData_Knowledge::Initialize_Implementation()
{
	for (const TTuple<FName, TObjectPtr<const UInworldTextEmbeddingAsset>>& Knowledge : InitialKnowledge)
	{
		AddKnowledgeTextEmbeddingAsset(Knowledge.Key, Knowledge.Value);
	}
	Super::Initialize_Implementation();
}

void UInworldGraphRuntimeData_Knowledge::Set(const TMap<FName, UInworldTextEmbeddingAsset*>& KnowledgeAssets)
{
	EnabledKnowledge = {};
	DisabledKnowledge = {};
	EnabledKnowledgeAssets = {};
	DisabledKnowledgeAssets = {};
	for (const TTuple<FName, UInworldTextEmbeddingAsset*>& Knowledge : KnowledgeAssets)
	{
		AddKnowledgeTextEmbeddingAsset(Knowledge.Key, Knowledge.Value);
	}
}

void UInworldGraphRuntimeData_Knowledge::AddKnowledgeTextEmbeddingAsset(
	FName Id, const UInworldTextEmbeddingAsset* Knowledge, bool Replace)
{
	const bool bIsEnabled{EnabledKnowledge.Contains(Id) || EnabledKnowledgeAssets.Contains(Id)};
	const bool bIsDisabled{DisabledKnowledge.Contains(Id) || DisabledKnowledgeAssets.Contains(Id)};
	const bool bHasKnowledgeId{bIsEnabled || bIsDisabled};
	if (bHasKnowledgeId && !Replace)
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Knowledge::AddKnowledgeSet - "
				 "Failed to add knowledge. A knowledge with ID %s already exists."),
			*Id.ToString());
		return;
	}
	if (bHasKnowledgeId)
	{
		RemoveKnowledge(Id);
	}

	if (bIsDisabled)
	{
		DisabledKnowledgeAssets.Add(Id, Knowledge);
	}
	else
	{
		EnabledKnowledgeAssets.Add(Id, Knowledge);
	}
}

void UInworldGraphRuntimeData_Knowledge::AddKnowledgeTextEmbeddings(
	FName Id, const TArray<FInworldTextEmbedding>& Knowledge, bool Replace)
{
	const bool bIsEnabled{EnabledKnowledge.Contains(Id) || EnabledKnowledgeAssets.Contains(Id)};
	const bool bIsDisabled{DisabledKnowledge.Contains(Id) || DisabledKnowledgeAssets.Contains(Id)};
	const bool bHasKnowledgeId{bIsEnabled || bIsDisabled};
	if (bHasKnowledgeId && !Replace)
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Knowledge::AddKnowledgeSet - "
				 "Failed to add knowledge. A knowledge with ID %s already exists."),
			*Id.ToString());
		return;
	}
	if (bHasKnowledgeId)
	{
		RemoveKnowledge(Id);
	}

	TSharedPtr<const TArray<FInworldTextEmbedding>> SharedTextEmbeddings =
		MakeShared<const TArray<FInworldTextEmbedding>>(Knowledge);
	if (bIsDisabled)
	{
		DisabledKnowledge.Add(Id, SharedTextEmbeddings);
	}
	else
	{
		EnabledKnowledge.Add(Id, SharedTextEmbeddings);
	}
}

void UInworldGraphRuntimeData_Knowledge::RemoveKnowledge(FName Id)
{
	if (EnabledKnowledge.Contains(Id))
	{
		EnabledKnowledge.Remove(Id);
	}
	else if (EnabledKnowledgeAssets.Contains(Id))
	{
		EnabledKnowledgeAssets.Remove(Id);
	}
	else if (DisabledKnowledge.Contains(Id))
	{
		DisabledKnowledge.Remove(Id);
	}
	else if (DisabledKnowledgeAssets.Contains(Id))
	{
		DisabledKnowledgeAssets.Remove(Id);
	}
	else
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Knowledge::RemoveKnowledgeSet - "
				 "A knowledge with ID %s could not be found."),
			*Id.ToString());
	}
}

void UInworldGraphRuntimeData_Knowledge::EnableKnowledge(FName Id)
{
	if (DisabledKnowledge.Contains(Id))
	{
		EnabledKnowledge.Add(Id, DisabledKnowledge[Id]);
		DisabledKnowledge.Remove(Id);
	}
	else if (DisabledKnowledgeAssets.Contains(Id))
	{
		EnabledKnowledgeAssets.Add(Id, DisabledKnowledgeAssets[Id]);
		DisabledKnowledgeAssets.Remove(Id);
	}
	else
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Knowledge::EnabledKnowledgeSet - "
				 "A disabled knowledge with ID %s could not be found."),
			*Id.ToString());
	}
}

void UInworldGraphRuntimeData_Knowledge::DisableKnowledge(FName Id)
{
	if (EnabledKnowledge.Contains(Id))
	{
		DisabledKnowledge.Add(Id, EnabledKnowledge[Id]);
		EnabledKnowledge.Remove(Id);
	}
	else if (EnabledKnowledgeAssets.Contains(Id))
	{
		DisabledKnowledgeAssets.Add(Id, EnabledKnowledgeAssets[Id]);
		EnabledKnowledgeAssets.Remove(Id);
	}
	else
	{
		UE_LOG(LogInworldGraph, Error,
			TEXT("UInworldGraphRuntimeData_Knowledge::DisableKnowledgeSet - "
				 "An enabled knowledge with ID %s could not be found."),
			*Id.ToString());
	}
}

TArray<FName> UInworldGraphRuntimeData_Knowledge::GetKnowledgeMapKeys() const
{
	TArray<FName> AllKeys;

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

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

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

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

	return AllKeys;
}
