// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause

#ifndef vtkWebGPUComputeFrustumCuller_h
#define vtkWebGPUComputeFrustumCuller_h

#include "vtkCallbackCommand.h" // for the bounds/camera recomputed callback
#include "vtkCuller.h"
#include "vtkNew.h"                   // for new macro
#include "vtkRenderingWebGPUModule.h" // For export macro
#include "vtkSmartPointer.h"          // for the pipeline smart pointer
#include "vtkWebGPUComputePipeline.h" // for the member compute pipeline

VTK_ABI_NAMESPACE_BEGIN

/**
 * This culler culls props to the camera view frustum using WebGPU compute shaders.
 *
 * To use this culler, simply instantiate it:
 *
 * vtkNew<vtkWebGPUComputeFrustumCuller> webgpuFrustumCuller;
 *
 * and add it to the cullers of your renderer. Note that by default, the renderer contains a
 * vtkFrustumCoverageCuller. You probably want to remove it from the renderer before adding the GPU
 * compute frustum culler as they are pretty much redundant with each other
 *
 * // Removing the default culler
 * renderer->GetCullers()->RemoveAllItems();
 *
 * // Adding the WebGPU compute shader frustum culler
 * renderer->GetCullers()->AddItem(webgpuFrustumCuller);
 */
class VTKRENDERINGWEBGPU_EXPORT vtkWebGPUComputeFrustumCuller : public vtkCuller
{
public:
  vtkTypeMacro(vtkWebGPUComputeFrustumCuller, vtkCuller);
  static vtkWebGPUComputeFrustumCuller* New();
  /**
   * This is called outside the render loop by vtkRenderer
   */
  virtual double Cull(vtkRenderer* ren, vtkProp** propList, int& listLength, int& initialized);

protected:
  vtkWebGPUComputeFrustumCuller();

private:
  vtkWebGPUComputeFrustumCuller(const vtkWebGPUComputeFrustumCuller&) = delete;
  void operator=(const vtkWebGPUComputeFrustumCuller&) = delete;

  /**
   * Creates the input bounds WebGPU buffer and adds it to the pipeline
   */
  void CreateInputBoundsBuffer(vtkProp** propList, int nbProps);

  /**
   * Creates the buffer that will contain the indicees of the objects that were not culled
   */
  void CreateOutputIndicesBuffer(int nbProps);

  /**
   * Creates the uniform buffer that contains the view-projection matrix data of the camera.
   * The given vector is expected to contain the data of the view-projection matrix in column major
   * order (so the matrices returned by a vtkCamera need to be transposed before going to WebGPU)
   */
  void CreateViewProjMatrixBuffer(const std::vector<float>& matrixData);

  /**
   * Reconfigures the culler so that it can handle a new number of props.
   * This encompasses the size of the WebGPU bounds buffer, the vector for the cached props
   * positions etc...
   */
  void ResizeCuller(vtkProp** propList, int propsCount);

  /**
   * Resizes the buffer that contains the bounds of the objects to be culled and uploads the new
   * bounds
   */
  void ResizeBoundsBuffer(vtkProp** propList, int newPropsCount);

  /**
   * Resizes the buffer that will contain the indices of the objects that were not culled
   */
  void ResizeOutputIndicesBuffer(vtkProp** propList, int newPropsCount);

  /**
   * Resizes the scratch list used by the OutputObjectIndicesMapCallback and fills it with the
   * adresses of the props of the propList
   */
  void ResizeScratchList(vtkProp** propList, int listLength);

  /**
   * Clears and updates the ActorToIndex map with the new props given in parameter
   */
  void UpdateActorToIndex(vtkProp** propList, int newPropsCount);

  /**
   * Adds observers to all the actor of the given that don't already have an observer
   */
  void AddBoundsObservers(vtkProp** propList, int propsCount);

  /**
   * Adds an observer to the camera of the renderer.
   */
  void AddCameraObserver(vtkRenderer* ren);

  /**
   * Bounds are lazily recomputed in VTK (only when GetBounds() is called) and in particular, they
   * are not recomputed when an actor's position changed. This means that an actor that is out of
   * the view frustum and that is culled, if moved into the view frustum will still be culled
   * because its bounds will not have been recomputed.
   */
  void TriggerBoundsRecomputation(vtkProp** propList, int listLength);

  /**
   * Reuploads the bounds of the actors that have changed since last time the culler was called
   */
  void UpdateBoundsBuffer(vtkProp** propList, int listLength);

  /**
   * Re-uploads the camera data to the GPU if necessary (if this->CameraDirty is true)
   */
  void UpdateCamera(vtkRenderer* ren);

