#include "Extensions/DruidsMCPBridgeSageExtension.h"

#include "DruidsMCPBridge.h"
#include "Editor.h"
#include "Editor/EditorEngine.h"
#include "JsonObjectConverter.h"
#include "ScopedTransaction.h"
#include "Serialization/JsonReader.h"
#include "Serialization/JsonSerializer.h"
#include "Serialization/JsonSerializer.h"
#include "Subsystems/AssetEditorSubsystem.h"
#include "UObject/Script.h"
#include "UObject/UnrealType.h"
#include "UObject/TextProperty.h"

UDruidsMCPBridgeSageExtension::UDruidsMCPBridgeSageExtension()
{
}

bool UDruidsMCPBridgeSageExtension::IsActive_Implementation() const
{
	return true;
}

void UDruidsMCPBridgeSageExtension::PostInitProperties()
{
	Super::PostInitProperties();

	if (HasAnyFlags(RF_ClassDefaultObject) && GetClass() != StaticClass())
	{
		RefreshExtension();
	}
}

void UDruidsMCPBridgeSageExtension::PostCDOCompiled(const FPostCDOCompiledContext& Context)
{
	Super::PostCDOCompiled(Context);

	if (GetClass() != StaticClass())
	{
		RefreshExtension();
	}
}

void UDruidsMCPBridgeSageExtension::RefreshExtension()
{
	if (!HasAnyFlags(RF_ClassDefaultObject) || GetClass() == StaticClass())
	{
		return;
	}

	Definition.ExtensionId = GetClass()->GetPathName();
	if (Definition.ExtensionName.IsEmpty())
	{
		Definition.ExtensionName = GetClass()->GetDisplayNameText().ToString();
	}

	DiscoverActionsAndQueries();
}

bool UDruidsMCPBridgeSageExtension::IsEnabled() const
{
	return bEnabled && IsActive();
}

bool UDruidsMCPBridgeSageExtension::IsSystemExtension() const
{
	return false;
}

bool UDruidsMCPBridgeSageExtension::MatchesContext(TWeakObjectPtr<UObject> ActiveObject) const
{
	if (Definition.ContextType == EDruidsMCPBridgeSageExtensionContextType::Any)
	{
		return true;
	}

	if (!ActiveObject.IsValid())
	{
		return false;
	}

	return Definition.SelectedAssetType.IsEmpty() || ActiveObject->GetClass()->GetName().Contains(Definition.SelectedAssetType);
}

TSharedPtr<FDruidsMCPBridgeSageExtensionDefinition> UDruidsMCPBridgeSageExtension::GetExtensionDefinition() const
{
	TSharedPtr<FDruidsMCPBridgeSageExtensionDefinition> ExtensionDef = MakeShared<FDruidsMCPBridgeSageExtensionDefinition>(Definition);
	if (ExtensionDef->ExtensionId.IsEmpty())
	{
		ExtensionDef->ExtensionId = GetClass()->GetPathName();
	}
	if (ExtensionDef->ExtensionName.IsEmpty())
	{
		ExtensionDef->ExtensionName = GetClass()->GetDisplayNameText().ToString();
	}
	ExtensionDef->bIsSystemExtension = IsSystemExtension();
	return ExtensionDef;
}

