// Copyright Epic Games, Inc. All Rights Reserved.

#if WITH_EDITOR

#include "SparseVolumeTextureOpenVDBUtility.h"
#include "SparseVolumeTextureOpenVDB.h"
#include "SparseVolumeTexture/SparseVolumeTexture.h"
#include "OpenVDBGridAdapter.h"
#include "OpenVDBImportOptions.h"

DEFINE_LOG_CATEGORY_STATIC(LogSparseVolumeTextureOpenVDBUtility, Log, All);

#if OPENVDB_AVAILABLE

namespace
{
	// Utility class acting as adapter between TArray64<uint8> and std::istream.
	// In order to work around a problem where a std::streambuf implementation is limited to
	// 2GB buffers, we need to manually update the pointers whenever we are done processing a 2GB chunk.
	class FArrayUint8StreamBuf : public std::streambuf
	{
	public:
		explicit FArrayUint8StreamBuf(TArray64<uint8>& Array)
		{
			char* Data = (char*)Array.GetData();
			size_t Num = Array.Num();
			FileBegin = Data;
			FileEnd = Data + Num;
			FileRead = Data;
		}

		// Calls setg() to with a set of pointers exposing a window into the input file.
		// Returns true if there are any more bytes to be read.
		bool UpdatePointers()
		{
			StreamBegin = FileRead;
			StreamEnd = FMath::Min(FileEnd, StreamBegin + ChunkSize);
			StreamRead = StreamBegin;
			FileRead = StreamRead;
			setg(StreamBegin, StreamRead, StreamEnd);
			const bool bHasBytesToRead = StreamBegin < StreamEnd;
			return bHasBytesToRead;
		}

		// This function is called by the parent class and is expected to be implemented by subclasses.
		// It requests n bytes to be copied into s. 
		std::streamsize xsgetn(char* s, std::streamsize n) override
		{
			for (std::streamsize ReadBytes = 0; ReadBytes < n;)
			{
				// We read all bytes of the input file. gptr() is the current read ptr and egptr() is the end ptr.
				if (gptr() == egptr() && !UpdatePointers())
				{
					check(gptr() == FileEnd);
					return ReadBytes;
				}
				// Try to read n bytes but make a smaller read if we would read past the current ptr window exposed to streambuf.
				const std::streamsize Available = FMath::Min(n - ReadBytes, static_cast<std::streamsize>(egptr() - gptr()));
				memcpy(s, gptr(), Available);
				// Advance pointers and the ReadBytes counter
				s += Available;
				ReadBytes += Available;
				StreamRead = gptr() + Available;
				FileRead = StreamRead;
				// Update the Next ptr in streambuf
				setg(StreamBegin, StreamRead, StreamEnd);
			}
			return n;
		}

		// Get the current character but don't advance the position. This is called by uflow() in the parent class when it runs out of bytes.
		// We use it to move the exposed window into the file data.
		int_type underflow() override
		{
			return (gptr() == egptr() && !UpdatePointers()) ? traits_type::eof() : *gptr();
		}

	private:
		static constexpr size_t ChunkSize = INT32_MAX;	// The size of the range to expose to std::streambuf with setg()
		char* FileBegin = nullptr;						// Begin ptr of the file data
		char* FileEnd = nullptr;						// One byte past the end of the file data
		char* FileRead = nullptr;						// The position up to which the file data has been exposed to/processed by the streambuf.
		char* StreamBegin = nullptr;					// The begin ptr of the currently exposed file chunk.
		char* StreamEnd = nullptr;						// The end ptr of the currently exposed file chunk.
		char* StreamRead = nullptr;						// The last value of the read/next ptr set with setg().
	};
}