  /**
   * When the bounds of some actors changed, we're going to have to reupload the new bounds to the
   * GPU.
   * If the number of actors' bound that need to be reuploaded is below the value returned by this
   * function, we're only going to reupload the bounds of the actors that were modified. If the
   * number of actors is above, we're going to reupload the whole buffer.
   *
   * For now, this function uses an arbitrarily fixed threshold.
   */
  static const int WHOLE_BUFFER_REUPLOAD_THRESHOLD = 10;
  int ReuploadBoundsBufferThreshold();

  /**
   * Callback function called when the bounds of an actor are recomputed. We will have to reupload
   * the bounds to the GPU and this function will prepare the list for it (this function does not do
   * the reupload itself)
   */
  static void BoundsRecomputedCallbackFunction(
    vtkObject* caller, unsigned long eid, void* clientdata, void* calldata);

  /**
   * Function called when the camera is modified. This callback sets this->CameraDirty to true
   */
  static void CameraModifiedCallbackFunction(
    vtkObject* caller, unsigned long eid, void* clientdata, void* calldata);

  /**
   * Callback that reads the number of objects that passed the culling test and that stores the
   * result in the listLength parameter of the Cull() method (passed through userdata).
   */
  static void OutputObjectCountMapCallback(const void* mappedData, void* userdata);

  /**
   * Reads the indices of the objects that passed the culling test and stores those indices at the
   * front of the vtkPropList passed in the OutputIndicesCallbackData structure passed through
   * userdata.
   */
  static void OutputObjectIndicesMapCallback(const void* mappedData, void* userdata);

  /**
   * Callback data passed to the OutputObjectIndicesMapCallback callback function.
   * indicesCount indicates how many objects passed the culling test. This value was retrieved
   * earlier by mapping the OutputObjectCountBuffer.
   * Because we're reading the props from the prop list and writing the results directly to the prop
   * list, we have a risk of overwriting the prop list before having the chance to read it. For
   * example:
   * - If the indices to copy are [0, 2, 1] and the prop list (list of addresses) is
   * [0x10, 0x20, 0x30], the final prop list is going to be [0x10, 0x30, 0x30] because the index '1'
   * at the end of the indices to copy now refers to 0x30 whereas it should have been refering to
   * 0x20.
   *
   * The scratch list helps us prevent this issue by keeping a "sane" list of props to read from
   */
  struct OutputIndicesCallbackData
  {
    // The list of props. This should be the same list as passed in parameter to the Cull() method
    vtkProp** propList;
    // How many props passed the culling test? This should be the value contained in the OutputCount
    // WebGPU buffer
    int* indicesCount;

    // The scratch list is a pointer on the pre-allocated vector stored in this compute culler
    std::vector<vtkProp*>* scratchList;
  };

  // How many props was the culler last configured for. This variable is used to determine if the
  // number of props to be culled has changed since last time, meaning that we have to recreate a
  // new bounds buffer, a new PropsBounds cache std::vector etc...
  int PreviousPropsCount = -1;

  // Compute pipeline used for the frustum culling compute shader
  vtkSmartPointer<vtkWebGPUComputePipeline> Pipeline;
  // Frustum culling compute shader pass
  vtkSmartPointer<vtkWebGPUComputePass> FrustumCullingPass;

  // Maps every actor that is cull-tested by this culler to its index in the PropsBounds buffer.
  // This map is used when an actor has had its bounds recomputed. We can then retrieve the position
  // of the actor whose bounds were modified within the whole buffer and thus only reupload its
  // bounds (and no other bounds) allowing for very fast bounds buffer updates.
  std::unordered_map<vtkActor*, int> ActorToIndex;

  // A set of the actors whose bounds will need to be reuploaded because they have changed since
  // last time the culler was called. We're
  std::vector<vtkActor*> ActorsToReupload;

  // Scratch list used by the OutputObjectIndicesMapCallback
  std::vector<vtkProp*> CallbackScratchList;

  // Index of the input bounds buffer in the compute pipeline
  int InputBoundsBufferIndex = -1;
  // Index of the buffer that contains the view-projection matrix of the camera
  int CameraViewProjMatrixBufferIndex = -1;
  // Index of the buffer that contains the indices of the actor that were not culled
  int OutputIndicesBufferIndex = -1;
  // Index of the buffer that contains the number of actors that were not culled
  int OutputObjectCountBufferIndex = -1;

  // Callback command for when the bounds of an actor are recomputed
  vtkSmartPointer<vtkCallbackCommand> ActorBoundsRecomputedCallback;
  // Callback command for when the camera is modified and we need to update the camera data used by
  // the shader. This callback sets the CameraDirty boolean to true.
  vtkSmartPointer<vtkCallbackCommand> CameraModifiedCallback;
  // If true, this means that the camera was modified and we need to poll and upload the new data to
  // the shader.
  bool CameraDirty = true;
};

VTK_ABI_NAMESPACE_END

#endif