void UDruidsMCPBridgeSageExtension::DiscoverActionsAndQueries()
{
	TMap<FString, FString> ExistingDescriptions;
	for (const FDruidsMCPBridgeSageActionDefinition& Action : Definition.Actions)
	{
		for (const FDruidsMCPBridgeSageParameterDefinition& Param : Action.ParameterDefinitions)
		{
			if (!Param.Description.IsEmpty())
			{
				ExistingDescriptions.Add(FString::Printf(TEXT("Action_%s_%s"), *Action.FunctionName.ToString(), *Param.Name.ToString()), Param.Description);
			}
		}
	}
	for (const FDruidsMCPBridgeSageQueryDefinition& Query : Definition.Queries)
	{
		for (const FDruidsMCPBridgeSageParameterDefinition& Param : Query.ParameterDefinitions)
		{
			if (!Param.Description.IsEmpty())
			{
				ExistingDescriptions.Add(FString::Printf(TEXT("Query_Param_%s_%s"), *Query.FunctionName.ToString(), *Param.Name.ToString()), Param.Description);
			}
		}
		for (const FDruidsMCPBridgeSageParameterDefinition& Result : Query.ResultDefinitions)
		{
			if (!Result.Description.IsEmpty())
			{
				ExistingDescriptions.Add(FString::Printf(TEXT("Query_Result_%s_%s"), *Query.FunctionName.ToString(), *Result.Name.ToString()), Result.Description);
			}
		}
	}

	Definition.Actions.Empty();
	Definition.Queries.Empty();

	for (TFieldIterator<UFunction> FuncIt(GetClass(), EFieldIteratorFlags::IncludeSuper); FuncIt; ++FuncIt)
	{
		DiscoverFunction(*FuncIt, ExistingDescriptions);
	}
}

void UDruidsMCPBridgeSageExtension::DiscoverFunction(UFunction* Function, const TMap<FString, FString>& ExistingDescriptions)
{
	if (!Function)
	{
		return;
	}

	const FString Category = Function->GetMetaData(TEXT("Category")).Replace(TEXT(" "), TEXT(""));
	const bool bMetaAction = Function->HasMetaData(TEXT("SageAction")) || Function->HasMetaData(TEXT("DruidsSageAction")) || Function->HasMetaData(TEXT("DruidSageAction"));
	const bool bMetaQuery = Function->HasMetaData(TEXT("SageQuery")) || Function->HasMetaData(TEXT("DruidsSageQuery")) || Function->HasMetaData(TEXT("DruidSageQuery"));
	const bool bIsAction = bMetaAction || Category.Contains(TEXT("SageActions")) || Category.Contains(TEXT("DruidSageActions")) || Category.Contains(TEXT("DruidsSageActions"));
	const bool bIsQuery = bMetaQuery || Category.Contains(TEXT("SageQueries")) || Category.Contains(TEXT("DruidSageQueries")) || Category.Contains(TEXT("DruidsSageQueries"));

	if (!bIsAction && !bIsQuery)
	{
		return;
	}

	const FString Description = Function->GetMetaData(TEXT("ToolTip"));

	if (bIsAction)
	{
		FDruidsMCPBridgeSageActionDefinition Action;
		Action.FunctionName = Function->GetFName();
		Action.Description = Description;

		for (TFieldIterator<FProperty> It(Function); It && (It->PropertyFlags & CPF_Parm); ++It)
		{
			FProperty* Prop = *It;
			if (Prop->HasAnyPropertyFlags(CPF_ReturnParm | CPF_OutParm))
			{
				continue;
			}

			FDruidsMCPBridgeSageParameterDefinition Param;
			Param.Name = Prop->GetFName();
			Param.ParameterType = FDruidsMCPBridgeSageParameterDefinition::GetParameterTypeForProperty(Prop);
			if (Param.ParameterType == EDruidsMCPBridgeSageExtensionParameterType::Array)
			{
				Param.ArrayElementType = GetArrayInnerType(CastField<FArrayProperty>(Prop));
			}
			else if (Param.ParameterType == EDruidsMCPBridgeSageExtensionParameterType::CustomStruct)
			{
				if (const FStructProperty* StructProp = CastField<FStructProperty>(Prop))
				{
					Param.StructElementType = StructProp->Struct;
				}
			}

			if (const FString* ExistingDesc = ExistingDescriptions.Find(FString::Printf(TEXT("Action_%s_%s"), *Action.FunctionName.ToString(), *Param.Name.ToString())))
			{
				Param.Description = *ExistingDesc;
			}
			Param.RefreshJsonString(Prop);
			Action.ParameterDefinitions.Add(Param);
		}

		Definition.Actions.Add(Action);
	}

	if (bIsQuery)
	{
		FDruidsMCPBridgeSageQueryDefinition Query;
		Query.FunctionName = Function->GetFName();
		Query.Description = Description;

		for (TFieldIterator<FProperty> It(Function); It && (It->PropertyFlags & CPF_Parm); ++It)
		{
			FProperty* Prop = *It;
			if (Prop->HasAnyPropertyFlags(CPF_ReturnParm))
			{
				continue;
			}

			FDruidsMCPBridgeSageParameterDefinition Param;
			Param.Name = Prop->GetFName();
			Param.ParameterType = FDruidsMCPBridgeSageParameterDefinition::GetParameterTypeForProperty(Prop);
			if (Param.ParameterType == EDruidsMCPBridgeSageExtensionParameterType::Array)
			{
				Param.ArrayElementType = GetArrayInnerType(CastField<FArrayProperty>(Prop));
			}
			else if (Param.ParameterType == EDruidsMCPBridgeSageExtensionParameterType::CustomStruct)
			{
				if (const FStructProperty* StructProp = CastField<FStructProperty>(Prop))
				{
					Param.StructElementType = StructProp->Struct;
				}
			}

			if (Prop->HasAnyPropertyFlags(CPF_OutParm))
			{
				if (const FString* ExistingDesc = ExistingDescriptions.Find(FString::Printf(TEXT("Query_Result_%s_%s"), *Query.FunctionName.ToString(), *Param.Name.ToString())))
				{
					Param.Description = *ExistingDesc;
				}
				Param.RefreshJsonString(Prop);
				Query.ResultDefinitions.Add(Param);
			}
			else
			{
				if (const FString* ExistingDesc = ExistingDescriptions.Find(FString::Printf(TEXT("Query_Param_%s_%s"), *Query.FunctionName.ToString(), *Param.Name.ToString())))
				{
					Param.Description = *ExistingDesc;
				}
				Param.RefreshJsonString(Prop);
				Query.ParameterDefinitions.Add(Param);
			}
		}

		Definition.Queries.Add(Query);
	}
}

