#include "MCPMaterialLibrary.h"

#include "Dom/JsonObject.h"
#include "Dom/JsonValue.h"
#include "Materials/Material.h"
#include "Materials/MaterialExpression.h"
#include "Materials/MaterialFunction.h"
#include "Serialization/JsonSerializer.h"
#include "Serialization/JsonWriter.h"

namespace
{
FString JsonStringifyMaterial(const TSharedRef<FJsonObject>& Object)
{
    FString Out;
    const TSharedRef<TJsonWriter<>> Writer = TJsonWriterFactory<>::Create(&Out);
    FJsonSerializer::Serialize(Object, Writer);
    return Out;
}

FString ExpressionId(const UMaterialExpression* Expression)
{
    if (!Expression)
    {
        return FString();
    }

#if WITH_EDITORONLY_DATA
    if (Expression->MaterialExpressionGuid.IsValid())
    {
        return Expression->MaterialExpressionGuid.ToString(EGuidFormats::DigitsWithHyphensInBraces);
    }
#endif

    return Expression->GetPathName();
}

FString OutputNameForIndex(UMaterialExpression* Expression, int32 OutputIndex)
{
    if (!Expression || OutputIndex < 0)
    {
        return FString();
    }

    TArray<FExpressionOutput>& Outputs = Expression->GetOutputs();
    if (!Outputs.IsValidIndex(OutputIndex))
    {
        return FString();
    }

    return Outputs[OutputIndex].OutputName.ToString();
}

int32 ResolveOutputIndex(UMaterialExpression* Expression, const FString& OutputNameOrIndex)
{
    if (!Expression)
    {
        return INDEX_NONE;
    }

    FString Needle = OutputNameOrIndex;
    Needle.TrimStartAndEndInline();
    if (Needle.IsEmpty() || Needle.Equals(TEXT("None"), ESearchCase::IgnoreCase))
    {
        return 0;
    }

    int32 NumericIndex = INDEX_NONE;
    if (LexTryParseString(NumericIndex, *Needle))
    {
        return Expression->GetOutputs().IsValidIndex(NumericIndex) ? NumericIndex : INDEX_NONE;
    }

    TArray<FExpressionOutput>& Outputs = Expression->GetOutputs();
    for (int32 OutputIndex = 0; OutputIndex < Outputs.Num(); ++OutputIndex)
    {
        if (Outputs[OutputIndex].OutputName.ToString().Equals(Needle, ESearchCase::IgnoreCase))
        {
            return OutputIndex;
        }
    }

    return Outputs.Num() == 1 ? 0 : INDEX_NONE;
}

int32 ResolveInputIndex(UMaterialExpression* Expression, const FString& InputNameOrIndex)
{
    if (!Expression)
    {
        return INDEX_NONE;
    }

    FString Needle = InputNameOrIndex;
    Needle.TrimStartAndEndInline();
    if (Needle.IsEmpty() || Needle.Equals(TEXT("None"), ESearchCase::IgnoreCase))
    {
        return Expression->CountInputs() == 1 ? 0 : INDEX_NONE;
    }

    int32 NumericIndex = INDEX_NONE;
    if (LexTryParseString(NumericIndex, *Needle))
    {
        return NumericIndex >= 0 && NumericIndex < Expression->CountInputs() ? NumericIndex : INDEX_NONE;
    }

    for (int32 InputIndex = 0; InputIndex < Expression->CountInputs(); ++InputIndex)
    {
        if (Expression->GetInputName(InputIndex).ToString().Equals(Needle, ESearchCase::IgnoreCase))
        {
            return InputIndex;
        }
    }

    return INDEX_NONE;
}

TSharedRef<FJsonObject> MakeStatusJson(bool bOk, const FString& Error = FString())
{
    TSharedRef<FJsonObject> Root = MakeShared<FJsonObject>();
    Root->SetBoolField(TEXT("ok"), bOk);
    if (!Error.IsEmpty())
    {
        Root->SetStringField(TEXT("error"), Error);
    }
    return Root;
}

void AppendExpressionJson(UMaterialExpression* Expression, TArray<TSharedPtr<FJsonValue>>& Nodes, TArray<TSharedPtr<FJsonValue>>& Edges)
{
    if (!Expression)
    {
        return;
    }

    TSharedRef<FJsonObject> Node = MakeShared<FJsonObject>();
    Node->SetStringField(TEXT("guid"), ExpressionId(Expression));
    Node->SetStringField(TEXT("name"), Expression->GetName());
    Node->SetStringField(TEXT("class"), Expression->GetClass()->GetPathName());
#if WITH_EDITORONLY_DATA
    Node->SetNumberField(TEXT("x"), Expression->MaterialExpressionEditorX);
    Node->SetNumberField(TEXT("y"), Expression->MaterialExpressionEditorY);
    if (!Expression->Desc.IsEmpty())
    {
        Node->SetStringField(TEXT("desc"), Expression->Desc);
    }
#endif

    TArray<TSharedPtr<FJsonValue>> Inputs;
    const int32 InputCount = Expression->CountInputs();
    for (int32 InputIndex = 0; InputIndex < InputCount; ++InputIndex)
    {
        const FExpressionInput* Input = Expression->GetInput(InputIndex);
        const FName InputName = Expression->GetInputName(InputIndex);

        TSharedRef<FJsonObject> InputJson = MakeShared<FJsonObject>();
        InputJson->SetStringField(TEXT("name"), InputName.ToString());
        InputJson->SetNumberField(TEXT("index"), InputIndex);
        InputJson->SetBoolField(TEXT("connected"), Input && Input->Expression);
        Inputs.Add(MakeShared<FJsonValueObject>(InputJson));

        if (Input && Input->Expression)
        {
            TSharedRef<FJsonObject> Edge = MakeShared<FJsonObject>();
            TArray<TSharedPtr<FJsonValue>> From;
            From.Add(MakeShared<FJsonValueString>(ExpressionId(Input->Expression)));
            From.Add(MakeShared<FJsonValueString>(OutputNameForIndex(Input->Expression, Input->OutputIndex)));
            TArray<TSharedPtr<FJsonValue>> To;
            To.Add(MakeShared<FJsonValueString>(ExpressionId(Expression)));
            To.Add(MakeShared<FJsonValueString>(InputName.ToString()));
            Edge->SetArrayField(TEXT("from"), From);
            Edge->SetArrayField(TEXT("to"), To);
            Edges.Add(MakeShared<FJsonValueObject>(Edge));
        }
    }
    Node->SetArrayField(TEXT("inputs"), Inputs);

    TArray<TSharedPtr<FJsonValue>> OutputsJson;
    TArray<FExpressionOutput>& Outputs = Expression->GetOutputs();
    for (int32 OutputIndex = 0; OutputIndex < Outputs.Num(); ++OutputIndex)
    {
        TSharedRef<FJsonObject> OutputJson = MakeShared<FJsonObject>();
        OutputJson->SetStringField(TEXT("name"), Outputs[OutputIndex].OutputName.ToString());
        OutputJson->SetNumberField(TEXT("index"), OutputIndex);
        OutputsJson.Add(MakeShared<FJsonValueObject>(OutputJson));
    }
    Node->SetArrayField(TEXT("outputs"), OutputsJson);

    Nodes.Add(MakeShared<FJsonValueObject>(Node));
}
}

