#pragma once

#include "CoreMinimal.h"
#include "MCPHandler.h"
#include "MCPAssets.h"
#include "MCPFetcher.h"
#include "MCPUtils.h"
#include "Materials/Material.h"
#include "Materials/MaterialFunction.h"
#include "Materials/MaterialExpression.h"
#include "Materials/MaterialExpressionScalarParameter.h"
#include "Materials/MaterialExpressionVectorParameter.h"
#include "Materials/MaterialExpressionConstant.h"
#include "Materials/MaterialExpressionConstant3Vector.h"
#include "Materials/MaterialExpressionConstant4Vector.h"
#include "Materials/MaterialExpressionTextureSample.h"
#include "Materials/MaterialExpressionTextureCoordinate.h"
#include "Materials/MaterialExpressionComponentMask.h"
#include "Materials/MaterialExpressionCustom.h"
#include "MaterialGraph/MaterialGraph.h"
#include "MaterialGraph/MaterialGraphNode.h"
#include "UMCPHandler_SetMaterialExpressionProperty.generated.h"


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

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

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

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

	UPROPERTY(meta=(Description="Expression node name (from FormatName)"))
	FString Node;

	virtual FString GetDescription() const override
	{
		return TEXT("Set the value or properties on a material expression node. "
			"The 'value' field in the JSON payload provides the new value, whose format depends on the expression type.");
	}

	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 (!Json->HasField(TEXT("value")))
		{
			Result.Append(TEXT("ERROR: Missing required field: value\n"));
			return;
		}

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

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

		if (MaterialObj) MCPUtils::EnsureMaterialGraph(MaterialObj);
		UEdGraph* Graph = MaterialObj ? (UEdGraph*)MaterialObj->MaterialGraph : (MatFunc ? MatFunc->MaterialGraph : nullptr);
		if (!Graph)
		{
			Result.Appendf(TEXT("ERROR: No material graph found\n"));
			return;
		}

		// Find the node by name using Identifies
		UMaterialGraphNode* TargetMatNode = nullptr;
		for (UEdGraphNode* GraphNode : Graph->Nodes)
		{
			if (!GraphNode) continue;
			if (!MCPUtils::Identifies(Node, GraphNode)) continue;
			TargetMatNode = Cast<UMaterialGraphNode>(GraphNode);
			if (TargetMatNode) break;
		}

		if (!TargetMatNode)
		{
			Result.Appendf(TEXT("ERROR: Node '%s' not found in material graph\n"), *Node);
			return;
		}

		UMaterialExpression* Expr = TargetMatNode->MaterialExpression;
		if (!Expr)
		{
			Result.Appendf(TEXT("ERROR: Node '%s' has no material expression\n"), *MCPUtils::FormatName(TargetMatNode));
			return;
		}

		UObject* Asset = MaterialObj ? (UObject*)MaterialObj : (UObject*)MatFunc;
		Asset->PreEditChange(nullptr);

		FString SetResult;
		if (!ApplyValue(Expr, Json, Result, SetResult))
		{
			Asset->PostEditChange();
			return;
		}

		Asset->PostEditChange();
		Asset->MarkPackageDirty();

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

		Result.Appendf(TEXT("%s = %s"), *MCPUtils::FormatName(Expr), *SetResult);
		if (!bSaved) Result.Append(TEXT(" (save failed)"));
		Result.Append(TEXT("\n"));
	}

