#include "HyperTwistUX/HyperTwistDictationCaptureComponent.h"

#include "HAL/PlatformFileManager.h"
#include "Interfaces/VoiceCapture.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "VoiceModule.h"

#include "HyperTwistRecognition/HyperTwistSpeechLibrary.h"
#include "HyperTwistUX/HyperTwistPlayerSettings.h"

namespace HyperTwistDictationCaptureComponentInternal
{
	constexpr int32 BytesPerPcm16Sample = 2;

	void AppendLittleEndian16(TArray<uint8>& Bytes, const uint16 Value)
	{
		Bytes.Add(static_cast<uint8>(Value & 0xff));
		Bytes.Add(static_cast<uint8>((Value >> 8) & 0xff));
	}

	void AppendLittleEndian32(TArray<uint8>& Bytes, const uint32 Value)
	{
		Bytes.Add(static_cast<uint8>(Value & 0xff));
		Bytes.Add(static_cast<uint8>((Value >> 8) & 0xff));
		Bytes.Add(static_cast<uint8>((Value >> 16) & 0xff));
		Bytes.Add(static_cast<uint8>((Value >> 24) & 0xff));
	}

	void AppendWaveHeader(
		TArray<uint8>& Bytes,
		const int32 PcmByteCount,
		const int32 SampleRateHz,
		const int32 ChannelCount
	)
	{
		const uint16 BitsPerSample = 16;
		const uint16 BlockAlign =
			static_cast<uint16>(ChannelCount * (BitsPerSample / 8));
		const uint32 BytesPerSecond =
			static_cast<uint32>(SampleRateHz) * BlockAlign;

		const ANSICHAR Riff[] = "RIFF";
		Bytes.Append(reinterpret_cast<const uint8*>(Riff), 4);
		AppendLittleEndian32(Bytes, 36U + static_cast<uint32>(PcmByteCount));
		const ANSICHAR Wave[] = "WAVE";
		Bytes.Append(reinterpret_cast<const uint8*>(Wave), 4);
		const ANSICHAR Format[] = "fmt ";
		Bytes.Append(reinterpret_cast<const uint8*>(Format), 4);
		AppendLittleEndian32(Bytes, 16);
		AppendLittleEndian16(Bytes, 1);
		AppendLittleEndian16(Bytes, static_cast<uint16>(ChannelCount));
		AppendLittleEndian32(Bytes, static_cast<uint32>(SampleRateHz));
		AppendLittleEndian32(Bytes, BytesPerSecond);
		AppendLittleEndian16(Bytes, BlockAlign);
		AppendLittleEndian16(Bytes, BitsPerSample);
		const ANSICHAR Data[] = "data";
		Bytes.Append(reinterpret_cast<const uint8*>(Data), 4);
		AppendLittleEndian32(Bytes, static_cast<uint32>(PcmByteCount));
	}

	bool IsCloudSpeechProvider(const FString& ProviderId)
	{
		return ProviderId == TEXT("openai")
			|| ProviderId == TEXT("groq")
			|| ProviderId == TEXT("deepgram")
			|| ProviderId == TEXT("elevenlabs");
	}
}

UHyperTwistDictationCaptureComponent::UHyperTwistDictationCaptureComponent()
{
	PrimaryComponentTick.bCanEverTick = true;
	PrimaryComponentTick.bStartWithTickEnabled = false;
}

void UHyperTwistDictationCaptureComponent::EndPlay(
	const EEndPlayReason::Type EndPlayReason
)
{
	CancelCapture();
	Super::EndPlay(EndPlayReason);
}

void UHyperTwistDictationCaptureComponent::TickComponent(
	const float DeltaTime,
	const ELevelTick TickType,
	FActorComponentTickFunction* ThisTickFunction
)
{
	Super::TickComponent(DeltaTime, TickType, ThisTickFunction);
	if (!bCaptureActive)
	{
		return;
	}

	FString FailureReason;
	if (!DrainCaptureBuffer(FailureReason))
	{
		ResetCaptureState(true);
		SetStatus(FailureReason, true);
		return;
	}

	if (FPlatformTime::Seconds() - CaptureStartedAtSeconds
		>= FMath::Max(static_cast<double>(MaximumCaptureSeconds), 1.0))
	{
		FString Transcript;
		EndCapture(Transcript, FailureReason);
		if (!FailureReason.IsEmpty())
		{
			SetStatus(FailureReason, true);
		}
	}
}

