#include "HyperTwistRecognition/HyperTwistSpeechLibrary.h"

#include "Components/AudioComponent.h"
#include "Engine/Engine.h"
#include "Engine/GameInstance.h"
#include "EngineUtils.h"
#include "HyperTwistBootstrap/HyperTwistContractLibrary.h"
#include "HyperTwistRecognition/HyperTwistVoiceClient.h"
#include "HyperTwistTraining/HyperTwistTrainingSubsystem.h"
#include "HyperTwistUX/HyperTwistPlayerSettings.h"
#include "TimerManager.h"

namespace HyperTwistSpeechLibraryInternal
{
	const FName SpeechAudioComponentTag(TEXT("HyperTwistSpeechAudio"));
	const FName MusicAudioComponentTag(TEXT("HyperTwistMusicAudio"));
	const FName EffectsAudioComponentTag(TEXT("HyperTwistEffectsAudio"));
	constexpr float DefaultDuckVolumeMultiplier = 0.15f;
	constexpr float RestoreDelaySeconds = 0.2f;

	struct FManagedAudioComponentEntry
	{
		TWeakObjectPtr<UAudioComponent> AudioComponent;
		EHyperTwistManagedAudioCategory Category =
			EHyperTwistManagedAudioCategory::Effects;
		float AuthoredVolumeMultiplier = 1.0f;
	};

	struct FManagedAudioRegistry
	{
		TArray<FManagedAudioComponentEntry> RegisteredComponents;
		TMap<TWeakObjectPtr<UAudioComponent>, float> OriginalVolumeMap;
		FTimerHandle RestoreTimerHandle;
		float MusicVolume = 1.0f;
		float EffectsVolume = 1.0f;
		float VoiceVolume = 1.0f;
		float ActiveDuckVolumeMultiplier = DefaultDuckVolumeMultiplier;
		bool bDuckingActive = false;
	};

	FManagedAudioRegistry GManagedAudioRegistry;

	UHyperTwistTrainingSubsystem* ResolveTrainingSubsystem(UObject* ContextObject)
	{
		if (UHyperTwistTrainingSubsystem* TrainingSubsystem = Cast<UHyperTwistTrainingSubsystem>(ContextObject))
		{
			return TrainingSubsystem;
		}

		if (ContextObject == nullptr)
		{
			return nullptr;
		}

		if (UWorld* World = ContextObject->GetWorld())
		{
			if (UGameInstance* GameInstance = World->GetGameInstance())
			{
				return GameInstance->GetSubsystem<UHyperTwistTrainingSubsystem>();
			}
		}

		return nullptr;
	}

	IHyperTwistVoiceClient* ResolveVoiceClient(UObject* VoiceClientObject)
	{
		return VoiceClientObject != nullptr ? Cast<IHyperTwistVoiceClient>(VoiceClientObject) : nullptr;
	}

	UObject* CreateVoiceClient(UObject* ContextObject, const bool bUseMockVoiceClient)
	{
		UObject* Outer = ContextObject != nullptr ? ContextObject : GetTransientPackage();
		if (bUseMockVoiceClient)
		{
			return NewObject<UHyperTwistMockVoiceClient>(Outer);
		}

		UHyperTwistHttpVoiceClient* Client =
			NewObject<UHyperTwistHttpVoiceClient>(Outer);
		Client->LoadConfig();
		const FHyperTwistPlayerPreferences Preferences =
			UHyperTwistPlayerSettingsLibrary::LoadPreferences();
		Client->ProviderLabel = Preferences.VoiceProviderId;
		Client->ServiceBaseUrl = Preferences.VoiceEndpoint;
		Client->AuthorizationToken.Reset();
		UHyperTwistPlayerSettingsLibrary::LoadProviderCredential(
			TEXT("voice-api-key"),
			Client->AuthorizationToken);
		return Client;
	}

	void CompactManagedAudioRegistry()
	{
		GManagedAudioRegistry.RegisteredComponents.RemoveAll(
			[](const FManagedAudioComponentEntry& Entry)
			{
				return !Entry.AudioComponent.IsValid();
			}
		);

		for (auto It = GManagedAudioRegistry.OriginalVolumeMap.CreateIterator(); It; ++It)
		{
			if (!It.Key().IsValid())
			{
				It.RemoveCurrent();
			}
		}
	}

