// Copyright 2017-2023, Nicholas Sharp and the Polyscope contributors. https://polyscope.run

#pragma once

#include <cstdint>
#include <functional>
#include <unordered_map>
#include <vector>

#include "polyscope/render/engine.h"
#include "polyscope/utilities.h"
#include "polyscope/weak_handle.h"

namespace polyscope {
namespace render {

// forward declaration
class ManagedBufferRegistry;

/*
 * This class is a wrapper which sits on top of data buffers in Polyscope, and handles common data-management concerns
 * of:
 *
 *      (a) mirroring the buffer to the GPU/rendering framework
 *      (b) allowing external users to update data either on the CPU- or GPU-side, and updating mirrored copies
 * appropriately
 *      (c) managed _indexed_ data, which gts expanded according to some index set for access at rendering time.
 *
 * Most often this class is used to wrap structure/quantity data passed in by the user, such as a scalar quantity, but
 * it is also sometimes wraps automatically-computed values within Polyscope, such as a vertex normal buffer for
 * rendering.
 *
 * This class offers functions and accessors which can (and generally, MUST) be used to interact with the underlying
 * data buffer.
 */
template <typename T>
class ManagedBuffer : public virtual WeakReferrable {
public:
  // === Constructors
  // (second variants are advanced versions which allow creation of multi-dimensional texture values)

  // Manage a buffer of data which is explicitly set externally.
  ManagedBuffer(ManagedBufferRegistry* registry, const std::string& name, std::vector<T>& data);


  // Manage a buffer of data which gets computed lazily
  ManagedBuffer(ManagedBufferRegistry* registry, const std::string& name, std::vector<T>& data,
                std::function<void()> computeFunc);


  ~ManagedBuffer();


  // === Core members

  // A meaningful name for the buffer
  std::string name;

  // A globally-unique ID
  const uint64_t uniqueID;

  // The registry in which it is tracked (can be null)
  ManagedBufferRegistry* registry;


  // The raw underlying buffer which this class wraps that holds the data.
  // It is assumed that it never changes length (although this class may clear it to empty).
  //
  // It is possible that data.size() == 0 if the data is lazily computed and has not been computed yet, or if this
  // host-side buffer is invalidated because it is being updated externally directly on the render device.
  //
  // External users can write directly for this buffer. The required order of operations for writing to this buffer is:
  //    buff.ensureHostBufferAllocated();
  //    buff.data = // fill .data with your values
  //    buff.markHostBufferUpdated();
  //
  //
  std::vector<T>& data;

  // == Members for computed data

  // This class somewhat weirdly tracks data from two separate sources:
  //    A) Values directly stored in `data` by an external source
  //    B) Values that get computed lazily by some function when needed
  //
  // When `dataGetsComputed = true`, we are in Case B, and computeFunc() must be set to a callback that does the
  // lazy computing.

  bool dataGetsComputed;             // if true, the value gets computed on-demand by calling computeFunc()
  std::function<void()> computeFunc; // (optional) callback which populates the `data` buffer

  // sanity check helper
  void checkInvalidValues();

  // mark as texture, set size
  void setTextureSize(uint32_t sizeX);
  void setTextureSize(uint32_t sizeX, uint32_t sizeY);
  void setTextureSize(uint32_t sizeX, uint32_t sizeY, uint32_t sizeZ);
  std::array<uint32_t, 3> getTextureSize() const;


  // == Members for indexed data

  // == Basic interactions

  // Ensure that the `data` member vector reference is populated with the current values. In the common-case where
  // the user sets data and it never changes, then this function will do nothing. However, if e.g. the value is
  // being updated directly from GPU memory, this will mirror the updates to the cpu-side vector. Also, if the value
  // is lazily computed by computeFunc(), it ensures that that function has been called.
  void ensureHostBufferPopulated();

  // Ensure that the `data` member has the proper size. This does _not_ populate the buffer with any particular data,
  // just ensures it is allocated. It is useful for when an external wants to fill the buffer with data.
  void ensureHostBufferAllocated();

  // Combines calling ensureHostBufferPopulated() and returning a reference to the `data` member
  std::vector<T>& getPopulatedHostBufferRef();

  // If the contents of `data` are updated, this function MUST be called. It internally handles concerns like
  // reflecting updates to the render buffer.
  void markHostBufferUpdated();

