// Copyright 2020-2024 CesiumGS, Inc. and Contributors

#pragma once

#include "CoreMinimal.h"

#if WITH_EDITOR
#include "Materials/MaterialFunctionMaterialLayer.h"
#endif

#include <memory>
#include <variant>

#if WITH_EDITOR
enum ECustomMaterialOutputType : int;

enum class ECesiumEncodedMetadataType : uint8;
enum class ECesiumEncodedMetadataComponentType : uint8;

class UMaterialExpressionParameter;
#endif

namespace GenerateMaterialUtility {
#if WITH_EDITOR
template <typename ObjClass>
static FORCEINLINE ObjClass* LoadObjFromPath(const FName& Path) {
  if (Path == NAME_None)
    return nullptr;

  return Cast<ObjClass>(
      StaticLoadObject(ObjClass::StaticClass(), nullptr, *Path.ToString()));
}

/**
 * @brief The message used for the description of autogenerated nodes.
 */
static const FString AutogeneratedMessage = "AUTOGENERATED DO NOT EDIT";

/**
 * @brief An increment constant that is used to space out the autogenerated
 * nodes.
 */
static const int32 Incr = 200;

/**
 * @brief Classification of nodes into categories, such that they can be handled
 * separately.
 */
struct MaterialNodeClassification {
  /** Nodes that were procedurally created with the AutogeneratedMessage. */
  TArray<UMaterialExpression*> AutoGeneratedNodes;
  /** Nodes without the AutogeneratedMessage (presumably added by the user). */
  TArray<UMaterialExpression*> UserAddedNodes;
};

/**
 * @brief A container that bundles various components used to procedurally
 * create or modify a material.
 */
struct MaterialGenerationState {
  /**
   * Nodes that are procedurally created and to be marked with
   * AutogeneratedMessage.
   */
  TArray<UMaterialExpression*> AutoGeneratedNodes;
  /**
   * Nodes that should only exist once in the material layer, e.g.,
   * SetMaterialAttributes. These will not be regenerated if already present.
   */
  TArray<UMaterialExpression*> OneTimeGeneratedNodes;
  /**
   * A map from autogenerated node names to the FExpressionInputs of user-added
   * nodes. These are the nodes to which they previously sent their outputs.
   */
  TMap<FString, TArray<FExpressionInput*>> UserConnectionOutputMap;
  /**
   * The details of an input connection from an autogenerated node. When nodes
   * are regenerated, this is used to reconstruct the connection.
   */
  struct AutogeneratedInput {
    FString nodeName;
    FString outputName;
  };
  /**
   * Describes a connection to a node placed by the user or to an
   * otherwise user-added connection between two autogenerated nodes (e.g.,
   * property values that were linked to the "If" that handles null feature
   * IDs).
   */
  typedef std::variant<const FExpressionInput*, AutogeneratedInput>
      UserConnectionInput;
  /**
   * A map between the names of parameters on a node and their user-added
   * connection details.
   */
  typedef TMap<FString, UserConnectionInput> ParameterConnections;
  /**
   * A map from autogenerated node names to a collection of parameters from
   * which they previously took inputs.
   */
  TMap<FString, ParameterConnections> UserConnectionInputMap;
};

/**
 * Creates a new material layer asset with the given name. If the specified
 * package doesn't exist, this function creates it.
 */
UMaterialFunctionMaterialLayer*
CreateMaterialLayer(const FString& PackageName, const FString& MaterialName);

/**
 * Moves the generated nodes from the material state into the given material
 * layer.
 */
void MoveNodesToMaterialLayer(
    MaterialGenerationState& MaterialState,
    UMaterialFunctionMaterialLayer* pMaterialLayer);

/**
 * Computes a scalar for spacing out material nodes. The actual computation is
 * rather arbitrary, but this prevents clumping when properties have extremely
 * long names.
 */
float GetNameLengthScalar(const FName& Name);
float GetNameLengthScalar(const FString& Name);

ECustomMaterialOutputType
GetOutputTypeForEncodedType(ECesiumEncodedMetadataType Type);

FString GetHlslTypeForEncodedType(
    ECesiumEncodedMetadataType Type,
    ECesiumEncodedMetadataComponentType ComponentType);

FString GetSwizzleForEncodedType(ECesiumEncodedMetadataType Type);

/**
 * @brief Generates a parameter node corresponding to the given encoded
 * metadata type.
 */
UMaterialExpressionParameter* GenerateParameterNode(
    UMaterialFunctionMaterialLayer* TargetMaterialLayer,
    const ECesiumEncodedMetadataType Type,
    const FString& Name,
    int32 NodeX,
    int32 NodeY);
#endif

} // namespace GenerateMaterialUtility