bool UHyperTwistDictationCaptureComponent::BeginCapture(FString& OutFailureReason)
{
	OutFailureReason.Reset();
	if (bCaptureActive)
	{
		return true;
	}

	const FHyperTwistPlayerPreferences Preferences =
		UHyperTwistPlayerSettingsLibrary::LoadPreferences();
	if (!Preferences.bSpeechInputEnabled)
	{
		OutFailureReason = TEXT("Enable speech-to-text in Settings > Voice & AI.");
		SetStatus(OutFailureReason, true);
		return false;
	}
	if (HyperTwistDictationCaptureComponentInternal::IsCloudSpeechProvider(
			Preferences.SpeechProviderId)
		&& !Preferences.bAllowCloudProviders)
	{
		OutFailureReason =
			TEXT("Enable cloud providers before using the selected transcription service.");
		SetStatus(OutFailureReason, true);
		return false;
	}
	if (!UHyperTwistPlayerSettingsLibrary::IsProviderEndpointAllowed(
			Preferences.SpeechEndpoint,
			Preferences.bAllowCloudProviders,
			OutFailureReason))
	{
		SetStatus(OutFailureReason, true);
		return false;
	}
	if (HyperTwistDictationCaptureComponentInternal::IsCloudSpeechProvider(
			Preferences.SpeechProviderId)
		&& !UHyperTwistPlayerSettingsLibrary::HasProviderCredential(
			TEXT("speech-api-key")))
	{
		OutFailureReason =
			TEXT("Save a protected transcription API key before starting dictation.");
		SetStatus(OutFailureReason, true);
		return false;
	}

	if (!UHyperTwistSpeechLibrary::StartDictationSession(
			this,
			OutFailureReason,
			true,
			AudioDuckMultiplier))
	{
		SetStatus(
			OutFailureReason.IsEmpty()
				? TEXT("The transcription service is not ready.")
				: OutFailureReason,
			true);
		return false;
	}

	if (!FVoiceModule::Get().DoesPlatformSupportVoiceCapture())
	{
		OutFailureReason = TEXT("Microphone capture is unsupported on this platform.");
		ResetCaptureState(true);
		SetStatus(OutFailureReason, true);
		return false;
	}

	const FString DeviceName = Preferences.SpeechMicrophoneId.Equals(
		TEXT("system-default"),
		ESearchCase::IgnoreCase)
		? FString()
		: Preferences.SpeechMicrophoneId;
	ActiveVoiceCapture = FVoiceModule::Get().CreateVoiceCapture(
		DeviceName,
		CaptureSampleRateHz,
		CaptureChannelCount);
	if (!ActiveVoiceCapture.IsValid() || !ActiveVoiceCapture->Start())
	{
		OutFailureReason =
			TEXT("HyperTwist could not open the selected microphone. Check Windows privacy settings or choose another device.");
		ResetCaptureState(true);
		SetStatus(OutFailureReason, true);
		return false;
	}

	CapturedPcmBytes.Reset();
	ActiveUtteranceId = FGuid::NewGuid().ToString(EGuidFormats::Digits);
	ActiveCapturePath.Reset();
	CaptureStartedAtSeconds = FPlatformTime::Seconds();
	bCaptureActive = true;
	SetComponentTickEnabled(true);
	SetStatus(TEXT("Listening... release the shortcut or select Stop dictation."));
	return true;
}

