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

#include "Conversation/InworldConversationGroup.h"

#include "Director/SpeakerRuleDirector.h"
#include "GameFramework/PlayerController.h"
#include "Graph/RuntimeData/InworldGraphRuntimeData_EventHistory.h"
#include "InworldBlueprintFunctionLibrary.h"
#include "InworldCharacter.h"
#include "InworldCharacterBFL.h"

void UInworldConversationGroup::CreateConversation(UObject* WorldContextObject, UObject* Owner,
	FOnInworldConversationGroupCreated Callback, UInworldGraphAsset* SpeakerSelectionGraphAsset)
{
	FOnInworldConversationGroupCreatedNative NativeCallback;
	if (Callback.IsBound())
	{
		NativeCallback.BindUFunction(Callback.GetUObject(), Callback.GetFunctionName());
	}
	CreateConversation(WorldContextObject, Owner, NativeCallback, SpeakerSelectionGraphAsset);
}

void UInworldConversationGroup::CreateConversation(UObject* WorldContextObject, UObject* Owner,
	FOnInworldConversationGroupCreatedNative NativeCallback, UInworldGraphAsset* SpeakerSelectionGraphAsset)
{
	UInworldConversationGroup* ConversationGroup = NewObject<UInworldConversationGroup>(Owner);
	ConversationGroup->EventHistory = NewObject<UInworldGraphRuntimeData_EventHistory>(ConversationGroup);
	if (IsValid(SpeakerSelectionGraphAsset))
	{
		SpeakerSelectionGraphAsset->GetGraphInstance(FOnGraphCompiledNative::CreateWeakLambda(ConversationGroup,
			[ConversationGroup, NativeCallback](UInworldGraph* Graph, bool bSuccess)
			{
				if (ConversationGroup && bSuccess)
				{
					ConversationGroup->SpeakerSelectionGraph = Graph;
					NativeCallback.ExecuteIfBound(ConversationGroup, bSuccess);
				}
				else
				{
					UE_LOG(LogInworldCharacter, Error,
						TEXT("UInworldConversationGroup::CreateConversation - Failed to create "
							 "Conversation. Failed to get Speaker Selection Graph."));
					NativeCallback.ExecuteIfBound(nullptr, false);
				}
			}));
	}
	else
	{
		NativeCallback.ExecuteIfBound(ConversationGroup, true);
	}
}

UInworldBasePlayerComponent* UInworldConversationGroup::GetPlayerTarget()
{
	return PlayerTarget.Get();
}

TArray<UInworldCharacterComponent*> UInworldConversationGroup::GetParticipants()
{
	TArray<TObjectPtr<UInworldCharacterComponent>> OutArray;
	ConversationParticipants.GenerateValueArray(OutArray);
	return OutArray;
}

void UInworldConversationGroup::AddCharacter(UInworldCharacterComponent* CharacterComponent)
{
	if (CharacterComponent == nullptr)
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::AddCharacter - failed to add character. "
				 "Character is null."));
		return;
	}

	if (!CharacterComponent->IsReady())
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::AddCharacter - failed to add character. "
				 "Character is not ready."));
		return;
	}

	if (CharacterComponent->ConversationGroup != nullptr)
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::AddCharacter - failed to add character. "
				 "Character is already in another group."));
		return;
	}

	FString CharacterId = CharacterComponent->GetId();
	if (CharacterId.IsEmpty())
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::AddCharacter - failed to add character. "
				 "CharacterId is empty."));
		return;
	}

	if (ConversationParticipants.Contains(CharacterId))
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::AddCharacter - failed to add character. "
				 "A character with the id `%s` already exists."),
			*CharacterId);
		return;
	}

	Execute_Interrupt(CharacterComponent);
	CharacterComponent->ConversationGroup = this;

	ConversationParticipants.Add(CharacterId, CharacterComponent);

	OnCharacterAdded.Broadcast(this, CharacterComponent);
	OnCharacterAddedNative.Broadcast(this, CharacterComponent);
}

void UInworldConversationGroup::RemoveCharacter(FString CharacterId)
{
	if (!ConversationParticipants.Contains(CharacterId))
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::RemoveCharacter - failed to remove character. "
				 "No character with id `%s` was found."),
			*CharacterId);
		return;
	}

	if (CurrentMessage.IsValid() && CurrentMessage.ResponderId == CharacterId)
	{
		CancelCurrentMessage();
		Execute_Interrupt(ConversationParticipants[CharacterId]);
	}

	ConversationParticipants[CharacterId]->ConversationGroup = nullptr;

	ConversationParticipants.Remove(CharacterId);

	OnCharacterRemoved.Broadcast(this, ConversationParticipants[CharacterId]);
	OnCharacterRemovedNative.Broadcast(this, ConversationParticipants[CharacterId]);
}