	EHyperTwistManagedAudioCategory ResolveTaggedCategory(
		const UAudioComponent* AudioComponent
	)
	{
		if (AudioComponent == nullptr)
		{
			return EHyperTwistManagedAudioCategory::Effects;
		}

		if (AudioComponent->ComponentHasTag(SpeechAudioComponentTag))
		{
			return EHyperTwistManagedAudioCategory::Voice;
		}
		if (AudioComponent->ComponentHasTag(MusicAudioComponentTag))
		{
			return EHyperTwistManagedAudioCategory::Music;
		}
		if (AudioComponent->ComponentHasTag(EffectsAudioComponentTag))
		{
			return EHyperTwistManagedAudioCategory::Effects;
		}

		for (const FManagedAudioComponentEntry& Entry : GManagedAudioRegistry.RegisteredComponents)
		{
			if (Entry.AudioComponent.Get() == AudioComponent)
			{
				return Entry.Category;
			}
		}

		return EHyperTwistManagedAudioCategory::Effects;
	}

	float ResolveCategoryVolume(const EHyperTwistManagedAudioCategory Category)
	{
		switch (Category)
		{
		case EHyperTwistManagedAudioCategory::Music:
			return GManagedAudioRegistry.MusicVolume;
		case EHyperTwistManagedAudioCategory::Voice:
			return GManagedAudioRegistry.VoiceVolume;
		default:
			return GManagedAudioRegistry.EffectsVolume;
		}
	}

	void ApplyCategoryTag(
		UAudioComponent* AudioComponent,
		const EHyperTwistManagedAudioCategory Category
	)
	{
		if (AudioComponent == nullptr)
		{
			return;
		}

		AudioComponent->ComponentTags.Remove(SpeechAudioComponentTag);
		AudioComponent->ComponentTags.Remove(MusicAudioComponentTag);
		AudioComponent->ComponentTags.Remove(EffectsAudioComponentTag);
		switch (Category)
		{
		case EHyperTwistManagedAudioCategory::Music:
			AudioComponent->ComponentTags.AddUnique(MusicAudioComponentTag);
			break;
		case EHyperTwistManagedAudioCategory::Voice:
			AudioComponent->ComponentTags.AddUnique(SpeechAudioComponentTag);
			break;
		default:
			AudioComponent->ComponentTags.AddUnique(EffectsAudioComponentTag);
			break;
		}
	}

	void ApplyCategoryVolume(FManagedAudioComponentEntry& Entry)
	{
		UAudioComponent* AudioComponent = Entry.AudioComponent.Get();
		if (AudioComponent == nullptr)
		{
			return;
		}

		const float UnduckedVolume =
			Entry.AuthoredVolumeMultiplier * ResolveCategoryVolume(Entry.Category);
		const TWeakObjectPtr<UAudioComponent> AudioComponentKey(AudioComponent);
		if (GManagedAudioRegistry.bDuckingActive
			&& Entry.Category != EHyperTwistManagedAudioCategory::Voice)
		{
			GManagedAudioRegistry.OriginalVolumeMap.Add(
				AudioComponentKey,
				UnduckedVolume);
			AudioComponent->SetVolumeMultiplier(
				UnduckedVolume
				* GManagedAudioRegistry.ActiveDuckVolumeMultiplier);
			return;
		}

		GManagedAudioRegistry.OriginalVolumeMap.Remove(AudioComponentKey);
		AudioComponent->SetVolumeMultiplier(UnduckedVolume);
	}