bool UHyperTwistDictationCaptureComponent::EndCapture(
	FString& OutTranscript,
	FString& OutFailureReason
)
{
	OutTranscript.Reset();
	OutFailureReason.Reset();
	if (!bCaptureActive)
	{
		OutFailureReason = TEXT("Dictation is not currently recording.");
		return false;
	}

	DrainCaptureBuffer(OutFailureReason);
	if (ActiveVoiceCapture.IsValid())
	{
		ActiveVoiceCapture->Stop();
	}
	DrainCaptureBuffer(OutFailureReason);

	const FHyperTwistTrainingCompanionSpeechSessionState SessionState =
		UHyperTwistSpeechLibrary::GetActiveDictationSessionState(this);
	if (CapturedPcmBytes.IsEmpty())
	{
		OutFailureReason =
			TEXT("No microphone audio was captured. Check the selected device and Windows microphone permission.");
		ResetCaptureState(true);
		SetStatus(OutFailureReason, true);
		return false;
	}

	ActiveCapturePath = BuildTransientCapturePath();
	if (!WriteTransientWaveFile(ActiveCapturePath, CapturedPcmBytes))
	{
		OutFailureReason = TEXT("HyperTwist could not prepare the temporary dictation audio.");
		ResetCaptureState(true);
		SetStatus(OutFailureReason, true);
		return false;
	}

	FHyperTwistSpeechUtteranceEnvelope Utterance;
	Utterance.SessionId = SessionState.ActiveSessionId;
	Utterance.UtteranceId = ActiveUtteranceId;
	Utterance.CapturedAtUtc = FDateTime::UtcNow().ToIso8601();
	Utterance.AudioRef = ActiveCapturePath;
	Utterance.SourceTimestampMs = FMath::RoundToInt(
		static_cast<float>(
			(FPlatformTime::Seconds() - CaptureStartedAtSeconds) * 1000.0));
	Utterance.SampleCount =
		CapturedPcmBytes.Num()
		/ HyperTwistDictationCaptureComponentInternal::BytesPerPcm16Sample;
	Utterance.SampleRateHz = CaptureSampleRateHz;
	Utterance.ChannelCount = CaptureChannelCount;
	Utterance.SpeechStartMs = 0;
	Utterance.SpeechEndMs = FMath::Max(Utterance.SourceTimestampMs, 1);
	Utterance.SilenceGapMs = 0;
	Utterance.EnergyThreshold = 0.0f;

	const FHyperTwistSpeechTranscriptResult Result =
		UHyperTwistSpeechLibrary::SubmitDictationUtterance(this, Utterance);
	FString CloseError;
	UHyperTwistSpeechLibrary::EndDictationSession(this, CloseError, true);
	OutTranscript = Result.TranscriptText.TrimStartAndEnd();

	ResetCaptureState(false);
	if (OutTranscript.IsEmpty())
	{
		OutFailureReason = Result.Warnings.IsEmpty()
			? TEXT("The transcription service returned no text.")
			: FString::Printf(
				TEXT("Transcription failed: %s"),
				*FString::Join(Result.Warnings, TEXT(", ")));
		SetStatus(OutFailureReason, true);
		return false;
	}
	if (!CloseError.IsEmpty())
	{
		SetStatus(
			FString::Printf(
				TEXT("Dictation complete; audio cleanup reported: %s"),
				*CloseError),
			true);
	}
	else
	{
		SetStatus(TEXT("Transcript ready."));
	}
	OnTranscriptReady.Broadcast(OutTranscript);
	return true;
}

void UHyperTwistDictationCaptureComponent::CancelCapture()
{
	if (!bCaptureActive && !ActiveVoiceCapture.IsValid())
	{
		return;
	}
	ResetCaptureState(true);
	SetStatus(TEXT("Dictation cancelled."));
}

bool UHyperTwistDictationCaptureComponent::IsCapturing() const
{
	return bCaptureActive;
}

FString UHyperTwistDictationCaptureComponent::GetStatus() const
{
	return CurrentStatus;
}