void UInworldConversationGroup::SetPlayerTarget(UInworldBasePlayerComponent* PlayerComponent)
{
	PlayerTarget = Cast<UInworldSimplePlayerComponent>(PlayerComponent);
	if (PlayerTarget == nullptr)
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::SetPlayerTarget - failed to set player target. "
				 "Player is not derived from type UInworldSimplePlayerComponent."));
	}
}

void UInworldConversationGroup::RemovePlayerTarget()
{
	PlayerTarget = nullptr;
}

void UInworldConversationGroup::Interrupt_Implementation()
{
	CancelCurrentMessage();
	for (TTuple<FString, TObjectPtr<UInworldBaseCharacterComponent>> ConversationParticipant : ConversationParticipants)
	{
		Execute_Interrupt(ConversationParticipant.Value);
	}
}

FString UInworldConversationGroup::SendMessage_Implementation(const FInworldDataHandle& Message)
{
	LastSpeaker = "Player";
	if (PlayerTarget.IsValid())
	{
		return Execute_SendDataMessage(
			this, Message, {{PlayerProfileRuntimeDataKey, FInworldData_Struct(PlayerTarget->PlayerProfile)}});
	}
	return Execute_SendDataMessage(this, Message, {{PlayerProfileRuntimeDataKey, FInworldData_Struct(FInworldPlayerProfile())}});
}

FString UInworldConversationGroup::SendPlayerMessage_Implementation(
	const FInworldDataHandle& Message, const FInworldPlayerProfile& PlayerProfile)
{
	LastSpeaker = PlayerProfile.Name;

	FString MessageId = Execute_SendDataMessage(this, Message, {{PlayerProfileRuntimeDataKey, FInworldData_Struct(PlayerProfile)}});

	OnPlayerMessage.Broadcast(this, PlayerProfile, MessageId, Message);
	OnPlayerMessageNative.Broadcast(this, PlayerProfile, MessageId, Message);

	return MessageId;
}

FString UInworldConversationGroup::SendDataMessage_Implementation(
	const FInworldDataHandle& Message, const TMap<FName, FInworldDataHandle>& RuntimeData)
{
	if (ConversationParticipants.Num() == 0)
	{
		UE_LOG(LogInworldCharacter, Warning,
			TEXT("UInworldConversationGroup::SendMessage - failed to send message. "
				 "No conversation participants."));
		return "";
	}

	if (CurrentMessage.IsValid())
	{
		Execute_Interrupt(this);
	}

	return SendMessageInternal(Message, RuntimeData);
}

void UInworldConversationGroup::InvokeNextResponse(const FString& CharacterId)
{
	if (!CharacterId.IsEmpty() && !ConversationParticipants.Contains(CharacterId))
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::AddCharacter - failed to invoke next response. "
				 "No character with id `%s` was found."),
			*CharacterId);
		return;
	}

	FInworldCharacterAction NextResponseAction;
	NextResponseAction.GoalName = "InvokeNextResponse";
	NextResponseAction.ResponseType = EInworldGoalResponseType::INSTRUCTION;
	NextResponseAction.Action = NextResponseInstruction;

	if (PlayerTarget.IsValid())
	{
		SendMessageInternal(FInworldData_Struct(FInworldCharacterAction(NextResponseAction)),
			{{PlayerProfileRuntimeDataKey, FInworldData_Struct(PlayerTarget->PlayerProfile)}}, CharacterId);
	}
	else
	{
		SendMessageInternal(FInworldData_Struct(FInworldCharacterAction(NextResponseAction)),
			{{PlayerProfileRuntimeDataKey, FInworldData_Struct(FInworldPlayerProfile())}}, CharacterId);
	}
}

FInworldConversationState UInworldConversationGroup::GetConversationState()
{
	TArray<FInworldConversationCharacter> Participants;
	for (auto Participant : GetParticipants())
	{
		Participants.Add({Participant->GetId(), Participant->CharacterProfile->Get().Name});
	}

	return FInworldConversationState{Participants, LastSpeaker, PlayerTarget != nullptr, true};
}

void UInworldConversationGroup::SetRuleDirector(USpeakerRuleDirector* InRuleDirector)
{
	RuleDirectorPtr = InRuleDirector;
}

void UInworldConversationGroup::BeginDestroy()
{
	Super::BeginDestroy();
}

