// Georgy Treshchev 2026 All Rights Reserved.

#include "RuntimeSpeechRecognizerVoskEditor.h"

#include "ISettingsModule.h"
#include "PropertyEditorModule.h"
#include "RuntimeVoskModelLibrary.h"
#include "RuntimeVoskModelSettingsDetails.h"
#include "RuntimeSpeechRecognizerVoskSettings.h"
#include "Settings/ProjectPackagingSettings.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "HAL/FileManager.h"
#include "Async/Async.h"
#include "Misc/EngineVersionComparison.h"

#if UE_VERSION_OLDER_THAN(5, 4, 0)
#include "VoskZipArchiveReader.h"
#else
#include "FileUtilities/ZipArchiveReader.h"
#endif

#define LOCTEXT_NAMESPACE "FRuntimeSpeechRecognizerVoskEditorModule"

DEFINE_LOG_CATEGORY(LogRuntimeSpeechRecognizerVoskEditor);

void FRuntimeSpeechRecognizerVoskEditorModule::StartupModule()
{
	FPropertyEditorModule& PropertyModule = FModuleManager::LoadModuleChecked<FPropertyEditorModule>("PropertyEditor");
	PropertyModule.RegisterCustomClassLayout(
		URuntimeSpeechRecognizerVoskSettings::StaticClass()->GetFName(),
		FOnGetDetailCustomizationInstance::CreateStatic(&FRuntimeSpeechRecognizerVoskSettingsDetails::MakeInstance)
	);
	PropertyModule.NotifyCustomizationModuleChanged();

	if (ISettingsModule* SettingsModule = FModuleManager::GetModulePtr<ISettingsModule>("Settings"))
	{
		SettingsModule->RegisterSettings(
			"Project", "Plugins", "Runtime Speech Recognizer Vosk",
			LOCTEXT("VoskEditorSettingsName", "Runtime Speech Recognizer Vosk"),
			LOCTEXT("VoskEditorSettingsDescription", "Browse, download, and manage Vosk speech recognition models"),
			GetMutableDefault<URuntimeSpeechRecognizerVoskSettings>());
	}

	IFileManager::Get().MakeDirectory(*URuntimeVoskModelLibrary::GetVoskModelsDirectory(), true);

	UpdatePackagingSettings();
}

void FRuntimeSpeechRecognizerVoskEditorModule::ShutdownModule()
{
	for (const auto& Pair : DownloadTasks)
	{
		if (Pair.Value.IsValid())
		{
			Pair.Value->CancelDownload();
		}
	}
	DownloadTasks.Empty();

	if (ISettingsModule* SettingsModule = FModuleManager::GetModulePtr<ISettingsModule>("Settings"))
	{
		SettingsModule->UnregisterSettings("Project", "Plugins", "Runtime Speech Recognizer Vosk");
	}

	if (FModuleManager::Get().IsModuleLoaded("PropertyEditor"))
	{
		FModuleManager::GetModuleChecked<FPropertyEditorModule>("PropertyEditor")
			.UnregisterCustomClassLayout(URuntimeSpeechRecognizerVoskSettings::StaticClass()->GetFName());
	}
}

bool FRuntimeSpeechRecognizerVoskEditorModule::UpdatePackagingSettings() const
{
	UProjectPackagingSettings* PackagingSettings = GetMutableDefault<UProjectPackagingSettings>();
	if (!PackagingSettings)
	{
		UE_LOG(LogRuntimeSpeechRecognizerVoskEditor, Error, TEXT("Cannot get the packaging settings"));
		return false;
	}

	const FString ModelsContentRelPath = URuntimeVoskModelLibrary::GetVoskModelsContentRelativePath();

	const bool bAlreadyInNonUFS = PackagingSettings->DirectoriesToAlwaysStageAsNonUFS.ContainsByPredicate(
		[&ModelsContentRelPath](const FDirectoryPath& DirPath)
		{
			return FPaths::IsSamePath(DirPath.Path, ModelsContentRelPath);
		});

	if (!bAlreadyInNonUFS)
	{
		FDirectoryPath DirPath;
		DirPath.Path = ModelsContentRelPath;
		PackagingSettings->DirectoriesToAlwaysStageAsNonUFS.Add(DirPath);
	}

#if UE_VERSION_OLDER_THAN(5, 0, 0)
	PackagingSettings->UpdateDefaultConfigFile();
#else
	PackagingSettings->TryUpdateDefaultConfigFile();
#endif

	return true;
}

