#include "Extensions/DruidsMCPBridgeSageExtensionTypes.h"

#include "Serialization/JsonReader.h"
#include "Serialization/JsonSerializer.h"
#include "Serialization/JsonWriter.h"
#include "UObject/TextProperty.h"

FDruidsMCPBridgeSageActionResult FDruidsMCPBridgeSageActionResult::NoAction()
{
	return FDruidsMCPBridgeSageActionResult();
}

FDruidsMCPBridgeSageActionResult FDruidsMCPBridgeSageActionResult::ExtensionSuccess(const FString& InResult)
{
	FDruidsMCPBridgeSageActionResult Result;
	Result.bActionCalled = true;
	Result.bSuccess = true;
	Result.Result = InResult;
	return Result;
}

FDruidsMCPBridgeSageActionResult FDruidsMCPBridgeSageActionResult::ExtensionFailure(const FString& InErrorMessage)
{
	FDruidsMCPBridgeSageActionResult Result;
	Result.bActionCalled = true;
	Result.bSuccess = false;
	Result.ErrorMessage = InErrorMessage;
	return Result;
}

FDruidsMCPBridgeSageActionResult FDruidsMCPBridgeSageActionResult::AsyncAction(const FString& InResult)
{
	FDruidsMCPBridgeSageActionResult Result;
	Result.bActionCalled = true;
	Result.bSuccess = true;
	Result.bIsAsync = true;
	Result.Result = InResult;
	return Result;
}

EDruidsMCPBridgeSageExtensionParameterType FDruidsMCPBridgeSageParameterDefinition::GetParameterTypeForProperty(const FProperty* Property)
{
	if (!Property)
	{
		return EDruidsMCPBridgeSageExtensionParameterType::None;
	}
	if (Property->IsA<FArrayProperty>())
	{
		return EDruidsMCPBridgeSageExtensionParameterType::Array;
	}
	if (Property->IsA<FStrProperty>() || Property->IsA<FNameProperty>() || Property->IsA<FTextProperty>())
	{
		return EDruidsMCPBridgeSageExtensionParameterType::String;
	}
	if (Property->IsA<FIntProperty>())
	{
		return EDruidsMCPBridgeSageExtensionParameterType::Int;
	}
	if (Property->IsA<FFloatProperty>() || Property->IsA<FDoubleProperty>())
	{
		return EDruidsMCPBridgeSageExtensionParameterType::Float;
	}
	if (Property->IsA<FBoolProperty>())
	{
		return EDruidsMCPBridgeSageExtensionParameterType::Bool;
	}
	if (const FStructProperty* StructProperty = CastField<FStructProperty>(Property))
	{
		if (StructProperty->Struct->GetFName() == NAME_Vector)
		{
			return EDruidsMCPBridgeSageExtensionParameterType::Vector3;
		}
		if (StructProperty->Struct->GetFName() == NAME_Color)
		{
			return EDruidsMCPBridgeSageExtensionParameterType::Color;
		}
		if (StructProperty->Struct->GetFName() == NAME_LinearColor)
		{
			return EDruidsMCPBridgeSageExtensionParameterType::LinearColor;
		}
		if (StructProperty->Struct->GetFName() == NAME_Rotator)
		{
			return EDruidsMCPBridgeSageExtensionParameterType::Rotator;
		}
		if (StructProperty->Struct->GetFName() == NAME_Transform)
		{
			return EDruidsMCPBridgeSageExtensionParameterType::Transform;
		}
		return EDruidsMCPBridgeSageExtensionParameterType::CustomStruct;
	}
	if (Property->IsA<FEnumProperty>() || Property->IsA<FByteProperty>())
	{
		return EDruidsMCPBridgeSageExtensionParameterType::Enum;
	}
	if (Property->IsA<FObjectProperty>())
	{
		return EDruidsMCPBridgeSageExtensionParameterType::Object;
	}

	return EDruidsMCPBridgeSageExtensionParameterType::None;
}