static FOpenVDBGridInfo GetOpenVDBGridInfo(openvdb::GridBase::Ptr Grid, uint32 GridIndex, bool bCreateStrings)
{
	openvdb::CoordBBox VolumeActiveAABB = Grid->evalActiveVoxelBoundingBox();
	openvdb::Coord VolumeActiveDim = Grid->evalActiveVoxelDim();
	openvdb::Vec3d VolumeVoxelSize = Grid->voxelSize();
	openvdb::math::MapBase::ConstPtr MapBase = Grid->constTransform().baseMap();
	openvdb::Vec3d VoxelSize = MapBase->voxelSize();
	openvdb::Mat4d GridTransformVDB = MapBase->getAffineMap()->getConstMat4();

	FOpenVDBGridInfo GridInfo;
	GridInfo.Index = GridIndex;
	GridInfo.NumComponents = 0;
	GridInfo.Type = EOpenVDBGridType::Unknown;
	GridInfo.VolumeActiveAABBMin = FIntVector3(VolumeActiveAABB.min().x(), VolumeActiveAABB.min().y(), VolumeActiveAABB.min().z());
	GridInfo.VolumeActiveAABBMax = FIntVector3(VolumeActiveAABB.max().x() + 1, VolumeActiveAABB.max().y() + 1, VolumeActiveAABB.max().z() + 1); // +1 because CoordBBox::Max is inclusive, but we want exclusive
	GridInfo.VolumeActiveDim = FIntVector3(VolumeActiveDim.x(), VolumeActiveDim.y(), VolumeActiveDim.z());
	GridInfo.VolumeVoxelSize = FVector(VolumeVoxelSize.x(), VolumeVoxelSize.y(), VolumeVoxelSize.z());
	GridInfo.bIsInWorldSpace = Grid->isInWorldSpace();
	GridInfo.bHasUniformVoxels = Grid->hasUniformVoxels();

	for (int i = 0; i < 4; ++i)
	{
		for (int j = 0; j < 4; ++j)
		{
			GridInfo.Transform.M[i][j] = static_cast<float>(GridTransformVDB[j][i]);
		}
	}

	// Figure out the type/format of the grid
	if (Grid->isType<FOpenVDBFloat1Grid>())
	{
		GridInfo.NumComponents = 1;
		GridInfo.Type = EOpenVDBGridType::Float;
	}
	else if (Grid->isType<FOpenVDBFloat2Grid>())
	{
		GridInfo.NumComponents = 2;
		GridInfo.Type = EOpenVDBGridType::Float2;
	}
	else if (Grid->isType<FOpenVDBFloat3Grid>())
	{
		GridInfo.NumComponents = 3;
		GridInfo.Type = EOpenVDBGridType::Float3;
	}
	else if (Grid->isType<FOpenVDBFloat4Grid>())
	{
		GridInfo.NumComponents = 4;
		GridInfo.Type = EOpenVDBGridType::Float4;
	}
	else if (Grid->isType<FOpenVDBDouble1Grid>())
	{
		GridInfo.NumComponents = 1;
		GridInfo.Type = EOpenVDBGridType::Double;
	}
	else if (Grid->isType<FOpenVDBDouble2Grid>())
	{
		GridInfo.NumComponents = 2;
		GridInfo.Type = EOpenVDBGridType::Double2;
	}
	else if (Grid->isType<FOpenVDBDouble3Grid>())
	{
		GridInfo.NumComponents = 3;
		GridInfo.Type = EOpenVDBGridType::Double3;
	}
	else if (Grid->isType<FOpenVDBDouble4Grid>())
	{
		GridInfo.NumComponents = 4;
		GridInfo.Type = EOpenVDBGridType::Double4;
	}

	if (bCreateStrings)
	{
		GridInfo.Name = Grid->getName().c_str();

		FStringFormatOrderedArguments FormatArgs;
		FormatArgs.Add(GridInfo.Index);
		FormatArgs.Add(OpenVDBGridTypeToString(GridInfo.Type));
		FormatArgs.Add(GridInfo.Name);

		GridInfo.DisplayString = FString::Format(TEXT("{0}. Type: {1}, Name: \"{2}\""), FormatArgs);
	}

	return GridInfo;
}

#endif

bool IsOpenVDBGridValid(const FOpenVDBGridInfo& GridInfo, const FString& Filename)
{
	if (GridInfo.VolumeActiveDim.X * GridInfo.VolumeActiveDim.Y * GridInfo.VolumeActiveDim.Z == 0)
	{
		// SVT_TODO we should gently handle that case
		UE_LOG(LogSparseVolumeTextureOpenVDBUtility, Warning, TEXT("OpenVDB grid is empty due to volume size being 0: %s"), *Filename);
		return false;
	}

	if (!GridInfo.bHasUniformVoxels)
	{
		UE_LOG(LogSparseVolumeTextureOpenVDBUtility, Warning, TEXT("OpenVDB importer cannot handle non uniform voxels: %s"), *Filename);
		return false;
	}
	return true;
}