	TArray<UAudioComponent*> ResolveManagedAudioComponents()
	{
		CompactManagedAudioRegistry();

		TArray<UAudioComponent*> Components;
		for (const FManagedAudioComponentEntry& Entry : GManagedAudioRegistry.RegisteredComponents)
		{
			if (UAudioComponent* AudioComponent = Entry.AudioComponent.Get())
			{
				Components.AddUnique(AudioComponent);
			}
		}

		if (GEngine == nullptr)
		{
			return Components;
		}

		for (const FWorldContext& WorldContext : GEngine->GetWorldContexts())
		{
			UWorld* World = WorldContext.World();
			if (World == nullptr || World->IsPreviewWorld())
			{
				continue;
			}

			for (TActorIterator<AActor> ActorIt(World); ActorIt; ++ActorIt)
			{
				TInlineComponentArray<UAudioComponent*> AudioComponents(*ActorIt);
				for (UAudioComponent* AudioComponent : AudioComponents)
				{
					if (AudioComponent != nullptr)
					{
						if (!Components.Contains(AudioComponent))
						{
							FManagedAudioComponentEntry Entry;
							Entry.AudioComponent = AudioComponent;
							Entry.Category = ResolveTaggedCategory(AudioComponent);
							Entry.AuthoredVolumeMultiplier =
								AudioComponent->VolumeMultiplier;
							GManagedAudioRegistry.RegisteredComponents.Add(Entry);
						}
						Components.AddUnique(AudioComponent);
					}
				}
			}
		}

		return Components;
	}

	int32 ApplyManagedAudioDucking(const float DuckVolumeMultiplier)
	{
		CompactManagedAudioRegistry();

		const float TargetVolume = FMath::Clamp(DuckVolumeMultiplier, 0.0f, 1.0f);
		const TArray<UAudioComponent*> AudioComponents = ResolveManagedAudioComponents();
		int32 DuckedComponentCount = 0;

		for (UAudioComponent* AudioComponent : AudioComponents)
		{
			if (AudioComponent == nullptr
				|| ResolveTaggedCategory(AudioComponent)
					== EHyperTwistManagedAudioCategory::Voice)
			{
				continue;
			}

			const TWeakObjectPtr<UAudioComponent> AudioComponentKey(AudioComponent);
			float UnduckedVolume = AudioComponent->VolumeMultiplier;
			if (const float* ExistingUnduckedVolume =
					GManagedAudioRegistry.OriginalVolumeMap.Find(AudioComponentKey))
			{
				UnduckedVolume = *ExistingUnduckedVolume;
			}
			GManagedAudioRegistry.OriginalVolumeMap.Add(
				AudioComponentKey,
				UnduckedVolume);
			AudioComponent->SetVolumeMultiplier(
				UnduckedVolume * TargetVolume);
			++DuckedComponentCount;
		}

		GManagedAudioRegistry.ActiveDuckVolumeMultiplier = TargetVolume;
		GManagedAudioRegistry.bDuckingActive = DuckedComponentCount > 0;
		return DuckedComponentCount;
	}

	int32 RestoreManagedAudioDucking()
	{
		CompactManagedAudioRegistry();

		int32 RestoredComponentCount = 0;
		for (auto It = GManagedAudioRegistry.OriginalVolumeMap.CreateIterator(); It; ++It)
		{
			if (UAudioComponent* AudioComponent = It.Key().Get())
			{
				AudioComponent->SetVolumeMultiplier(It.Value());
				++RestoredComponentCount;
			}
			It.RemoveCurrent();
		}

		GManagedAudioRegistry.bDuckingActive = false;
		return RestoredComponentCount;
	}

	void ScheduleManagedAudioRestore(UObject* ContextObject)
	{
		if (ContextObject != nullptr)
		{
			if (UWorld* World = ContextObject->GetWorld())
			{
				World->GetTimerManager().ClearTimer(GManagedAudioRegistry.RestoreTimerHandle);
				World->GetTimerManager().SetTimer(
					GManagedAudioRegistry.RestoreTimerHandle,
					[]()
					{
						RestoreManagedAudioDucking();
					},
					RestoreDelaySeconds,
					false
				);
				return;
			}
		}

		RestoreManagedAudioDucking();
	}

	uint32 ReadLittleEndianUint32(const TArray<uint8>& Bytes, const int32 Offset)
	{
		if (!Bytes.IsValidIndex(Offset + 3))
		{
			return 0;
		}

		return static_cast<uint32>(Bytes[Offset])
			| (static_cast<uint32>(Bytes[Offset + 1]) << 8)
			| (static_cast<uint32>(Bytes[Offset + 2]) << 16)
			| (static_cast<uint32>(Bytes[Offset + 3]) << 24);
	}