FDruidsMCPBridgeSageActionResult UDruidsMCPBridgeSageExtension::ExecuteAction(const TSharedPtr<FJsonObject>& ActionRequestDetails)
{
	ClearExtensionError();
	ClearExtensionResults();

	if (!ActionRequestDetails.IsValid())
	{
		return FDruidsMCPBridgeSageActionResult::NoAction();
	}

	FString ActionName;
	if (!ActionRequestDetails->TryGetStringField(TEXT("action_name"), ActionName))
	{
		return FDruidsMCPBridgeSageActionResult::ExtensionFailure(TEXT("Action request is missing action_name."));
	}

	FName FunctionName;
	if (!ResolveFunctionName(ActionName, true, FunctionName))
	{
		return FDruidsMCPBridgeSageActionResult::ExtensionFailure(FString::Printf(TEXT("Action not found: %s"), *ActionName));
	}

	UFunction* Function = GetClass()->FindFunctionByName(FunctionName);
	if (!Function)
	{
		return FDruidsMCPBridgeSageActionResult::ExtensionFailure(TEXT("Action function not found."));
	}

	TSharedPtr<FJsonObject> Parameters = MakeShared<FJsonObject>();
	const TSharedPtr<FJsonObject>* ParametersObject = nullptr;
	if (ActionRequestDetails->TryGetObjectField(TEXT("parameters"), ParametersObject) && ParametersObject && ParametersObject->IsValid())
	{
		Parameters = *ParametersObject;
	}

	TArray<uint8> ParamBuffer;
	if (!BuildParamsFromJson(Function, Parameters, ParamBuffer))
	{
		return FDruidsMCPBridgeSageActionResult::ExtensionFailure(TEXT("Could not marshal action parameters from JSON."));
	}

	{
		const FScopedTransaction Transaction(FText::Format(NSLOCTEXT("DruidsMCPBridge", "ExecuteSageAction", "Execute Sage Action: {0}"), FText::FromString(ActionName)));
		const FEditorScriptExecutionGuard ScriptGuard;
		ProcessEvent(Function, ParamBuffer.GetData());
	}

	for (TFieldIterator<FProperty> It(Function); It && (It->PropertyFlags & CPF_Parm); ++It)
	{
		It->DestroyValue_InContainer(ParamBuffer.GetData());
	}

	if (HasExtensionError())
	{
		return FDruidsMCPBridgeSageActionResult::ExtensionFailure(GetExtensionError());
	}
	if (HasExtensionResults())
	{
		return FDruidsMCPBridgeSageActionResult::ExtensionSuccess(GetExtensionResults());
	}

	return FDruidsMCPBridgeSageActionResult::ExtensionSuccess();
}