FString FDruidsMCPBridgeSageParameterDefinition::GetParameterStringForType(EDruidsMCPBridgeSageExtensionParameterType ParameterType)
{
	switch (ParameterType)
	{
	case EDruidsMCPBridgeSageExtensionParameterType::String: return TEXT("String");
	case EDruidsMCPBridgeSageExtensionParameterType::Int: return TEXT("Int");
	case EDruidsMCPBridgeSageExtensionParameterType::Float: return TEXT("Float");
	case EDruidsMCPBridgeSageExtensionParameterType::Bool: return TEXT("Bool");
	case EDruidsMCPBridgeSageExtensionParameterType::Vector3: return TEXT("Vector3");
	case EDruidsMCPBridgeSageExtensionParameterType::Color: return TEXT("Color");
	case EDruidsMCPBridgeSageExtensionParameterType::Rotator: return TEXT("Rotator");
	case EDruidsMCPBridgeSageExtensionParameterType::LinearColor: return TEXT("LinearColor");
	case EDruidsMCPBridgeSageExtensionParameterType::Enum: return TEXT("Enum");
	case EDruidsMCPBridgeSageExtensionParameterType::Array: return TEXT("Array");
	case EDruidsMCPBridgeSageExtensionParameterType::Transform: return TEXT("Transform");
	case EDruidsMCPBridgeSageExtensionParameterType::CustomStruct: return TEXT("CustomStruct");
	case EDruidsMCPBridgeSageExtensionParameterType::Object: return TEXT("Object");
	default: return TEXT("Unknown");
	}
}

void FDruidsMCPBridgeSageParameterDefinition::RefreshJsonString(const FProperty* Property)
{
	TSharedPtr<FJsonObject> ParamObject = MakeShared<FJsonObject>();
	ParamObject->SetStringField(TEXT("name"), Name.ToString());
	ParamObject->SetStringField(TEXT("description"), Description);
	InsertPropertyAsJson(ParamObject, Property);

	ParameterTypeJsonString.Reset();
	const TSharedRef<TJsonWriter<>> Writer = TJsonWriterFactory<>::Create(&ParameterTypeJsonString);
	FJsonSerializer::Serialize(ParamObject.ToSharedRef(), Writer);
}

void FDruidsMCPBridgeSageParameterDefinition::InsertPropertyAsJson(const TSharedPtr<FJsonObject>& ParamObject, const FProperty* Property)
{
	const EDruidsMCPBridgeSageExtensionParameterType Type = GetParameterTypeForProperty(Property);
	ParamObject->SetStringField(TEXT("type"), GetParameterStringForType(Type));

	if (Type == EDruidsMCPBridgeSageExtensionParameterType::Array)
	{
		const FArrayProperty* ArrayProp = CastField<FArrayProperty>(Property);
		ParamObject->SetObjectField(TEXT("arrayElement"), GetArrayTypeAsJson(ArrayProp ? ArrayProp->Inner : nullptr));
	}
	else if (Type == EDruidsMCPBridgeSageExtensionParameterType::CustomStruct)
	{
		const FStructProperty* StructProp = CastField<FStructProperty>(Property);
		ParamObject->SetObjectField(TEXT("structElement"), GetStructTypeInfo(StructProp ? StructProp->Struct : nullptr));
	}
}

TSharedPtr<FJsonObject> FDruidsMCPBridgeSageParameterDefinition::GetArrayTypeAsJson(const FProperty* ArrayType)
{
	TSharedPtr<FJsonObject> ArrayInfo = MakeShared<FJsonObject>();
	InsertPropertyAsJson(ArrayInfo, ArrayType);
	return ArrayInfo;
}

TSharedPtr<FJsonObject> FDruidsMCPBridgeSageParameterDefinition::GetStructTypeInfo(TWeakObjectPtr<UScriptStruct> StructType)
{
	TSharedPtr<FJsonObject> StructInfo = MakeShared<FJsonObject>();
	if (!StructType.IsValid())
	{
		return StructInfo;
	}

	StructInfo->SetStringField(TEXT("structType"), StructType->GetAuthoredName());
	TArray<TSharedPtr<FJsonValue>> Parameters;
	for (TFieldIterator<FProperty> PropIt(StructType.Get()); PropIt; ++PropIt)
	{
		FProperty* StructMemberProp = *PropIt;
		TSharedPtr<FJsonObject> PropObject = MakeShared<FJsonObject>();
		PropObject->SetStringField(TEXT("name"), StructMemberProp->GetAuthoredName());
		InsertPropertyAsJson(PropObject, StructMemberProp);
		Parameters.Add(MakeShared<FJsonValueObject>(PropObject));
	}
	StructInfo->SetArrayField(TEXT("parameters"), Parameters);
	return StructInfo;
}