bool UHyperTwistDictationCaptureComponent::DrainCaptureBuffer(
	FString& OutFailureReason
)
{
	OutFailureReason.Reset();
	if (!ActiveVoiceCapture.IsValid())
	{
		OutFailureReason = TEXT("The microphone capture route is unavailable.");
		return false;
	}

	uint32 AvailableBytes = 0;
	const EVoiceCaptureState::Type CaptureState =
		ActiveVoiceCapture->GetCaptureState(AvailableBytes);
	if (CaptureState == EVoiceCaptureState::NoData
		|| CaptureState == EVoiceCaptureState::Stopping
		|| CaptureState == EVoiceCaptureState::NotCapturing)
	{
		return true;
	}
	if (CaptureState != EVoiceCaptureState::Ok
		&& CaptureState != EVoiceCaptureState::BufferTooSmall)
	{
		OutFailureReason = FString::Printf(
			TEXT("Microphone capture failed (%s)."),
			EVoiceCaptureState::ToString(CaptureState));
		return false;
	}
	if (AvailableBytes == 0)
	{
		return true;
	}

	const int64 MaximumByteCount =
		static_cast<int64>(CaptureSampleRateHz)
		* FMath::Max(CaptureChannelCount, 1)
		* HyperTwistDictationCaptureComponentInternal::BytesPerPcm16Sample
		* FMath::Max(FMath::CeilToInt(MaximumCaptureSeconds), 1);
	if (static_cast<int64>(CapturedPcmBytes.Num()) + AvailableBytes
		> MaximumByteCount)
	{
		OutFailureReason = TEXT("Dictation reached its safe recording limit.");
		return false;
	}

	const int32 BufferOffset = CapturedPcmBytes.Num();
	CapturedPcmBytes.AddUninitialized(static_cast<int32>(AvailableBytes));
	uint32 ReadBytes = 0;
	const EVoiceCaptureState::Type ReadState = ActiveVoiceCapture->GetVoiceData(
		CapturedPcmBytes.GetData() + BufferOffset,
		AvailableBytes,
		ReadBytes);
	if (ReadState != EVoiceCaptureState::Ok
		&& ReadState != EVoiceCaptureState::NoData)
	{
		CapturedPcmBytes.SetNum(BufferOffset);
		OutFailureReason = FString::Printf(
			TEXT("Microphone read failed (%s)."),
			EVoiceCaptureState::ToString(ReadState));
		return false;
	}

	CapturedPcmBytes.SetNum(BufferOffset + static_cast<int32>(ReadBytes));
	return true;
}

bool UHyperTwistDictationCaptureComponent::WriteTransientWaveFile(
	const FString& FilePath,
	const TArray<uint8>& PcmBytes
) const
{
	IPlatformFile& PlatformFile = FPlatformFileManager::Get().GetPlatformFile();
	if (!PlatformFile.CreateDirectoryTree(*FPaths::GetPath(FilePath)))
	{
		return false;
	}

	TArray<uint8> WaveBytes;
	WaveBytes.Reserve(44 + PcmBytes.Num());
	HyperTwistDictationCaptureComponentInternal::AppendWaveHeader(
		WaveBytes,
		PcmBytes.Num(),
		CaptureSampleRateHz,
		CaptureChannelCount);
	WaveBytes.Append(PcmBytes);
	return FFileHelper::SaveArrayToFile(WaveBytes, *FilePath);
}

void UHyperTwistDictationCaptureComponent::ResetCaptureState(
	const bool bCloseDictationSession
)
{
	if (ActiveVoiceCapture.IsValid())
	{
		ActiveVoiceCapture->Stop();
		ActiveVoiceCapture.Reset();
	}
	if (bCloseDictationSession)
	{
		FString CloseError;
		UHyperTwistSpeechLibrary::EndDictationSession(this, CloseError, true);
	}
	if (!ActiveCapturePath.IsEmpty())
	{
		FPlatformFileManager::Get().GetPlatformFile().DeleteFile(*ActiveCapturePath);
	}

	SetComponentTickEnabled(false);
	CapturedPcmBytes.Reset();
	ActiveUtteranceId.Reset();
	ActiveCapturePath.Reset();
	CaptureStartedAtSeconds = 0.0;
	bCaptureActive = false;
}

void UHyperTwistDictationCaptureComponent::SetStatus(
	const FString& NewStatus,
	const bool bError
)
{
	CurrentStatus = NewStatus;
	OnStateChanged.Broadcast(CurrentStatus, bError);
}

FString UHyperTwistDictationCaptureComponent::BuildTransientCapturePath() const
{
	return FPaths::Combine(
		FPaths::ProjectSavedDir(),
		TEXT("VoiceCaptures"),
		TEXT("Transient"),
		FString::Printf(TEXT("%s.wav"), *ActiveUtteranceId));
}