TSharedPtr<FJsonValue> UDruidsMCPBridgeSageExtension::ExecuteQuery(const TSharedPtr<FJsonObject>& QueryRequestDetails)
{
	if (!QueryRequestDetails.IsValid())
	{
		TSharedPtr<FJsonObject> Error = MakeShared<FJsonObject>();
		Error->SetStringField(TEXT("error"), TEXT("Query request invalid."));
		return MakeShared<FJsonValueObject>(Error);
	}

	FString QueryName;
	if (!QueryRequestDetails->TryGetStringField(TEXT("query_name"), QueryName))
	{
		TSharedPtr<FJsonObject> Error = MakeShared<FJsonObject>();
		Error->SetStringField(TEXT("error"), TEXT("Query request is missing query_name."));
		return MakeShared<FJsonValueObject>(Error);
	}

	FName FunctionName;
	if (!ResolveFunctionName(QueryName, false, FunctionName))
	{
		TSharedPtr<FJsonObject> Error = MakeShared<FJsonObject>();
		Error->SetStringField(TEXT("error"), FString::Printf(TEXT("Query not found: %s"), *QueryName));
		return MakeShared<FJsonValueObject>(Error);
	}

	UFunction* Function = GetClass()->FindFunctionByName(FunctionName);
	if (!Function)
	{
		TSharedPtr<FJsonObject> Error = MakeShared<FJsonObject>();
		Error->SetStringField(TEXT("error"), TEXT("Query function not found."));
		return MakeShared<FJsonValueObject>(Error);
	}

	TSharedPtr<FJsonObject> Parameters = MakeShared<FJsonObject>();
	const TSharedPtr<FJsonObject>* ParametersObject = nullptr;
	if (QueryRequestDetails->TryGetObjectField(TEXT("parameters"), ParametersObject) && ParametersObject && ParametersObject->IsValid())
	{
		Parameters = *ParametersObject;
	}

	TArray<uint8> ParamBuffer;
	if (!BuildParamsFromJson(Function, Parameters, ParamBuffer))
	{
		TSharedPtr<FJsonObject> Error = MakeShared<FJsonObject>();
		Error->SetStringField(TEXT("error"), TEXT("Could not marshal query parameters from JSON."));
		return MakeShared<FJsonValueObject>(Error);
	}

	const FEditorScriptExecutionGuard ScriptGuard;
	ProcessEvent(Function, ParamBuffer.GetData());

	TSharedPtr<FJsonObject> Results = MakeShared<FJsonObject>();
	for (TFieldIterator<FProperty> It(Function); It && (It->PropertyFlags & CPF_Parm); ++It)
	{
		FProperty* Prop = *It;
		if (!Prop->HasAnyPropertyFlags(CPF_OutParm) || Prop->HasAnyPropertyFlags(CPF_ReturnParm))
		{
			continue;
		}

		void* ValuePtr = Prop->ContainerPtrToValuePtr<void>(ParamBuffer.GetData());
		TSharedPtr<FJsonValue> JsonValue = GetJsonValueFromProperty(Prop, ValuePtr);
		if (JsonValue.IsValid())
		{
			Results->SetField(Prop->GetName(), JsonValue);
		}
	}

	for (TFieldIterator<FProperty> It(Function); It && (It->PropertyFlags & CPF_Parm); ++It)
	{
		It->DestroyValue_InContainer(ParamBuffer.GetData());
	}

	return MakeShared<FJsonValueObject>(Results);
}