FString UInworldConversationGroup::SendMessageInternal(
	const FInworldDataHandle& MessageData, const TMap<FName, FInworldDataHandle>& RuntimeData, const FString& ResponderId)
{
	if (CurrentMessage.IsValid())
	{
		UE_LOG(LogInworldCharacter, Warning,
			TEXT("UInworldConversationGroup::SendMessageInternal - failed to send "
				 "message. A previous message is still sending."));
		return "";
	}

	CurrentMessage.MessageId = FGuid::NewGuid().ToString();
	CurrentMessage.DataHandle = MessageData;
	CurrentMessage.RuntimeData = RuntimeData;

	if (!ResponderId.IsEmpty())
	{
		CurrentMessage.ResponderId = ResponderId;
		UInworldCharacterComponent* CharacterToRespond = ConversationParticipants[ResponderId];
		OnCharacterGraphResultHandle =
			CharacterToRespond->OnDialogueResponseNative.AddUObject(this, &UInworldConversationGroup::OnCharacterDialogueResponse);
		CurrentMessage.ResponderMessageId = Execute_SendDataMessage(CharacterToRespond, MessageData, RuntimeData);
	}
	else
	{
		GetCharacterToRespond(
			MessageData, FOnInworldCharacterToRespond::CreateLambda(
							 [WeakThis = TWeakObjectPtr<UInworldConversationGroup>(this), Message = &CurrentMessage](
								 UInworldCharacterComponent* CharacterToRespond, bool bSuccess)
							 {
								 if (!WeakThis.IsValid() || !Message->IsValid())
								 {
									 return;
								 }
								 if (!bSuccess || CharacterToRespond == nullptr)
								 {
									 WeakThis->CancelCurrentMessage();
									 UE_LOG(LogInworldCharacter, Warning,
										 TEXT("UInworldConversationGroup::BroadcastMessage - "
											  "broadcast failed. Failed to get character from Speaker Selection graph."))
									 return;
								 }
								 Message->ResponderId = CharacterToRespond->Id;
								 WeakThis->OnCharacterGraphResultHandle = CharacterToRespond->OnDialogueResponseNative.AddUObject(
									 WeakThis.Get(), &UInworldConversationGroup::OnCharacterDialogueResponse);
								 Message->ResponderMessageId =
									 Execute_SendDataMessage(CharacterToRespond, Message->DataHandle, Message->RuntimeData);
							 }));
	}
	return CurrentMessage.MessageId;
}

void UInworldConversationGroup::UpdateHistory(
	UInworldSimpleCharacterComponent* RespondingCharacter, const FInworldCharacterResponse& Response)
{
	if (!CurrentMessage.IsValid())
	{
		return;
	}

	RespondingCharacter->OnDialogueResponseNative.Remove(OnCharacterGraphResultHandle);

	if (Response.UnSafe)
	{
		CurrentMessage.Reset();
		return;
	}

	if (EventHistory != nullptr)
	{
		if (!Response.Query.IsEmpty())
		{
			if (FInworldPlayerProfile PlayerProfile;
				CurrentMessage.GetPlayerProfileRuntimeData(PlayerProfileRuntimeDataKey, PlayerProfile))
			{
				EventHistory->AddSpeechEvent({PlayerProfile.Name, Response.Query});
			}
			else
			{
				EventHistory->AddSpeechEvent({"Player", Response.Query});
			}
		}
		if (!Response.Result.IsEmpty())
		{
			EventHistory->AddSpeechEvent({RespondingCharacter->CharacterProfile->Get().Name, Response.Result});
		}

		if (IsValid(Memory))
		{
			Memory->HandleEventHistoryAddition(EventHistory);
		}
	}
	else
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::UpdateHistory - failed to update event "
				 "history. Missing Event History runtime data."))
	}

	LastSpeaker = RespondingCharacter->CharacterProfile->Get().Name;
	for (TTuple<FString, TObjectPtr<UInworldCharacterComponent>> ConversationParticipant : ConversationParticipants)
	{
		UInworldCharacterComponent* Character = ConversationParticipant.Value;
		if (Character->GetId() != RespondingCharacter->GetId())
		{
			if (!Response.Query.IsEmpty())
			{
				if (FInworldPlayerProfile PlayerProfile;
					CurrentMessage.GetPlayerProfileRuntimeData(PlayerProfileRuntimeDataKey, PlayerProfile))
				{
					Character->AddHistoryEvent({PlayerProfile.Name, Response.Query});
				}
				else
				{
					Character->AddHistoryEvent({"Player", Response.Query});
				}
			}
			Character->AddHistoryEvent({LastSpeaker, Response.Result});
		}
	}
	CurrentMessage.Reset();
}

void UInworldConversationGroup::OnCharacterDialogueResponse(
	UInworldSimpleCharacterComponent* CharacterComponent, const FString& MessageId, const FInworldCharacterResponse& Response)
{
	if (!CurrentMessage.IsValid() || CurrentMessage.ResponderMessageId != MessageId)
	{
		return;
	}

	if (!Response.UnSafe)
	{
		OnCharacterMessage.Broadcast(this, CharacterComponent, CurrentMessage.MessageId, Response.Result);
		OnCharacterMessageNative.Broadcast(this, CharacterComponent, CurrentMessage.MessageId, Response.Result);
	}
	UpdateHistory(CharacterComponent, Response);
}

