#pragma once

#include "CoreMinimal.h"
#include "MCPHandler.h"
#include "MCPAssets.h"
#include "MCPUtils.h"
#include "Materials/Material.h"
#include "Materials/MaterialFunction.h"
#include "Materials/MaterialExpression.h"
#include "MaterialGraph/MaterialGraph.h"
#include "MaterialGraph/MaterialGraphNode.h"
#include "UObject/UObjectIterator.h"
#include "UMCPHandler_AddMaterialExpression.generated.h"


// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------

UCLASS(meta=(Group="Unclassified"))
class UMCPHandler_AddMaterialExpression : public UObject, public IMCPHandler
{
	GENERATED_BODY()

public:
	UPROPERTY(meta=(Optional, Description="Material name or package path (specify this or materialFunction, not both)"))
	FString Material;

	UPROPERTY(meta=(Optional, Description="Material function name or package path (specify this or material, not both)"))
	FString MaterialFunction;

	UPROPERTY(meta=(Description="Expression class name without 'MaterialExpression' prefix (e.g. 'Constant', 'ScalarParameter', 'Add', 'Multiply', 'Lerp')"))
	FString ExpressionClass;

	UPROPERTY(meta=(Optional, Description="X position in the material graph editor"))
	int32 PosX = 0;

	UPROPERTY(meta=(Optional, Description="Y position in the material graph editor"))
	int32 PosY = 0;

	virtual FString GetDescription() const override
	{
		return TEXT("Add a new expression node to a material or material function graph.");
	}

	virtual void Handle(const FJsonObject* Json, FStringBuilderBase& Result) override
	{
		if (Material.IsEmpty() && MaterialFunction.IsEmpty())
		{
			Result.Append(TEXT("ERROR: Specify 'material' or 'materialFunction'\n"));
			return;
		}
		if (!Material.IsEmpty() && !MaterialFunction.IsEmpty())
		{
			Result.Append(TEXT("ERROR: Specify 'material' or 'materialFunction', not both\n"));
			return;
		}

		// Resolve the expression class
		UClass* ExprClass = ResolveExpressionClass(Result);
		if (!ExprClass) return;

		// Load material or material function
		UMaterial* MaterialObj = nullptr;
		UMaterialFunction* MatFunc = nullptr;
		UObject* Owner = nullptr;

		if (!MaterialFunction.IsEmpty())
		{
			MCPAssets<UMaterialFunction> Assets;
			if (!Assets.Exact(MaterialFunction).Errors(Result).ENone().ETwo().Load()) return;
			MatFunc = Assets.Object();
			Owner = MatFunc;
		}
		else
		{
			MCPAssets<UMaterial> Assets;
			if (!Assets.Exact(Material).Errors(Result).ENone().ETwo().Load()) return;
			MaterialObj = Assets.Object();
			Owner = MaterialObj;
		}

		// Ensure the MaterialGraph exists (commandlet mode doesn't auto-create it)
		if (MaterialObj) MCPUtils::EnsureMaterialGraph(MaterialObj);

		// Create the expression
		UMaterialExpression* NewExpr = NewObject<UMaterialExpression>(Owner, ExprClass);
		if (!NewExpr)
		{
			Result.Append(TEXT("ERROR: Failed to create material expression object\n"));
			return;
		}
		NewExpr->MaterialExpressionEditorX = PosX;
		NewExpr->MaterialExpressionEditorY = PosY;

		if (MaterialObj)
		{
			MaterialObj->GetExpressionCollection().AddExpression(NewExpr);
			if (MaterialObj->MaterialGraph)
				MaterialObj->MaterialGraph->RebuildGraph();
			MaterialObj->PreEditChange(nullptr);
			MaterialObj->PostEditChange();
			MaterialObj->MarkPackageDirty();
		}
		else
		{
			MatFunc->GetExpressionCollection().AddExpression(NewExpr);
			MatFunc->PreEditChange(nullptr);
			MatFunc->PostEditChange();
			MatFunc->MarkPackageDirty();
		}

		// Save
		bool bSaved = MaterialObj
			? MCPUtils::SaveMaterialPackage(MaterialObj)
			: MCPUtils::SaveGenericPackage(MatFunc);

		// Output
		Result.Appendf(TEXT("Added %s\n"), *MCPUtils::FormatName(NewExpr));

		// Find the graph node GUID (only for materials with a material graph)
		if (MaterialObj && MaterialObj->MaterialGraph)
		{
			for (UEdGraphNode* Node : MaterialObj->MaterialGraph->Nodes)
			{
				UMaterialGraphNode* MatNode = Cast<UMaterialGraphNode>(Node);
				if (MatNode && MatNode->MaterialExpression == NewExpr)
				{
					Result.Appendf(TEXT("NodeId: %s\n"), *Node->NodeGuid.ToString());
					break;
				}
			}
		}

		if (!bSaved)
			Result.Append(TEXT("WARNING: Failed to save package\n"));
	}

private:
	UClass* ResolveExpressionClass(FStringBuilderBase& Result)
	{
		// Convenience aliases
		static TMap<FString, FString> Aliases = {
			{TEXT("Lerp"), TEXT("LinearInterpolate")},
		};

		FString LookupName = ExpressionClass;
		if (const FString* Alias = Aliases.Find(ExpressionClass))
			LookupName = *Alias;

		FString FullClassName = FString::Printf(TEXT("MaterialExpression%s"), *LookupName);
		for (TObjectIterator<UClass> It; It; ++It)
		{
			if (It->GetName() == FullClassName && It->IsChildOf(UMaterialExpression::StaticClass()))
			{
				if (It->HasAnyClassFlags(CLASS_Abstract))
				{
					Result.Appendf(TEXT("ERROR: Expression class '%s' is abstract\n"), *ExpressionClass);
					return nullptr;
				}
				return *It;
			}
		}

		Result.Appendf(TEXT("ERROR: Unknown expression class '%s'. Use the UMaterialExpression subclass name without the 'MaterialExpression' prefix (e.g. 'Constant', 'ScalarParameter', 'Add', 'Multiply', 'Lerp')\n"),
			*ExpressionClass);
		return nullptr;
	}
};