void FRuntimeSpeechRecognizerVoskEditorModule::DownloadModel(const FVoskModelVariant& ModelVariant)
{
	const FString ModelId = ModelVariant.ModelId;

	if (URuntimeVoskModelLibrary::IsVoskModelOnDisk(ModelId))
	{
		OnModelOperationFailed.Broadcast(ModelId, LOCTEXT("VoskModelAlreadyPresent", "The Vosk model is already present on disk"));
		return;
	}
	if (ModelVariant.DownloadURL.IsEmpty())
	{
		OnModelOperationFailed.Broadcast(ModelId, LOCTEXT("VoskModelNoURL", "No download URL specified for this model"));
		return;
	}
	if (DownloadTasks.Contains(ModelId))
	{
		OnModelOperationFailed.Broadcast(ModelId, LOCTEXT("VoskModelDownloadInProgress", "A download for this model is already in progress"));
		return;
	}

	TSharedPtr<FRuntimeChunkDownloader_Recognizer> Downloader = MakeShared<FRuntimeChunkDownloader_Recognizer>();
	DownloadTasks.Add(ModelId, Downloader);

	Downloader->DownloadFile(ModelVariant.DownloadURL, 60.0f, TEXT("application/zip"), 20 * 1024 * 1024,
		[this, ModelId](int64 BytesReceived, int64 ContentSize)
		{
			const float Progress = ContentSize > 0 ? static_cast<float>(BytesReceived) / static_cast<float>(ContentSize) : 0.0f;
			AsyncTask(ENamedThreads::GameThread, [this, ModelId, Progress]()
			{
				OnModelDownloadProgress.Broadcast(ModelId, Progress * 0.9f);
			});
		}
	).Next([this, ModelId](FRuntimeChunkDownloaderResult_Recognizer Result)
	{
		AsyncTask(ENamedThreads::GameThread, [this, ModelId, Result = MoveTemp(Result)]() mutable
		{
			if (Result.Result == EDownloadToMemoryResult_Recognizer::Cancelled)
			{
				DownloadTasks.Remove(ModelId);
				OnModelDownloadCanceled.Broadcast(ModelId);
				return;
			}

			const bool bDownloadOk = (Result.Result == EDownloadToMemoryResult_Recognizer::Success || Result.Result == EDownloadToMemoryResult_Recognizer::SucceededByPayload)
				&& Result.Data.Num() > 0;

			if (!bDownloadOk)
			{
				DownloadTasks.Remove(ModelId);
				OnModelOperationFailed.Broadcast(ModelId, LOCTEXT("VoskModelDownloadFailed", "Failed to download the Vosk model"));
				return;
			}

			const FString TempZipPath = FPaths::Combine(FPaths::ProjectIntermediateDir(), TEXT("RuntimeSpeechRecognizerVosk"), ModelId + TEXT(".zip"));
			IFileManager::Get().MakeDirectory(*FPaths::GetPath(TempZipPath), true);

			bool bWroteZip = false;
			{
				TUniquePtr<FArchive> ZipWriter(IFileManager::Get().CreateFileWriter(*TempZipPath));
				if (ZipWriter)
				{
					ZipWriter->Serialize(Result.Data.GetData(), Result.Data.Num());
					bWroteZip = !ZipWriter->IsError();
					ZipWriter->Close();
				}
			}

			Result.Data.Empty();

			if (!bWroteZip)
			{
				DownloadTasks.Remove(ModelId);
				IFileManager::Get().Delete(*TempZipPath);
				OnModelOperationFailed.Broadcast(ModelId, LOCTEXT("VoskModelTempWriteFailed", "Downloaded the model but failed to write a temporary file for extraction"));
				return;
			}

			OnModelDownloadProgress.Broadcast(ModelId, 0.95f);

			FText ExtractError;
			const bool bExtracted = ExtractModelZip(TempZipPath, ModelId, ExtractError);

			IFileManager::Get().Delete(*TempZipPath);
			DownloadTasks.Remove(ModelId);

			if (!bExtracted)
			{
				URuntimeVoskModelLibrary::DeleteVoskModelFiles(ModelId);
				OnModelOperationFailed.Broadcast(ModelId, ExtractError);
				return;
			}

			const FString MarkerPath = FPaths::Combine(URuntimeVoskModelLibrary::GetVoskModelFolderPath(ModelId), URuntimeVoskModelLibrary::GetModelCompleteMarkerFileName());
			FFileHelper::SaveStringToFile(FDateTime::UtcNow().ToIso8601(), *MarkerPath);

			OnModelDownloadProgress.Broadcast(ModelId, 1.0f);
			OnModelDownloadComplete.Broadcast(ModelId);

			UE_LOG(LogRuntimeSpeechRecognizerVoskEditor, Log, TEXT("Successfully downloaded and extracted Vosk model '%s'"), *ModelId);
		});
	});
}