bool GetOpenVDBGridInfo(TArray64<uint8>& SourceFile, bool bCreateStrings, TArray<FOpenVDBGridInfo>* OutGridInfo)
{
#if OPENVDB_AVAILABLE
	FArrayUint8StreamBuf StreamBuf(SourceFile);
	std::istream IStream(&StreamBuf);
	openvdb::io::Stream Stream(IStream, false /*delayLoad*/);

	openvdb::GridPtrVecPtr Grids = Stream.getGrids();

	OutGridInfo->Empty(Grids->size());

	uint32 GridIndex = 0;
	for (openvdb::GridBase::Ptr& Grid : *Grids)
	{
		OutGridInfo->Add(GetOpenVDBGridInfo(Grid, GridIndex, bCreateStrings));
		++GridIndex;
	}

	return true;

#endif // OPENVDB_AVAILABLE
	return false;
}

static EPixelFormat GetMultiComponentFormat(ESparseVolumeAttributesFormat Format, uint32 NumComponents)
{
	switch (Format)
	{
	case ESparseVolumeAttributesFormat::Unorm8:
	{
		switch (NumComponents)
		{
		case 1: return PF_R8;
		case 2: return PF_R8G8;
		case 3:
		case 4: return PF_R8G8B8A8;
		}
		break;
	}
	case ESparseVolumeAttributesFormat::Float16:
	{
		switch (NumComponents)
		{
		case 1: return PF_R16F;
		case 2: return PF_G16R16F;
		case 3:
		case 4: return PF_FloatRGBA;
		}
		break;
	}
	case ESparseVolumeAttributesFormat::Float32:
	{
		switch (NumComponents)
		{
		case 1: return PF_R32_FLOAT;
		case 2: return PF_G32R32F;
		case 3:
		case 4: return PF_A32B32G32R32F;
		}
		break;
	}
	}
	return PF_Unknown;
}


#if OPENVDB_AVAILABLE

class FSparseVolumeRawSourceConstructionOpenVDBAdapter : public ISparseVolumeRawSourceConstructionAdapter
{
public:

	bool Initialize(TArray64<uint8>& SourceFile, const FOpenVDBImportOptions& ImportOptions, const FIntVector3& InVolumeBoundsMin)
	{
		Attributes = ImportOptions.Attributes;
		VolumeBoundsMin = InVolumeBoundsMin;

		// Compute some basic info about the number of components and which format to use
		EPixelFormat MultiCompFormat[NumAttributesDescs] = {};
		bool bNormalizedFormat[NumAttributesDescs] = {};
		bool bHasValidSourceGrids[NumAttributesDescs] = {};
		bool bAnySourceGridIndicesValid = false;

		for (int32 AttributesIdx = 0; AttributesIdx < NumAttributesDescs; ++AttributesIdx)
		{
			int32 NumRequiredComponents = 0;
			for (int32 ComponentIdx = 0; ComponentIdx < 4; ++ComponentIdx)
			{
				if (Attributes[AttributesIdx].Mappings[ComponentIdx].SourceGridIndex != INDEX_NONE)
				{
					check(Attributes[AttributesIdx].Mappings[ComponentIdx].SourceComponentIndex != INDEX_NONE);
					NumRequiredComponents = FMath::Max(ComponentIdx + 1, NumRequiredComponents);
					bHasValidSourceGrids[AttributesIdx] = true;
					bAnySourceGridIndicesValid = true;
				}
			}

			if (bHasValidSourceGrids[AttributesIdx])
			{
				NumComponents[AttributesIdx] = NumRequiredComponents == 3 ? 4 : NumRequiredComponents; // We don't support formats with only 3 components
				bNormalizedFormat[AttributesIdx] = Attributes[AttributesIdx].Format == ESparseVolumeAttributesFormat::Unorm8;
				MultiCompFormat[AttributesIdx] = GetMultiComponentFormat(Attributes[AttributesIdx].Format, NumComponents[AttributesIdx]);

				if (MultiCompFormat[AttributesIdx] == PF_Unknown)
				{
					UE_LOG(LogSparseVolumeTextureOpenVDBUtility, Warning, TEXT("SparseVolumeTexture is set to use an unsupported format: %i"), (int32)Attributes[AttributesIdx].Format);
					return false;
				}
			}
		}

		// All source grid indices are INDEX_NONE, so nothing was selected for import
		if (!bAnySourceGridIndicesValid)
		{
			UE_LOG(LogSparseVolumeTextureOpenVDBUtility, Warning, TEXT("SparseVolumeTexture has all components set to <None>, so there is nothing to import."));
			return false;
		}

		// Load file
		FArrayUint8StreamBuf StreamBuf(SourceFile);
		std::istream IStream(&StreamBuf);
		openvdb::io::Stream Stream(IStream, false /*delayLoad*/);

		// Check that the source grid indices are valid
		openvdb::GridPtrVecPtr Grids = Stream.getGrids();
		const size_t NumSourceGrids = Grids ? Grids->size() : 0;
		for (const FOpenVDBSparseVolumeAttributesDesc& AttributesDesc : Attributes)
		{
			for (const FOpenVDBSparseVolumeComponentMapping& Mapping : AttributesDesc.Mappings)
			{
				const int32 SourceGridIndex = Mapping.SourceGridIndex;
				if (SourceGridIndex != INDEX_NONE && SourceGridIndex >= (int32)NumSourceGrids)
				{
					UE_LOG(LogSparseVolumeTextureOpenVDBUtility, Warning, TEXT("SparseVolumeTexture has invalid index into the array of grids in the source file: %i"), SourceGridIndex);
					return false;
				}
			}
		}

		AttributesInfo[0].Format = MultiCompFormat[0];
		AttributesInfo[1].Format = MultiCompFormat[1];
		AttributesInfo[0].bNormalized = Attributes[0].bRemapInputForUnorm;
		AttributesInfo[1].bNormalized = Attributes[1].bRemapInputForUnorm;

		FIntVector3 SmallestAABBMin = FIntVector3(INT32_MAX);
		FIntVector3 LargestAABBMax = FIntVector3(INT32_MIN);

		// Compute per source grid data of up to 4 different grids (one per component)
		UniqueGridAdapters.SetNum((int32)NumSourceGrids);
		GridToComponentMappings.SetNum((int32)NumSourceGrids);
		for (int32 AttributesIdx = 0; AttributesIdx < NumAttributesDescs; ++AttributesIdx)
		{
			for (int32 CompIdx = 0; CompIdx < 4; ++CompIdx)
			{
				AttributesInfo[AttributesIdx].NormalizeScale[CompIdx] = 1.0f;
				const uint32 SourceGridIndex = Attributes[AttributesIdx].Mappings[CompIdx].SourceGridIndex;
				const uint32 SourceComponentIndex = Attributes[AttributesIdx].Mappings[CompIdx].SourceComponentIndex;
				if (SourceGridIndex == INDEX_NONE)
				{
					continue;
				}

				openvdb::GridBase::Ptr GridBase = (*Grids)[SourceGridIndex];

				// Try to reuse adapters. Internally they use caching to accelerate read accesses, 
				// so using three different adapters to access the three components of a single grid would be wasteful.
				if (UniqueGridAdapters[SourceGridIndex] == nullptr)
				{
					UniqueGridAdapters[SourceGridIndex] = CreateOpenVDBGridAdapter(GridBase);
					if (!UniqueGridAdapters[SourceGridIndex])
					{
						return false;
					}
				}

				FOpenVDBGridInfo GridInfo = GetOpenVDBGridInfo(GridBase, 0, false);
				if (!IsOpenVDBGridValid(GridInfo, TEXT("")))
				{
					return false;
				}

				SmallestAABBMin.X = FMath::Min(SmallestAABBMin.X, GridInfo.VolumeActiveAABBMin.X);
				SmallestAABBMin.Y = FMath::Min(SmallestAABBMin.Y, GridInfo.VolumeActiveAABBMin.Y);
				SmallestAABBMin.Z = FMath::Min(SmallestAABBMin.Z, GridInfo.VolumeActiveAABBMin.Z);
				LargestAABBMax.X = FMath::Max(LargestAABBMax.X, GridInfo.VolumeActiveAABBMax.X);
				LargestAABBMax.Y = FMath::Max(LargestAABBMax.Y, GridInfo.VolumeActiveAABBMax.Y);
				LargestAABBMax.Z = FMath::Max(LargestAABBMax.Z, GridInfo.VolumeActiveAABBMax.Z);

				AttributesInfo[AttributesIdx].FallbackValue[CompIdx] = UniqueGridAdapters[SourceGridIndex]->GetBackgroundValue(SourceComponentIndex);
				if (bNormalizedFormat[AttributesIdx] && Attributes[AttributesIdx].bRemapInputForUnorm)
				{
					float MinVal = 0.0f;
					float MaxVal = 0.0f;
					UniqueGridAdapters[SourceGridIndex]->GetMinMaxValue(SourceComponentIndex, &MinVal, &MaxVal);
					const float Diff = MaxVal - MinVal;
					AttributesInfo[AttributesIdx].NormalizeScale[CompIdx] = MaxVal > SMALL_NUMBER ? (1.0f / Diff) : 1.0f;
					AttributesInfo[AttributesIdx].NormalizeBias[CompIdx] = -MinVal * AttributesInfo[AttributesIdx].NormalizeScale[CompIdx];
				}

				FSingleGridToComponentMapping Mapping{};
				Mapping.AttributesIdx = (int32)AttributesIdx;
				Mapping.ComponentIdx = (int32)CompIdx;
				Mapping.GridComponentIdx = (int32)SourceComponentIndex;
				GridToComponentMappings[SourceGridIndex].Add(Mapping);
			}
		}

		AABBMin = SmallestAABBMin - VolumeBoundsMin;
		AABBMax = LargestAABBMax - VolumeBoundsMin;
		Resolution = LargestAABBMax - SmallestAABBMin;

		return true;
	}