bool UDruidsMCPBridgeSageExtension::ResolveFunctionName(const FString& RequestedName, bool bAction, FName& OutFunctionName) const
{
	const TArray<FDruidsMCPBridgeSageActionDefinition>& Actions = Definition.Actions;
	const TArray<FDruidsMCPBridgeSageQueryDefinition>& Queries = Definition.Queries;

	if (bAction)
	{
		for (const FDruidsMCPBridgeSageActionDefinition& Action : Actions)
		{
			if (Action.FunctionName.ToString().Equals(RequestedName, ESearchCase::IgnoreCase))
			{
				OutFunctionName = Action.FunctionName;
				return true;
			}
		}
	}
	else
	{
		for (const FDruidsMCPBridgeSageQueryDefinition& Query : Queries)
		{
			if (Query.FunctionName.ToString().Equals(RequestedName, ESearchCase::IgnoreCase))
			{
				OutFunctionName = Query.FunctionName;
				return true;
			}
		}
	}

	for (TFieldIterator<UFunction> FuncIt(GetClass(), EFieldIteratorFlags::IncludeSuper); FuncIt; ++FuncIt)
	{
		if ((*FuncIt)->GetFName().ToString().Equals(RequestedName, ESearchCase::IgnoreCase))
		{
			OutFunctionName = (*FuncIt)->GetFName();
			return true;
		}
	}

	return false;
}

bool UDruidsMCPBridgeSageExtension::BuildParamsFromJson(const UFunction* Function, const TSharedPtr<const FJsonObject>& JsonObject, TArray<uint8>& OutParamBuffer)
{
	if (!Function || !JsonObject.IsValid())
	{
		return false;
	}

	OutParamBuffer.SetNumZeroed(Function->ParmsSize);
	void* Params = OutParamBuffer.GetData();

	for (TFieldIterator<FProperty> It(Function); It && (It->PropertyFlags & CPF_Parm); ++It)
	{
		FProperty* Prop = *It;
		Prop->InitializeValue_InContainer(Params);

		if (Prop->HasAnyPropertyFlags(CPF_ReturnParm | CPF_OutParm))
		{
			continue;
		}

		const FString Key = Prop->GetAuthoredName();
#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION >= 8
		const TSharedPtr<FJsonValue>* JsonValue = JsonObject->Values.Find(FJsonObject::FStringType(Key));
#else
		const TSharedPtr<FJsonValue>* JsonValue = JsonObject->Values.Find(Key);
#endif
		if (!JsonValue)
		{
#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION >= 8
			JsonValue = JsonObject->Values.Find(FJsonObject::FStringType(Prop->GetName()));
#else
			JsonValue = JsonObject->Values.Find(Prop->GetName());
#endif
		}
		if (!JsonValue || !JsonValue->IsValid())
		{
			continue;
		}

		void* DestValuePtr = Prop->ContainerPtrToValuePtr<void>(Params);
		if (const FStructProperty* StructProp = CastField<FStructProperty>(Prop))
		{
			TSharedPtr<FJsonObject> StructObject = (*JsonValue)->AsObject();
			if (!StructObject.IsValid())
			{
				TSharedPtr<FJsonObject> ParsedObject;
				const TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create((*JsonValue)->AsString());
				if (FJsonSerializer::Deserialize(Reader, ParsedObject))
				{
					StructObject = ParsedObject;
				}
			}

			if (!StructObject.IsValid() || !FJsonObjectConverter::JsonObjectToUStruct(StructObject.ToSharedRef(), StructProp->Struct, DestValuePtr))
			{
				return false;
			}
		}
		else if (!FJsonObjectConverter::JsonValueToUProperty(*JsonValue, Prop, DestValuePtr))
		{
			return false;
		}
	}

	return true;
}

