// Copyright UShaderLab. All Rights Reserved.

#include "ShaderLabVSCodeButton.h"

#include "Framework/MultiBox/MultiBoxBuilder.h"
#include "Framework/MultiBox/MultiBoxExtender.h"
#include "HAL/PlatformProcess.h"
#include "MaterialEditorModule.h"
#include "Misc/Paths.h"
#include "Modules/ModuleManager.h"
#include "ShaderLabGraphBuilder.h"
#include "ShaderLabMaterialInstanceConstant.h"
#include "Styling/AppStyle.h"
#include "Textures/SlateIcon.h"
#include "Toolkits/AssetEditorToolkit.h"

#define LOCTEXT_NAMESPACE "ShaderLabVSCode"

DEFINE_LOG_CATEGORY_STATIC(LogShaderLabVSCode, Log, All);

namespace ShaderLabVSCodeButton
{
	static FDelegateHandle GExtenderDelegateHandle;

	/** Reverse the ShaderLabPath virtual path (e.g. "/Project/Examples/Basic.usl") to an absolute disk path
	 *  via the shared shader-mapping resolver (generic — /Project, /Plugin/<any>, /Engine). Empty if unresolved. */
	static FString ResolveVirtualToDisk(const FString& VirtualPath)
	{
		FString Disk;
		if (!FShaderLabGraphBuilder::ResolveVirtualShaderFile(VirtualPath, Disk))
		{
			return FString();
		}
		return FPaths::ConvertRelativePathToFull(Disk);
	}

	static void OpenInVSCode(TWeakObjectPtr<UShaderLabMaterialInstanceConstant> MICWeak)
	{
		UShaderLabMaterialInstanceConstant* MIC = MICWeak.Get();
		if (!MIC)
		{
			return;
		}
		const FString VirtualPath = UShaderLabMaterialInstanceConstant::GetShaderLabPath(MIC);
		if (VirtualPath.IsEmpty())
		{
			UE_LOG(LogShaderLabVSCode, Warning, TEXT("'%s' has no ShaderLabPath."), *MIC->GetName());
			return;
		}
		const FString DiskPath = ResolveVirtualToDisk(VirtualPath);
		if (DiskPath.IsEmpty() || !FPaths::FileExists(DiskPath))
		{
			UE_LOG(LogShaderLabVSCode, Warning, TEXT("Cannot resolve '%s' to an existing file (got '%s')."), *VirtualPath, *DiskPath);
			return;
		}

		// `code` is a .cmd on Windows, so launch it through the command interpreter; it opens/reuses VSCode.
#if PLATFORM_WINDOWS
		const FString Exe = TEXT("cmd.exe");
		const FString Args = FString::Printf(TEXT("/c code \"%s\""), *DiskPath);
#else
		const FString Exe = TEXT("code");
		const FString Args = FString::Printf(TEXT("\"%s\""), *DiskPath);
#endif
		FProcHandle Proc = FPlatformProcess::CreateProc(*Exe, *Args, /*bLaunchDetached*/ true,
			/*bLaunchHidden*/ true, /*bLaunchReallyHidden*/ true, nullptr, 0, nullptr, nullptr);
		if (Proc.IsValid())
		{
			FPlatformProcess::CloseProc(Proc);
		}
		else
		{
			UE_LOG(LogShaderLabVSCode, Warning, TEXT("Failed to launch VSCode ('code' on PATH?) for '%s'."), *DiskPath);
		}
	}

	/** Context-aware extender: add the button only when a ShaderLab MIC is being edited. */
	static TSharedRef<FExtender> OnExtendToolbar(const TSharedRef<FUICommandList> Commands, const TArray<UObject*> EditedObjects)
	{
		TSharedRef<FExtender> Extender = MakeShared<FExtender>();

		UShaderLabMaterialInstanceConstant* MIC = nullptr;
		for (UObject* Object : EditedObjects)
		{
			if (UShaderLabMaterialInstanceConstant* Found = Cast<UShaderLabMaterialInstanceConstant>(Object))
			{
				MIC = Found;
				break;
			}
		}
		if (!MIC || UShaderLabMaterialInstanceConstant::GetShaderLabPath(MIC).IsEmpty())
		{
			return Extender;
		}

		TWeakObjectPtr<UShaderLabMaterialInstanceConstant> MICWeak(MIC);
		Extender->AddToolBarExtension("Asset", EExtensionHook::After, Commands,
			FToolBarExtensionDelegate::CreateLambda([MICWeak](FToolBarBuilder& Builder)
			{
				Builder.AddToolBarButton(
					FUIAction(FExecuteAction::CreateStatic(&OpenInVSCode, MICWeak)),
					NAME_None,
					LOCTEXT("OpenSourceLabel", "Open .usl"),
					LOCTEXT("OpenSourceTooltip", "Open this ShaderLab material's source .usl in VSCode"),
					FSlateIcon(FAppStyle::GetAppStyleSetName(), "Icons.Edit"));
			}));
		return Extender;
	}

	void Register()
	{
		IMaterialEditorModule& MaterialEditor = FModuleManager::LoadModuleChecked<IMaterialEditorModule>("MaterialEditor");
		TArray<FAssetEditorExtender>& Delegates = MaterialEditor.GetToolBarExtensibilityManager()->GetExtenderDelegates();
		Delegates.Add(FAssetEditorExtender::CreateStatic(&OnExtendToolbar));
		GExtenderDelegateHandle = Delegates.Last().GetHandle();
	}

	void Unregister()
	{
		if (IMaterialEditorModule* MaterialEditor = FModuleManager::GetModulePtr<IMaterialEditorModule>("MaterialEditor"))
		{
			MaterialEditor->GetToolBarExtensibilityManager()->GetExtenderDelegates().RemoveAll(
				[](const FAssetEditorExtender& Delegate) { return Delegate.GetHandle() == GExtenderDelegateHandle; });
		}
	}
}

#undef LOCTEXT_NAMESPACE