	TArray<uint8> ExtractPcm16Payload(const TArray<uint8>& AudioBytes, const FString& AudioEncoding)
	{
		if (!AudioEncoding.Contains(TEXT("wav"), ESearchCase::IgnoreCase)
			|| AudioBytes.Num() < 12
			|| AudioBytes[0] != 'R'
			|| AudioBytes[1] != 'I'
			|| AudioBytes[2] != 'F'
			|| AudioBytes[3] != 'F'
			|| AudioBytes[8] != 'W'
			|| AudioBytes[9] != 'A'
			|| AudioBytes[10] != 'V'
			|| AudioBytes[11] != 'E')
		{
			return AudioBytes;
		}

		int32 ChunkOffset = 12;
		while (ChunkOffset + 8 <= AudioBytes.Num())
		{
			const uint32 ChunkSize = ReadLittleEndianUint32(AudioBytes, ChunkOffset + 4);
			const int32 ChunkDataOffset = ChunkOffset + 8;
			const int32 ChunkEndOffset = ChunkDataOffset + static_cast<int32>(ChunkSize);
			if (ChunkEndOffset > AudioBytes.Num())
			{
				break;
			}

			if (AudioBytes[ChunkOffset] == 'd'
				&& AudioBytes[ChunkOffset + 1] == 'a'
				&& AudioBytes[ChunkOffset + 2] == 't'
				&& AudioBytes[ChunkOffset + 3] == 'a')
			{
				TArray<uint8> Payload;
				Payload.Append(AudioBytes.GetData() + ChunkDataOffset, static_cast<int32>(ChunkSize));
				return Payload;
			}

			ChunkOffset = ChunkEndOffset + (ChunkSize % 2 == 0 ? 0 : 1);
		}

		return AudioBytes;
	}
}

bool UHyperTwistSpeechLibrary::StartDictationSession(
	UObject* WorldContextObject,
	FString& OutError,
	const bool bDuckAudio,
	const float DuckVolumeMultiplier
)
{
	if (UHyperTwistTrainingSubsystem* TrainingSubsystem =
			HyperTwistSpeechLibraryInternal::ResolveTrainingSubsystem(WorldContextObject))
	{
		if (TrainingSubsystem->OpenActiveCompanionSpeechSession(OutError))
		{
			if (bDuckAudio)
			{
				HyperTwistSpeechLibraryInternal::ApplyManagedAudioDucking(DuckVolumeMultiplier);
			}
			return true;
		}
	}
	else
	{
		OutError = TEXT("training-subsystem-unavailable");
	}

	return false;
}

FHyperTwistSpeechTranscriptResult UHyperTwistSpeechLibrary::SubmitDictationUtterance(
	UObject* WorldContextObject,
	const FHyperTwistSpeechUtteranceEnvelope& Utterance
)
{
	if (UHyperTwistTrainingSubsystem* TrainingSubsystem =
			HyperTwistSpeechLibraryInternal::ResolveTrainingSubsystem(WorldContextObject))
	{
		const FHyperTwistTrainingCompanionSpeechSessionState SessionState =
			TrainingSubsystem->GetActiveCompanionSpeechSessionState();
		if (!SessionState.bSessionOpen)
		{
			FString OpenError;
			StartDictationSession(WorldContextObject, OpenError, true, HyperTwistSpeechLibraryInternal::DefaultDuckVolumeMultiplier);
		}

		return TrainingSubsystem->SubmitActiveCompanionSpeechUtterance(Utterance);
	}

	FHyperTwistSpeechTranscriptResult Result;
	Result.Warnings.Add(TEXT("training-subsystem-unavailable"));
	return Result;
}

bool UHyperTwistSpeechLibrary::EndDictationSession(
	UObject* WorldContextObject,
	FString& OutError,
	const bool bRestoreAudio
)
{
	if (UHyperTwistTrainingSubsystem* TrainingSubsystem =
			HyperTwistSpeechLibraryInternal::ResolveTrainingSubsystem(WorldContextObject))
	{
		const bool bClosed = TrainingSubsystem->CloseActiveCompanionSpeechSession(OutError);
		if (bClosed && bRestoreAudio)
		{
			HyperTwistSpeechLibraryInternal::ScheduleManagedAudioRestore(WorldContextObject);
		}
		return bClosed;
	}

	OutError = TEXT("training-subsystem-unavailable");
	return false;
}