	void GetAttributesInfo(FAttributesInfo& OutInfoA, FAttributesInfo& OutInfoB) const override
	{
		OutInfoA = AttributesInfo[0];
		OutInfoB = AttributesInfo[1];
	}

	FIntVector3 GetAABBMin() const override
	{
		return AABBMin;
	}

	FIntVector3 GetAABBMax() const override
	{
		return AABBMax;
	}

	FIntVector3 GetResolution() const override
	{
		return Resolution;
	}

	void IteratePhysicalSource(TFunctionRef<void(const FIntVector3& Coord, int32 AttributesIdx, int32 ComponentIdx, float VoxelValue)> OnVisit) const override
	{
		for (int32 GridIdx = 0; GridIdx < UniqueGridAdapters.Num(); ++GridIdx)
		{
			if (!UniqueGridAdapters[GridIdx])
			{
				continue;
			}

			UniqueGridAdapters[GridIdx]->IteratePhysical(
				[&](const FIntVector3& Coord, uint32 NumVoxelComponents, float* VoxelValues)
				{
					for (const FSingleGridToComponentMapping& Mapping : GridToComponentMappings[GridIdx])
					{
						OnVisit(Coord - VolumeBoundsMin, Mapping.AttributesIdx, Mapping.ComponentIdx, VoxelValues[Mapping.GridComponentIdx]);
					}
				});
		}
	}

private:

	struct FSingleGridToComponentMapping
	{
		int32 AttributesIdx;
		int32 ComponentIdx;
		int32 GridComponentIdx;
	};

	static constexpr int32 NumAttributesDescs = 2;
	TArray<TSharedPtr<IOpenVDBGridAdapterBase>> UniqueGridAdapters;
	TArray<TArray<FSingleGridToComponentMapping, TInlineAllocator<4>>> GridToComponentMappings;
	TStaticArray<FOpenVDBSparseVolumeAttributesDesc, NumAttributesDescs> Attributes;
	TStaticArray<uint32, NumAttributesDescs> NumComponents;
	TStaticArray<FAttributesInfo, NumAttributesDescs> AttributesInfo;
	FIntVector3 AABBMin;
	FIntVector3 AABBMax;
	FIntVector3 VolumeBoundsMin;
	FIntVector3 Resolution;
};

#endif // OPENVDB_AVAILABLE

bool ConvertOpenVDBToSparseVolumeTexture(TArray64<uint8>& SourceFile, const FOpenVDBImportOptions& ImportOptions, const FIntVector3& VolumeBoundsMin, FSparseVolumeRawSource& OutResult)
{
#if OPENVDB_AVAILABLE
	FSparseVolumeRawSourceConstructionOpenVDBAdapter Adapter;
	if (!Adapter.Initialize(SourceFile, ImportOptions, VolumeBoundsMin))
	{
		return false;
	}
	if (!OutResult.Construct(Adapter))
	{
		return false;
	}
	return true;
#else
	return false;
#endif // OPENVDB_AVAILABLE
}

const TCHAR* OpenVDBGridTypeToString(EOpenVDBGridType Type)
{
	switch (Type)
	{
	case EOpenVDBGridType::Float:
		return TEXT("Float");
	case EOpenVDBGridType::Float2:
		return TEXT("Float2");
	case EOpenVDBGridType::Float3:
		return TEXT("Float3");
	case EOpenVDBGridType::Float4:
		return TEXT("Float4");
	case EOpenVDBGridType::Double:
		return TEXT("Double");
	case EOpenVDBGridType::Double2:
		return TEXT("Double2");
	case EOpenVDBGridType::Double3:
		return TEXT("Double3");
	case EOpenVDBGridType::Double4:
		return TEXT("Double4");
	default:
		return TEXT("Unknown");
	}
}

#endif // WITH_EDITOR