TSharedPtr<FJsonValue> UDruidsMCPBridgeSageExtension::GetJsonValueFromProperty(FProperty* Property, void* ValuePtr)
{
	if (!Property || !ValuePtr)
	{
		return nullptr;
	}

	TSharedPtr<FJsonValue> JsonValue = FJsonObjectConverter::UPropertyToJsonValue(Property, ValuePtr);
	if (JsonValue.IsValid())
	{
		return JsonValue;
	}

	if (const FNameProperty* NameProp = CastField<FNameProperty>(Property))
	{
		return MakeShared<FJsonValueString>(NameProp->GetPropertyValue(ValuePtr).ToString());
	}
#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION >= 8
    else if (const FTextProperty* TextProp = CastField<FTextProperty>(Property))
#else
    if (const FTextProperty* TextProp = CastField<FTextProperty>(Property))
#endif
	{
		return MakeShared<FJsonValueString>(TextProp->GetPropertyValue(ValuePtr).ToString());
	}

	return nullptr;
}

EDruidsMCPBridgeSageExtensionParameterType UDruidsMCPBridgeSageExtension::GetArrayInnerType(const FArrayProperty* ArrayProperty)
{
	return ArrayProperty
		? FDruidsMCPBridgeSageParameterDefinition::GetParameterTypeForProperty(ArrayProperty->Inner)
		: EDruidsMCPBridgeSageExtensionParameterType::None;
}

UWorld* UDruidsMCPBridgeSageExtension::GetWorldEditorContext()
{
	return GEditor ? GEditor->GetEditorWorldContext().World() : nullptr;
}

UObject* UDruidsMCPBridgeSageExtension::GetCurrentlyEditedAsset()
{
	if (!GEditor)
	{
		return nullptr;
	}

	if (UAssetEditorSubsystem* AssetEditorSubsystem = GEditor->GetEditorSubsystem<UAssetEditorSubsystem>())
	{
		TArray<UObject*> EditedAssets = AssetEditorSubsystem->GetAllEditedAssets();
		return EditedAssets.Num() > 0 ? EditedAssets[0] : nullptr;
	}

	return nullptr;
}

FString& UDruidsMCPBridgeSageExtension::GetCurrentExtensionError()
{
	static thread_local FString CurrentExtensionError;
	return CurrentExtensionError;
}

FString& UDruidsMCPBridgeSageExtension::GetCurrentExtensionResults()
{
	static thread_local FString CurrentExtensionResults;
	return CurrentExtensionResults;
}

void UDruidsMCPBridgeSageExtension::ReportExtensionError(const FString& ErrorMessage)
{
	GetCurrentExtensionError() = ErrorMessage;
	UE_LOG(LogDruidsMCPBridge, Warning, TEXT("Sage extension error: %s"), *ErrorMessage);
}

void UDruidsMCPBridgeSageExtension::ReportExtensionResults(const FString& ResultsMessage)
{
	GetCurrentExtensionResults() = ResultsMessage;
	UE_LOG(LogDruidsMCPBridge, Display, TEXT("Sage extension result: %s"), *ResultsMessage);
}

void UDruidsMCPBridgeSageExtension::ClearExtensionError()
{
	GetCurrentExtensionError().Reset();
}

FString UDruidsMCPBridgeSageExtension::GetExtensionError()
{
	return GetCurrentExtensionError();
}

bool UDruidsMCPBridgeSageExtension::HasExtensionError()
{
	return !GetCurrentExtensionError().IsEmpty();
}

void UDruidsMCPBridgeSageExtension::ClearExtensionResults()
{
	GetCurrentExtensionResults().Reset();
}

FString UDruidsMCPBridgeSageExtension::GetExtensionResults()
{
	return GetCurrentExtensionResults();
}

bool UDruidsMCPBridgeSageExtension::HasExtensionResults()
{
	return !GetCurrentExtensionResults().IsEmpty();
}