  // Get the value at index `i`. It may be dynamically fetched from either the cpu-side `data` member or the render
  // buffer, depending on where the data currently lives.
  // If the data lives only on the device-side render buffer, this function is expensive, so don't call it in a
  // loop.
  T getValue(size_t ind);
  T getValue(size_t indX, size_t indY);              // only valid for 2d texture data
  T getValue(size_t indX, size_t indY, size_t indZ); // only valid for 3d texture data

  // If computeFunc() has already been called to populate the stored data, call it again to recompute the data, and
  // re-fill the buffer if necessary. This function is only meaningful in the case where `dataGetsComputed = true`.
  void recomputeIfPopulated();

  bool hasData(); // true if there is valid data on either the host or device
  size_t size();  // size of the data (number of entries)

  // Is it an attribute, texture1d, texture2d, etc?
  DeviceBufferType getDeviceBufferType();

  std::string summaryString(); // for debugging

  // ========================================================================
  // == Direct access to the GPU (device-side) render attribute buffer
  // ========================================================================

  // NOTE: this is only for attribute-accessed buffers (DeviceBufferType::Attribute). See the variants below for
  // textures.

  // NOTE: This class follows the policy that once the render buffer is allocated, it is always immediately kept
  // updated to reflect any external changes.

  // Get a reference to the underlying GPU-side attribute buffer
  // Once this reference is created, it will always be immediately updated to reflect any external changes to the
  // data. (note that if you write to this buffer externally, you MUST call markRenderAttributeBufferUpdated()
  // below)
  std::shared_ptr<render::AttributeBuffer> getRenderAttributeBuffer();

  // Tell Polyscope that you wrote updated data into the render buffer. This MUST be called after externally writing
  // to the buffer from getRenderBuffer() above.
  void markRenderAttributeBufferUpdated();

  // ========================================================================
  // == Indexed views
  // ========================================================================

  // For some data (e.g. values a vertices of a mesh), we store the data in a canonical ordering (one value per
  // vertex), but need to index it at render time according to some specified list of indices (one value per
  // triangle corner). Rendering frameworks _do_ have explicit support for this case (e.g. IndexedArray DrawMode in
  // openGL), but for various reasons it does not make sense to use those features. Instead, we explicitly expand
  // out the data into a render buffer for drawing as view[i] = data[indices[i]].
  //
  // If such an index is used, it should be set to the `indices` argument below, and this class will automatically
  // handle expanding out the indexed data to populate the returned buffer. External callers can still update the
  // data directly on the host, and this class will handle updating an indexed version of the data for drawing.
  //
  // Internally, these indexed views are cached. It is safe to call this function many times, after the first the
  // same view will be returned repeatedly at no additional cost.
  std::shared_ptr<render::AttributeBuffer> getIndexedRenderAttributeBuffer(ManagedBuffer<uint32_t>& indices);

  // Get a copy of the data viewed through an index, such that view[i] = data[indices[i]].
  //
  // This follows the same logic as above, but rather than returning a render buffer it simply returns a host-side
  // copy (which is not cached).
  std::vector<T> getIndexedView(ManagedBuffer<uint32_t>& indices);

  // ========================================================================
  // == Direct access to the GPU (device-side) render texture buffer
  // ========================================================================
  //
  // NOTE: these follow the same semantics as the attribute version above, but these apply when the buffer is a texture
  // (DeviceBufferType::Texture1d, etc).

  std::shared_ptr<render::TextureBuffer> getRenderTextureBuffer();
  void markRenderTextureBufferUpdated();


protected:
  // == Internal members

  bool hostBufferIsPopulated; // true if the host buffer contains currently-valid data

  std::shared_ptr<render::AttributeBuffer> renderAttributeBuffer;
  std::shared_ptr<render::TextureBuffer> renderTextureBuffer;

  // For storing as textures

  // For data that can be interpreted as a 1/2/3 dimensional texture
  DeviceBufferType deviceBufferType = DeviceBufferType::Attribute; // this gets set when you call setTextureSize
  uint32_t sizeX = 0;
  uint32_t sizeY = 0; // holds 0 if texture dim < 2
  uint32_t sizeZ = 0; // holds 0 if texture dim < 3