TArray<UMaterialExpression*> UMCPMaterialLibrary::GetMaterialExpressions(UObject* MaterialOrFunction)
{
    TArray<UMaterialExpression*> Out;
    if (UMaterial* Material = Cast<UMaterial>(MaterialOrFunction))
    {
        for (TObjectPtr<UMaterialExpression> Expression : Material->GetExpressionCollection().Expressions)
        {
            if (Expression)
            {
                Out.Add(Expression.Get());
            }
        }
    }
    else if (UMaterialFunction* Function = Cast<UMaterialFunction>(MaterialOrFunction))
    {
        for (TObjectPtr<UMaterialExpression> Expression : Function->GetExpressionCollection().Expressions)
        {
            if (Expression)
            {
                Out.Add(Expression.Get());
            }
        }
    }
    return Out;
}

UMaterialExpression* UMCPMaterialLibrary::FindMaterialExpression(UObject* MaterialOrFunction, const FString& NodeId)
{
    for (UMaterialExpression* Expression : GetMaterialExpressions(MaterialOrFunction))
    {
        if (!Expression)
        {
            continue;
        }

        if (ExpressionId(Expression) == NodeId || Expression->GetName() == NodeId || Expression->GetPathName() == NodeId)
        {
            return Expression;
        }
    }
    return nullptr;
}