FHyperTwistTrainingCompanionSpeechSessionState UHyperTwistSpeechLibrary::GetActiveDictationSessionState(
	UObject* WorldContextObject
)
{
	if (UHyperTwistTrainingSubsystem* TrainingSubsystem =
			HyperTwistSpeechLibraryInternal::ResolveTrainingSubsystem(WorldContextObject))
	{
		return TrainingSubsystem->GetActiveCompanionSpeechSessionState();
	}

	return FHyperTwistTrainingCompanionSpeechSessionState();
}

void UHyperTwistSpeechLibrary::RegisterManagedAudioComponent(
	UAudioComponent* AudioComponent,
	const bool bTreatAsSpeechAudio
)
{
	RegisterCategorizedAudioComponent(
		AudioComponent,
		bTreatAsSpeechAudio
			? EHyperTwistManagedAudioCategory::Voice
			: EHyperTwistManagedAudioCategory::Effects);
}

void UHyperTwistSpeechLibrary::RegisterCategorizedAudioComponent(
	UAudioComponent* AudioComponent,
	const EHyperTwistManagedAudioCategory Category
)
{
	if (AudioComponent == nullptr)
	{
		return;
	}

	HyperTwistSpeechLibraryInternal::CompactManagedAudioRegistry();
	for (HyperTwistSpeechLibraryInternal::FManagedAudioComponentEntry& Entry :
		HyperTwistSpeechLibraryInternal::GManagedAudioRegistry.RegisteredComponents)
	{
		if (Entry.AudioComponent.Get() == AudioComponent)
		{
			Entry.Category = Category;
			HyperTwistSpeechLibraryInternal::ApplyCategoryTag(
				AudioComponent,
				Category);
			HyperTwistSpeechLibraryInternal::ApplyCategoryVolume(Entry);
			return;
		}
	}

	HyperTwistSpeechLibraryInternal::FManagedAudioComponentEntry Entry;
	Entry.AudioComponent = AudioComponent;
	Entry.Category = Category;
	Entry.AuthoredVolumeMultiplier = AudioComponent->VolumeMultiplier;
	HyperTwistSpeechLibraryInternal::GManagedAudioRegistry.RegisteredComponents.Add(Entry);
	HyperTwistSpeechLibraryInternal::ApplyCategoryTag(AudioComponent, Category);
	HyperTwistSpeechLibraryInternal::ApplyCategoryVolume(
		HyperTwistSpeechLibraryInternal::GManagedAudioRegistry.RegisteredComponents.Last());
}

void UHyperTwistSpeechLibrary::UnregisterManagedAudioComponent(UAudioComponent* AudioComponent)
{
	if (AudioComponent == nullptr)
	{
		return;
	}

	HyperTwistSpeechLibraryInternal::GManagedAudioRegistry.RegisteredComponents.RemoveAll(
		[AudioComponent](const HyperTwistSpeechLibraryInternal::FManagedAudioComponentEntry& Entry)
		{
			return Entry.AudioComponent.Get() == AudioComponent;
		}
	);
	HyperTwistSpeechLibraryInternal::GManagedAudioRegistry.OriginalVolumeMap.Remove(
		TWeakObjectPtr<UAudioComponent>(AudioComponent));
}

int32 UHyperTwistSpeechLibrary::ApplyManagedAudioCategoryVolumes(
	const float MusicVolume,
	const float EffectsVolume,
	const float VoiceVolume
)
{
	using namespace HyperTwistSpeechLibraryInternal;
	GManagedAudioRegistry.MusicVolume = FMath::Clamp(MusicVolume, 0.0f, 1.0f);
	GManagedAudioRegistry.EffectsVolume = FMath::Clamp(EffectsVolume, 0.0f, 1.0f);
	GManagedAudioRegistry.VoiceVolume = FMath::Clamp(VoiceVolume, 0.0f, 1.0f);
	ResolveManagedAudioComponents();

	int32 AppliedComponentCount = 0;
	for (FManagedAudioComponentEntry& Entry : GManagedAudioRegistry.RegisteredComponents)
	{
		if (Entry.AudioComponent.IsValid())
		{
			ApplyCategoryVolume(Entry);
			++AppliedComponentCount;
		}
	}
	return AppliedComponentCount;
}