TSharedPtr<FJsonObject> FDruidsMCPBridgeSageExtensionDefinition::GetExtensionDefinitionJson() const
{
	TSharedPtr<FJsonObject> ExtensionObject = MakeShared<FJsonObject>();
	ExtensionObject->SetStringField(TEXT("extension_id"), ExtensionId);
	ExtensionObject->SetStringField(TEXT("title"), ExtensionName);
	ExtensionObject->SetStringField(TEXT("description"), ExtensionDescription);
	ExtensionObject->SetStringField(TEXT("contextType"), UEnum::GetValueAsString(ContextType));
	ExtensionObject->SetStringField(TEXT("assetType"), SelectedAssetType);
	ExtensionObject->SetBoolField(TEXT("is_system_extension"), bIsSystemExtension);

	TArray<TSharedPtr<FJsonValue>> ActionsArray;
	for (const FDruidsMCPBridgeSageActionDefinition& Action : Actions)
	{
		TSharedPtr<FJsonObject> ActionObject = MakeShared<FJsonObject>();
		ActionObject->SetStringField(TEXT("name"), Action.FunctionName.ToString());
		ActionObject->SetStringField(TEXT("description"), Action.Description);

		TArray<TSharedPtr<FJsonValue>> ParametersArray;
		for (const FDruidsMCPBridgeSageParameterDefinition& Param : Action.ParameterDefinitions)
		{
			TSharedPtr<FJsonObject> ParamObject = MakeShared<FJsonObject>();
			const TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create(Param.ParameterTypeJsonString);
			if (FJsonSerializer::Deserialize(Reader, ParamObject) && ParamObject.IsValid())
			{
				ParametersArray.Add(MakeShared<FJsonValueObject>(ParamObject));
			}
		}
		ActionObject->SetArrayField(TEXT("parameters"), ParametersArray);
		ActionsArray.Add(MakeShared<FJsonValueObject>(ActionObject));
	}
	ExtensionObject->SetArrayField(TEXT("actions"), ActionsArray);

	TArray<TSharedPtr<FJsonValue>> QueriesArray;
	for (const FDruidsMCPBridgeSageQueryDefinition& Query : Queries)
	{
		TSharedPtr<FJsonObject> QueryObject = MakeShared<FJsonObject>();
		QueryObject->SetStringField(TEXT("name"), Query.FunctionName.ToString());
		QueryObject->SetStringField(TEXT("description"), Query.Description);

		TArray<TSharedPtr<FJsonValue>> ParametersArray;
		for (const FDruidsMCPBridgeSageParameterDefinition& Param : Query.ParameterDefinitions)
		{
			TSharedPtr<FJsonObject> ParamObject = MakeShared<FJsonObject>();
			const TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create(Param.ParameterTypeJsonString);
			if (FJsonSerializer::Deserialize(Reader, ParamObject) && ParamObject.IsValid())
			{
				ParametersArray.Add(MakeShared<FJsonValueObject>(ParamObject));
			}
		}
		QueryObject->SetArrayField(TEXT("parameters"), ParametersArray);

		TArray<TSharedPtr<FJsonValue>> ReturnParametersArray;
		for (const FDruidsMCPBridgeSageParameterDefinition& Param : Query.ResultDefinitions)
		{
			TSharedPtr<FJsonObject> ParamObject = MakeShared<FJsonObject>();
			const TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create(Param.ParameterTypeJsonString);
			if (FJsonSerializer::Deserialize(Reader, ParamObject) && ParamObject.IsValid())
			{
				ReturnParametersArray.Add(MakeShared<FJsonValueObject>(ParamObject));
			}
		}
		QueryObject->SetArrayField(TEXT("returnParameters"), ReturnParametersArray);
		QueriesArray.Add(MakeShared<FJsonValueObject>(QueryObject));
	}
	ExtensionObject->SetArrayField(TEXT("queries"), QueriesArray);

	return ExtensionObject;
}