private:
	// Apply the value from JSON to the expression. Returns false on error (with message in Result).
	// On success, fills SetResult with a human-readable summary of the new value.
	bool ApplyValue(UMaterialExpression* Expr, const FJsonObject* Json, FStringBuilderBase& Result, FString& SetResult)
	{
		if (auto* E = Cast<UMaterialExpressionConstant>(Expr))
		{
			double Value = Json->GetNumberField(TEXT("value"));
			E->R = (float)Value;
			SetResult = FString::Printf(TEXT("%g"), Value);
			return true;
		}

		if (auto* E = Cast<UMaterialExpressionConstant3Vector>(Expr))
		{
			FLinearColor C;
			if (!ParseColorValue(Json, C, false, Result)) return false;
			E->Constant = C;
			SetResult = FString::Printf(TEXT("(%.3f, %.3f, %.3f)"), C.R, C.G, C.B);
			return true;
		}

		if (auto* E = Cast<UMaterialExpressionConstant4Vector>(Expr))
		{
			FLinearColor C;
			if (!ParseColorValue(Json, C, true, Result)) return false;
			E->Constant = C;
			SetResult = FString::Printf(TEXT("(%.3f, %.3f, %.3f, %.3f)"), C.R, C.G, C.B, C.A);
			return true;
		}

		if (auto* E = Cast<UMaterialExpressionScalarParameter>(Expr))
		{
			double Value = Json->GetNumberField(TEXT("value"));
			E->DefaultValue = (float)Value;
			SetResult = FString::Printf(TEXT("%g"), Value);
			ApplyParameterName(E, Json);
			return true;
		}

		if (auto* E = Cast<UMaterialExpressionVectorParameter>(Expr))
		{
			FLinearColor C;
			if (!ParseColorValue(Json, C, true, Result)) return false;
			E->DefaultValue = C;
			SetResult = FString::Printf(TEXT("(%.3f, %.3f, %.3f, %.3f)"), C.R, C.G, C.B, C.A);
			ApplyParameterName(E, Json);
			return true;
		}

		if (auto* E = Cast<UMaterialExpressionTextureCoordinate>(Expr))
		{
			const TSharedPtr<FJsonObject>* ValueObj = nullptr;
			if (!Json->TryGetObjectField(TEXT("value"), ValueObj) || !ValueObj || !(*ValueObj).IsValid())
			{
				Result.Append(TEXT("ERROR: TextureCoordinate requires value as {coordinateIndex, uTiling, vTiling}\n"));
				return false;
			}
			double CoordIndex = 0, UTiling = 1, VTiling = 1;
			(*ValueObj)->TryGetNumberField(TEXT("coordinateIndex"), CoordIndex);
			(*ValueObj)->TryGetNumberField(TEXT("uTiling"), UTiling);
			(*ValueObj)->TryGetNumberField(TEXT("vTiling"), VTiling);
			E->CoordinateIndex = (int32)CoordIndex;
			E->UTiling = (float)UTiling;
			E->VTiling = (float)VTiling;
			SetResult = FString::Printf(TEXT("index=%d uTiling=%g vTiling=%g"), (int32)CoordIndex, UTiling, VTiling);
			return true;
		}

		if (auto* E = Cast<UMaterialExpressionCustom>(Expr))
		{
			FString Code;
			if (Json->TryGetStringField(TEXT("code"), Code))
			{
				E->Code = Code;
			}
			else if (Json->HasField(TEXT("value")))
			{
				FString ValueStr = Json->GetStringField(TEXT("value"));
				if (!ValueStr.IsEmpty()) E->Code = ValueStr;
			}
			SetResult = FString::Printf(TEXT("code: %d chars"), E->Code.Len());

			FString OutputTypeStr;
			if (Json->TryGetStringField(TEXT("outputType"), OutputTypeStr) && !OutputTypeStr.IsEmpty())
			{
				ECustomMaterialOutputType OutType;
				if (MCPUtils::StringToEnum(OutputTypeStr, OutType, MCPErrorCallback(Result)))
					E->OutputType = OutType;
			}
			return true;
		}

		if (auto* E = Cast<UMaterialExpressionComponentMask>(Expr))
		{
			const TSharedPtr<FJsonObject>* ValueObj = nullptr;
			if (!Json->TryGetObjectField(TEXT("value"), ValueObj) || !ValueObj || !(*ValueObj).IsValid())
			{
				Result.Append(TEXT("ERROR: ComponentMask requires value as {r, g, b, a} (booleans)\n"));
				return false;
			}
			bool bR = false, bG = false, bB = false, bA = false;
			(*ValueObj)->TryGetBoolField(TEXT("r"), bR);
			(*ValueObj)->TryGetBoolField(TEXT("g"), bG);
			(*ValueObj)->TryGetBoolField(TEXT("b"), bB);
			(*ValueObj)->TryGetBoolField(TEXT("a"), bA);
			E->R = bR ? 1 : 0;
			E->G = bG ? 1 : 0;
			E->B = bB ? 1 : 0;
			E->A = bA ? 1 : 0;
			SetResult = FString::Printf(TEXT("R=%s G=%s B=%s A=%s"),
				bR ? TEXT("true") : TEXT("false"),
				bG ? TEXT("true") : TEXT("false"),
				bB ? TEXT("true") : TEXT("false"),
				bA ? TEXT("true") : TEXT("false"));
			return true;
		}

		Result.Appendf(TEXT("ERROR: Expression type '%s' does not support value setting. Supported: "
			"Constant, Constant3Vector, Constant4Vector, ScalarParameter, VectorParameter, "
			"TextureCoordinate, Custom, ComponentMask\n"),
			*Expr->GetClass()->GetName());
		return false;
	}

	// Parse {r, g, b[, a]} from the "value" JSON field.
	bool ParseColorValue(const FJsonObject* Json, FLinearColor& OutColor, bool bHasAlpha, FStringBuilderBase& Result)
	{
		const TSharedPtr<FJsonObject>* ValueObj = nullptr;
		if (!Json->TryGetObjectField(TEXT("value"), ValueObj) || !ValueObj || !(*ValueObj).IsValid())
		{
			Result.Appendf(TEXT("ERROR: requires value as {r, g, b%s}\n"), bHasAlpha ? TEXT(", a") : TEXT(""));
			return false;
		}
		double R = 0, G = 0, B = 0, A = 1;
		(*ValueObj)->TryGetNumberField(TEXT("r"), R);
		(*ValueObj)->TryGetNumberField(TEXT("g"), G);
		(*ValueObj)->TryGetNumberField(TEXT("b"), B);
		if (bHasAlpha) (*ValueObj)->TryGetNumberField(TEXT("a"), A);
		OutColor = FLinearColor((float)R, (float)G, (float)B, (float)A);
		return true;
	}

	// If the JSON has a "parameterName" field, apply it to a parameter expression.
	void ApplyParameterName(UMaterialExpressionParameter* Param, const FJsonObject* Json)
	{
		FString ParamName;
		if (Json->TryGetStringField(TEXT("parameterName"), ParamName) && !ParamName.IsEmpty())
			Param->ParameterName = FName(*ParamName);
	}
};