void UInworldConversationGroup::GetCharacterToRespond(const FInworldDataHandle& Data, FOnInworldCharacterToRespond Callback)
{
	if (RuleDirectorPtr.IsValid())
	{
		TArray<FString> ParticipantIDs;
		ConversationParticipants.GetKeys(ParticipantIDs);

		TOptional<FString> NextSpeakerID = RuleDirectorPtr->GetNextSpeaker(LastSpeaker, ParticipantIDs);

		if (NextSpeakerID.IsSet())
		{
			const FString& SpeakerID = NextSpeakerID.GetValue();
			if (TObjectPtr<UInworldCharacterComponent>* FoundCharacter = ConversationParticipants.Find(SpeakerID))
			{
				UE_LOG(LogInworldCharacter, Log, TEXT("SpeakerDirector: Overriding AI selection. Rule-based next speaker is: %s"),
					*SpeakerID);

				Callback.ExecuteIfBound(*FoundCharacter, true);
				return;
			}
			else
			{
				UE_LOG(LogInworldCharacter, Warning,
					TEXT(
						"SpeakerDirector: Rule specified speaker '%s' but they are not in the ConversationParticipants map. Falling back to AI."),
					*SpeakerID);
			}
		}
	}

	if (ConversationParticipants.Num() == 0)
	{
		Callback.ExecuteIfBound(nullptr, false);
		return;
	}

	if (SpeakerSelectionGraph == nullptr)
	{
		TArray<FString> Keys;
		ConversationParticipants.GetKeys(Keys);
		Callback.ExecuteIfBound(ConversationParticipants[Keys[FMath::RandRange(0, ConversationParticipants.Num() - 1)]], true);
		return;
	}

	CurrentSpeakerSelectionGraphExecutionId = SpeakerSelectionGraph->Execute(this, Data,
		{EventHistory->GetDataPair(), {ConversationStateRuntimeDataKey, FInworldData_Struct(GetConversationState())}},
		FOnGraphResultNative::CreateLambda(
			[WeakThis = TWeakObjectPtr<UInworldConversationGroup>(this), Callback](
				const FString& ExecutionId, const FString& NodeId, const FInworldDataHandle& DataHandle)
			{
				if (!WeakThis.IsValid() || WeakThis->CurrentSpeakerSelectionGraphExecutionId != ExecutionId)
				{
					Callback.ExecuteIfBound(nullptr, false);
					return;
				}

				if (FInworldData_Text* TextData = DataHandle.Unwrap<FInworldData_Text>(); TextData != nullptr)
				{
					int32 Index = FCString::Atoi(*TextData->Text);
					if (Index-- <= 0 || Index > WeakThis->ConversationParticipants.Num() - 1)
					{
						UE_LOG(LogInworldCharacter, Warning,
							TEXT("UInworldConversationGroup::GetCharacterToRespond - "
								 "failed to get character from Speaker Selection graph. Returning random participant."))
						// Return random participant
						Index = FMath::RandRange(0, WeakThis->ConversationParticipants.Num() - 1);
					}
					TArray<FString> Keys;
					WeakThis->ConversationParticipants.GetKeys(Keys);
					Callback.ExecuteIfBound(WeakThis->ConversationParticipants[Keys[Index]], true);
				}
				else
				{
					Callback.ExecuteIfBound(nullptr, false);
				}
				WeakThis->CurrentSpeakerSelectionGraphExecutionId = "";
			}));
}

void UInworldConversationGroup::CancelCurrentMessage()
{
	if (CurrentMessage.IsValid())
	{
		if (SpeakerSelectionGraph && !CurrentSpeakerSelectionGraphExecutionId.IsEmpty())
		{
			SpeakerSelectionGraph->CancelExecution(CurrentSpeakerSelectionGraphExecutionId);
			CurrentSpeakerSelectionGraphExecutionId = "";
		}
		CurrentMessage.Reset();
	}
}

void UInworldConversationGroup::ClearHistory(bool IncludeCharacters) const
{
	if (IsValid(EventHistory))
	{
		EventHistory->Clear();
	}
	else
	{
		UE_LOG(LogInworldCharacter, Error,
			TEXT("UInworldConversationGroup::ClearHistory - Failed to clear "
				 "Event History. Event History runtime data is null."));
	}

	if (IncludeCharacters)
	{
		for (TTuple<FString, TObjectPtr<UInworldCharacterComponent>> ConversationParticipant : ConversationParticipants)
		{
			ConversationParticipant.Value->ClearHistory();
		}
	}
}