  // == Internal representation of indexed views
  // NOTE: this seems like a problem, we are storing pointers as keys in a cache. Here, it works out because if the
  // key ptr becomes invalid, the value weak_ptr must also be invalid, and we check that before dereferencing the
  // key.
  std::vector<std::tuple<render::ManagedBuffer<uint32_t>*, std::weak_ptr<render::AttributeBuffer>>>
      existingIndexedViews;
  void updateIndexedViews();
  void removeDeletedIndexedViews();

  // == Internal helper functions

  void invalidateHostBuffer();
  bool deviceBufferTypeIsTexture();
  void checkDeviceBufferTypeIs(DeviceBufferType targetType);
  void checkDeviceBufferTypeIsTexture();

  enum class CanonicalDataSource { HostData = 0, NeedsCompute, RenderBuffer };
  CanonicalDataSource currentCanonicalDataSource();

  // Manage the program which copies indexed data from the renderBuffer to the indexed views
  void ensureHaveBufferIndexCopyProgram();
  void invokeBufferIndexCopyProgram();
  std::shared_ptr<render::ShaderProgram> bufferIndexCopyProgram;
};


// == Manage a store of all registered managed buffers

// These registries are set up to be static: once a buffer is added it is never removed.

// Helper class representing a map [name] --> [buffer]
template <typename T>
struct ManagedBufferMap {
public:
  void addManagedBuffer(ManagedBuffer<T>* buffer);
  void removeManagedBuffer(ManagedBuffer<T>* buffer);

  ManagedBuffer<T>& getManagedBuffer(std::string name);
  bool hasManagedBuffer(std::string name);

  // internal helper for template things
  static ManagedBufferMap<T>& getManagedBufferMapRef(ManagedBufferRegistry* r);

private:
  // the actual store
  // NOTE: we do NOT support removal
  std::vector<ManagedBuffer<T>*> allBuffers;
};


// A registry of buffer of various types
// Classes (structures, quantities) can extend this to track their own buffers
class ManagedBufferRegistry {
public:
  // get a reference to any buffer that currently exists in Polyscope, by name
  template <typename T>
  ManagedBuffer<T>& getManagedBuffer(std::string name);

  // check for the existence of a managed buffer with a given name
  template <typename T>
  bool hasManagedBuffer(std::string name);

  // checks for a managed buffer with the given name of any type
  // return value is (bool, type), the bool indicates whether the buffer was found, and if so the type indicates what
  // type it was
  std::tuple<bool, ManagedBufferType> hasManagedBufferType(std::string name);

  template <typename T>
  void addManagedBuffer(ManagedBuffer<T>* buffer);

  template <typename T>
  void removeManagedBuffer(ManagedBuffer<T>* buffer);

  // clang-format off
  ManagedBufferMap<float>        managedBufferMap_float;
  ManagedBufferMap<double>       managedBufferMap_double;
  ManagedBufferMap<glm::vec2>    managedBufferMap_vec2;
  ManagedBufferMap<glm::vec3>    managedBufferMap_vec3;
  ManagedBufferMap<glm::vec4>    managedBufferMap_vec4;
  ManagedBufferMap<std::array<glm::vec3,2>> managedBufferMap_arr2vec3;
  ManagedBufferMap<std::array<glm::vec3,3>> managedBufferMap_arr3vec3;
  ManagedBufferMap<std::array<glm::vec3,4>> managedBufferMap_arr4vec3;
  ManagedBufferMap<int32_t>      managedBufferMap_int32;
  ManagedBufferMap<glm::ivec2>   managedBufferMap_ivec2;
  ManagedBufferMap<glm::ivec3>   managedBufferMap_ivec3;
  ManagedBufferMap<glm::ivec4>   managedBufferMap_ivec4;
  ManagedBufferMap<uint32_t>     managedBufferMap_uint32;
  ManagedBufferMap<glm::uvec2>   managedBufferMap_uvec2;
  ManagedBufferMap<glm::uvec3>   managedBufferMap_uvec3;
  ManagedBufferMap<glm::uvec4>   managedBufferMap_uvec4;
  // clang-format on
};

} // namespace render

std::string typeName(ManagedBufferType type);

} // namespace polyscope

#include "polyscope/render/managed_buffer.ipp"