int32 UHyperTwistSpeechLibrary::DuckManagedAudioComponents(const float DuckVolumeMultiplier)
{
	return HyperTwistSpeechLibraryInternal::ApplyManagedAudioDucking(DuckVolumeMultiplier);
}

int32 UHyperTwistSpeechLibrary::RestoreManagedAudioComponents()
{
	return HyperTwistSpeechLibraryInternal::RestoreManagedAudioDucking();
}

int32 UHyperTwistSpeechLibrary::GetManagedAudioComponentCount()
{
	HyperTwistSpeechLibraryInternal::CompactManagedAudioRegistry();
	return HyperTwistSpeechLibraryInternal::GManagedAudioRegistry.RegisteredComponents.Num();
}

TArray<FHyperTwistVoiceProfileSummary> UHyperTwistSpeechLibrary::ListVoiceProfiles(
	UObject* ContextObject,
	const bool bUseMockVoiceClient
)
{
	TStrongObjectPtr<UObject> VoiceClientObject(
		HyperTwistSpeechLibraryInternal::CreateVoiceClient(ContextObject, bUseMockVoiceClient)
	);
	if (IHyperTwistVoiceClient* VoiceClient =
			HyperTwistSpeechLibraryInternal::ResolveVoiceClient(VoiceClientObject.Get()))
	{
		return VoiceClient->ListVoiceProfiles();
	}

	return TArray<FHyperTwistVoiceProfileSummary>();
}

FHyperTwistNarrationSynthesisResult UHyperTwistSpeechLibrary::SynthesizeNarration(
	UObject* ContextObject,
	const FHyperTwistNarrationSynthesisRequest& Request,
	const bool bUseMockVoiceClient
)
{
	TStrongObjectPtr<UObject> VoiceClientObject(
		HyperTwistSpeechLibraryInternal::CreateVoiceClient(ContextObject, bUseMockVoiceClient)
	);
	if (IHyperTwistVoiceClient* VoiceClient =
			HyperTwistSpeechLibraryInternal::ResolveVoiceClient(VoiceClientObject.Get()))
	{
		return VoiceClient->SynthesizeNarration(Request);
	}

	FHyperTwistNarrationSynthesisResult Result;
	Result.RequestId = Request.RequestId;
	Result.Warnings.Add(TEXT("voice-client-unavailable"));
	return Result;
}

TArray<uint8> UHyperTwistSpeechLibrary::Synthesize(
	UObject* ContextObject,
	const FString& Text,
	const FString& VoiceId,
	const bool bUseMockVoiceClient
)
{
	FHyperTwistNarrationSynthesisRequest Request =
		UHyperTwistContractLibrary::MakeSampleNarrationSynthesisRequest();
	Request.RequestId = FGuid::NewGuid().ToString(EGuidFormats::DigitsWithHyphensLower);
	Request.ScriptText = Text;
	Request.SubtitleSeedText = Text;
	if (!VoiceId.IsEmpty())
	{
		Request.VoiceProfileId = VoiceId;
	}

	const FHyperTwistNarrationSynthesisResult Result =
		SynthesizeNarration(ContextObject, Request, bUseMockVoiceClient);
	return HyperTwistSpeechLibraryInternal::ExtractPcm16Payload(Result.AudioBytes, Result.AudioEncoding);
}

TArray<uint8> UHyperTwistSpeechLibrary::ExtractPcm16PayloadFromAudioBytes(
	const TArray<uint8>& AudioBytes,
	const FString& AudioEncoding
)
{
	return HyperTwistSpeechLibraryInternal::ExtractPcm16Payload(AudioBytes, AudioEncoding);
}

FHyperTwistVoiceServiceHealth UHyperTwistSpeechLibrary::GetVoiceServiceHealth(
	UObject* ContextObject,
	const bool bUseMockVoiceClient
)
{
	TStrongObjectPtr<UObject> VoiceClientObject(
		HyperTwistSpeechLibraryInternal::CreateVoiceClient(ContextObject, bUseMockVoiceClient)
	);
	if (IHyperTwistVoiceClient* VoiceClient =
			HyperTwistSpeechLibraryInternal::ResolveVoiceClient(VoiceClientObject.Get()))
	{
		return VoiceClient->GetVoiceServiceHealth();
	}

	return FHyperTwistVoiceServiceHealth();
}