FString UMCPMaterialLibrary::DumpMaterialGraphToJson(UObject* MaterialOrFunction)
{
    TSharedRef<FJsonObject> Root = MakeShared<FJsonObject>();
    if (!MaterialOrFunction)
    {
        Root->SetBoolField(TEXT("ok"), false);
        Root->SetStringField(TEXT("error"), TEXT("null material/function"));
        return JsonStringifyMaterial(Root);
    }

    Root->SetBoolField(TEXT("ok"), true);
    Root->SetStringField(TEXT("asset"), MaterialOrFunction->GetPathName());
    Root->SetStringField(TEXT("class"), MaterialOrFunction->GetClass()->GetPathName());

    TArray<TSharedPtr<FJsonValue>> Nodes;
    TArray<TSharedPtr<FJsonValue>> Edges;
    for (UMaterialExpression* Expression : GetMaterialExpressions(MaterialOrFunction))
    {
        AppendExpressionJson(Expression, Nodes, Edges);
    }
    Root->SetArrayField(TEXT("nodes"), Nodes);
    Root->SetArrayField(TEXT("edges"), Edges);

    if (UMaterial* Material = Cast<UMaterial>(MaterialOrFunction))
    {
        TSharedRef<FJsonObject> Outputs = MakeShared<FJsonObject>();
        for (int32 PropertyIndex = 0; PropertyIndex < static_cast<int32>(MP_MAX); ++PropertyIndex)
        {
            const EMaterialProperty Property = static_cast<EMaterialProperty>(PropertyIndex);
            FExpressionInput* Input = Material->GetExpressionInputForProperty(Property);
            if (!Input || !Input->Expression)
            {
                continue;
            }

            TSharedRef<FJsonObject> Ref = MakeShared<FJsonObject>();
            TArray<TSharedPtr<FJsonValue>> From;
            From.Add(MakeShared<FJsonValueString>(ExpressionId(Input->Expression)));
            From.Add(MakeShared<FJsonValueString>(OutputNameForIndex(Input->Expression, Input->OutputIndex)));
            Ref->SetArrayField(TEXT("from"), From);
            Outputs->SetObjectField(UEnum::GetValueAsString(Property), Ref);
        }
        Root->SetObjectField(TEXT("outputs"), Outputs);
    }

    return JsonStringifyMaterial(Root);
}

FString UMCPMaterialLibrary::ConnectMaterialExpressionsRaw(UObject* MaterialOrFunction, UMaterialExpression* FromExpression, const FString& FromOutput, UMaterialExpression* ToExpression, const FString& ToInput)
{
    if (!MaterialOrFunction)
    {
        return JsonStringifyMaterial(MakeStatusJson(false, TEXT("null material/function")));
    }
    if (!FromExpression)
    {
        return JsonStringifyMaterial(MakeStatusJson(false, TEXT("null source expression")));
    }
    if (!ToExpression)
    {
        return JsonStringifyMaterial(MakeStatusJson(false, TEXT("null target expression")));
    }

    const int32 OutputIndex = ResolveOutputIndex(FromExpression, FromOutput);
    if (OutputIndex == INDEX_NONE)
    {
        return JsonStringifyMaterial(MakeStatusJson(false, FString::Printf(TEXT("source output not found: %s.%s"), *FromExpression->GetName(), *FromOutput)));
    }

    const int32 InputIndex = ResolveInputIndex(ToExpression, ToInput);
    if (InputIndex == INDEX_NONE)
    {
        return JsonStringifyMaterial(MakeStatusJson(false, FString::Printf(TEXT("target input not found: %s.%s"), *ToExpression->GetName(), *ToInput)));
    }

    FExpressionInput* Input = ToExpression->GetInput(InputIndex);
    if (!Input)
    {
        return JsonStringifyMaterial(MakeStatusJson(false, FString::Printf(TEXT("target input is null: %s.%s"), *ToExpression->GetName(), *ToInput)));
    }

    Input->Expression = FromExpression;
    Input->OutputIndex = OutputIndex;

    if (UMaterial* Material = Cast<UMaterial>(MaterialOrFunction))
    {
        Material->PreEditChange(nullptr);
        Material->PostEditChange();
        Material->MarkPackageDirty();
    }
    else if (UMaterialFunction* Function = Cast<UMaterialFunction>(MaterialOrFunction))
    {
        Function->PreEditChange(nullptr);
        Function->PostEditChange();
        Function->MarkPackageDirty();
    }

    TSharedRef<FJsonObject> Root = MakeStatusJson(true);
    Root->SetStringField(TEXT("from"), FromExpression->GetName());
    Root->SetStringField(TEXT("from_output"), OutputNameForIndex(FromExpression, OutputIndex));
    Root->SetStringField(TEXT("to"), ToExpression->GetName());
    Root->SetStringField(TEXT("to_input"), ToExpression->GetInputName(InputIndex).ToString());
    return JsonStringifyMaterial(Root);
}