material.h

// Copyright 2024 The Khronos Group
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include "material/textureLoader.h"

// pxr
#include <pxr/base/tf/token.h>
#include <pxr/imaging/hd/enums.h>
#include <pxr/imaging/hd/instancer.h>
#include <pxr/imaging/hd/material.h>
#include <pxr/imaging/hd/materialNetwork2Interface.h>
#include <pxr/imaging/hd/mesh.h>
#include <pxr/imaging/hd/renderDelegate.h>
#include <pxr/imaging/hd/sceneDelegate.h>
#include <pxr/usd/sdf/path.h>

// std
#include <anari/anari_cpp.hpp>
#include <map>

PXR_NAMESPACE_OPEN_SCOPE

struct HdAnariMaterial : public HdMaterial
{
 public:
  using PrimvarBinding =
      std::map<TfToken, TfToken>; // Maps a primvar to an anari attribute name
  using PrimvarMapping =
      std::map<SdfPath, TfToken>; // Maps an input to a primvar name
  using TextureDescMapping = std::map<SdfPath,
      HdAnariTextureLoader::TextureDesc>; // Maps a path to a usduvtexture and
                                          // related info

  using SamplerMapping = std::map<SdfPath,
      anari::Sampler>; // Maps a usduvtexture path to the matching sampler

  HdAnariMaterial(anari::Device d, const SdfPath &id);
  ~HdAnariMaterial() override;

  // Get a material instance.
  anari::Material GetAnariMaterial() const;

  // Create thet transient materialnetwork2 representation. To be used by the
  // above ProcessMaterialNetwork call.
  void Sync(HdSceneDelegate *sceneDelegate,
      HdRenderParam *renderParam,
      HdDirtyBits *dirtyBits) override;

  // Finalize
  void Finalize(HdRenderParam *renderParam) override;

  // Get initial dirty bits mask
  HdDirtyBits GetInitialDirtyBitsMask() const override;

  const PrimvarBinding &GetPrimvarBinding() const
  {
    return attributes_;
  }

 protected:
  // Load and returns all textures associated with this material.
  virtual TextureDescMapping EnumerateTextures(
      const HdMaterialNetwork2Interface &materialNetworkInterface,
      TfToken terminal) const = 0;
  virtual PrimvarMapping EnumeratePrimvars(
      const HdMaterialNetwork2Interface &materialNetworkInterface,
      TfToken terminal) const = 0;

  virtual anari::Material GetOrCreateMaterial(
      const HdMaterialNetwork2Interface &materialNetworkIface,
      const PrimvarBinding &primvarBinding,
      const PrimvarMapping &primvarMapping,
      const SamplerMapping &samplerMapping) const = 0;

  TextureDescMapping textures_;
  SamplerMapping samplers_;
  PrimvarMapping primvars_;
  PrimvarBinding attributes_;

  HdMaterialNetwork2 materialNetwork2_;
  anari::Device device_;
  anari::Material material_;

 private:
  // Convert the given material network to the newer HdMaterialNetwork2
  static HdMaterialNetwork2 convertToHdMaterialNetwork2(
      const HdMaterialNetworkMap &hdNetworkMap);

  // Return a primvar to anari parameter mapping that can be used to correctly
  // expose primvars on the meshes.
  static PrimvarBinding BuildPrimvarBinding(const PrimvarMapping &primvarNames);

  static SamplerMapping CreateSamplers(
      anari::Device device, const TextureDescMapping &textureDescs);
  static void ReleaseSamplers(anari::Device, const SamplerMapping &samplers);
};

PXR_NAMESPACE_CLOSE_SCOPE