bool FRuntimeSpeechRecognizerVoskEditorModule::ExtractModelZip(const FString& ZipFilePath, const FString& ModelId, FText& OutError) const
{
	IFileHandle* ZipFileHandle = FPlatformFileManager::Get().GetPlatformFile().OpenRead(*ZipFilePath);
	if (!ZipFileHandle)
	{
		OutError = LOCTEXT("VoskModelZipOpenFailed", "Failed to open the downloaded archive for extraction");
		return false;
	}

	FZipArchiveReader ZipReader(ZipFileHandle);
	if (!ZipReader.IsValid())
	{
		OutError = LOCTEXT("VoskModelZipInvalid", "The downloaded archive is not a valid zip file");
		return false;
	}

	const TArray<FString> FileNames = ZipReader.GetFileNames();
	if (FileNames.Num() == 0)
	{
		OutError = LOCTEXT("VoskModelZipEmpty", "The downloaded archive did not contain any files");
		return false;
	}

	const FString ModelsDir = URuntimeVoskModelLibrary::GetVoskModelsDirectory();

	for (const FString& EntryName : FileNames)
	{
		if (EntryName.EndsWith(TEXT("/")))
		{
			continue;
		}

		TArray<uint8> FileData;
		if (!ZipReader.TryReadFile(EntryName, FileData))
		{
			OutError = FText::Format(LOCTEXT("VoskModelZipReadFailed", "Failed to read '{0}' from the downloaded archive"), FText::FromString(EntryName));
			return false;
		}

		const FString OutputPath = FPaths::Combine(ModelsDir, EntryName);
		IFileManager::Get().MakeDirectory(*FPaths::GetPath(OutputPath), true);

		if (!FFileHelper::SaveArrayToFile(FileData, *OutputPath))
		{
			OutError = FText::Format(LOCTEXT("VoskModelExtractWriteFailed", "Failed to write '{0}' while extracting the model"), FText::FromString(OutputPath));
			return false;
		}
	}

	if (!IFileManager::Get().DirectoryExists(*URuntimeVoskModelLibrary::GetVoskModelFolderPath(ModelId)))
	{
		OutError = FText::Format(LOCTEXT("VoskModelUnexpectedLayout", "The archive for '{0}' did not extract into the expected folder layout"), FText::FromString(ModelId));
		return false;
	}

	return true;
}

void FRuntimeSpeechRecognizerVoskEditorModule::CancelDownloadModel(const FString& ModelId)
{
	if (TSharedPtr<FRuntimeChunkDownloader_Recognizer>* Downloader = DownloadTasks.Find(ModelId))
	{
		if (Downloader->IsValid())
		{
			(*Downloader)->CancelDownload();
		}
	}
}

bool FRuntimeSpeechRecognizerVoskEditorModule::IsDownloadActive(const FString& ModelId) const
{
	return DownloadTasks.Contains(ModelId);
}

bool FRuntimeSpeechRecognizerVoskEditorModule::DeleteModel(const FString& ModelId)
{
	CancelDownloadModel(ModelId);
	const bool bDeleted = URuntimeVoskModelLibrary::DeleteVoskModelFiles(ModelId);
	UE_LOG(LogRuntimeSpeechRecognizerVoskEditor, Log, TEXT("Deleted Vosk model '%s': %s"), *ModelId, bDeleted ? TEXT("success") : TEXT("failed or not present"));
	return bDeleted;
}

#undef LOCTEXT_NAMESPACE

IMPLEMENT_MODULE(FRuntimeSpeechRecognizerVoskEditorModule, RuntimeSpeechRecognizerVoskEditor)