// Copyright 1998-2014 Epic Games, Inc. All Rights Reserved.

#include "Core.h"
#include "CoreMisc.h"
#include "ImageCore.h"
#include "ModuleInterface.h"
#include "ModuleManager.h"
#include "TargetPlatform.h"
#include "TextureCompressorModule.h"
#include "PixelFormat.h"
#include "GenericPlatformProcess.h"

DEFINE_LOG_CATEGORY_STATIC(LogTextureFormatPVR, Log, All);

/**
 * Macro trickery for supported format names.
 */
#define ENUM_SUPPORTED_FORMATS(op) \
	op(PVRTC2) \
	op(PVRTC4) \
	op(PVRTCN) \
	op(AutoPVRTC)

#define DECL_FORMAT_NAME(FormatName) static FName GTextureFormatName##FormatName = FName(TEXT(#FormatName));
ENUM_SUPPORTED_FORMATS(DECL_FORMAT_NAME);
#undef DECL_FORMAT_NAME

#define DECL_FORMAT_NAME_ENTRY(FormatName) GTextureFormatName##FormatName ,
static FName GSupportedTextureFormatNames[] =
{
	ENUM_SUPPORTED_FORMATS(DECL_FORMAT_NAME_ENTRY)
};
#undef DECL_FORMAT_NAME_ENTRY

#undef ENUM_SUPPORTED_FORMATS

// PVR file header format
#if PLATFORM_SUPPORTS_PRAGMA_PACK
#pragma pack(push, 4)
#endif
struct FPVRHeader
{
	uint32 Version;
	uint32 Flags;
	uint64 PixelFormat ;
	uint32 ColorSpace;
	uint32 ChannelType;
	uint32 Height;
	uint32 Width;
	uint32 Depth;
	uint32 NumSurfaces;
	uint32 NumFaces;
	uint32 NumMipmaps;
	uint32 MetaDataSize;
};
#if PLATFORM_SUPPORTS_PRAGMA_PACK
#pragma pack(pop)
#endif

/**
 * Converts a power-of-two image to a square format (ex: 256x512 -> 512x512). Be wary of memory waste when too many texture are not square.
 *
 * @param Image The image to be converted to a square
 * @return true if the image was converted successfully, else false
 */
static bool SquarifyImage(FImage& Image, uint32 MinSquareSize)
{
	// Early out
	if (Image.SizeX == Image.SizeY && Image.SizeX >= int32(MinSquareSize))
	{
		return true;
	}

    // Figure out the squarified size
    uint32 SquareSize = FMath::Max(Image.SizeX, Image.SizeY);
	if(SquareSize < MinSquareSize)
	{
		SquareSize = MinSquareSize;
	}
    
    // Calculate how many times to duplicate each row of column
    uint32 MultX = SquareSize / Image.SizeX;
    uint32 MultY = SquareSize / Image.SizeY;

	// Only give memory overhead warning if we're actually going to use a larger image
	// Small mips that have to be upscaled for compression only save the smaller mip for use
	if(MultX == 1 || MultY == 1)
	{
		float FOverhead = float(FMath::Min(Image.SizeX, Image.SizeY)) / float(SquareSize);
		int32 POverhead = FMath::RoundToInt(100.0f - (FOverhead * 100.0f));
		UE_LOG(LogTextureFormatPVR, Warning, TEXT("Expanding mip (%d,%d) to (%d, %d). Memory overhead: ~%d%%"),
			Image.SizeX, Image.SizeY, SquareSize, SquareSize, POverhead);
	}
	else if (MultX != MultY)
	{
		float FOverhead = float(FMath::Min(Image.SizeX, Image.SizeY)) / float(FMath::Max(Image.SizeX, Image.SizeY));
		int32 POverhead = FMath::RoundToInt(100.0f - (FOverhead * 100.0f));
		UE_LOG(LogTextureFormatPVR, Warning, TEXT("Expanding mip (%d,%d) to (%d, %d). Memory overhead: ~%d%%"),
			Image.SizeX, Image.SizeY, FMath::Max(Image.SizeX, Image.SizeY), FMath::Max(Image.SizeX, Image.SizeY), POverhead);
	}
    
    // Allocate room to fill out into
    TArray<uint32> SquareRawData;
	SquareRawData.Init(SquareSize * SquareSize * Image.NumSlices);
    
	int32 SourceSliceSize = Image.SizeX * Image.SizeY;
	int32 DestSliceSize = SquareSize * SquareSize;
	for ( int32 SliceIndex=0; SliceIndex < Image.NumSlices; ++SliceIndex )
	{
		uint32* RectData = ((uint32*)Image.RawData.GetData()) + SliceIndex * SourceSliceSize;
		uint32* SquareData = ((uint32*)SquareRawData.GetData()) + SliceIndex * DestSliceSize;
		for ( int32 Y = 0; Y < Image.SizeY; ++Y )
		{
			for ( int32 X = 0; X < Image.SizeX; ++X )
			{
				uint32 SourceColor = *(RectData + Y * Image.SizeX + X);
            
				for ( uint32 YDup = 0; YDup < MultY; ++YDup )
				{
					for ( uint32 XDup = 0; XDup < MultX; ++XDup )
					{
						uint32* DestColor = SquareData + ((Y * MultY + YDup) * SquareSize + (X * MultX + XDup));
						*DestColor = SourceColor;
					}
				}
			}
		}
	}

	// Put the new image data into the existing Image (copying from uint32 array to uint8 array)
	Image.RawData.Empty(SquareSize * SquareSize * Image.NumSlices * sizeof(uint32));
	Image.RawData.Init(SquareSize * SquareSize * Image.NumSlices * sizeof(uint32));
	uint32* FinalData = (uint32*)Image.RawData.GetData();
	FMemory::Memcpy(Image.RawData.GetData(), SquareRawData.GetData(), SquareSize * SquareSize * Image.NumSlices * sizeof(uint32));

    Image.SizeX = SquareSize;
    Image.SizeY = SquareSize;

    return true;
}

static void DeriveNormalZ(FImage& Image)
{
	int32 SliceSize = Image.SizeX * Image.SizeY;
	for ( int32 SliceIndex=0; SliceIndex < Image.NumSlices; ++SliceIndex )
	{
		FColor* RectData = (FColor*)Image.RawData.GetData() + SliceIndex * SliceSize;
		for(int32 Y = 0; Y < Image.SizeY; ++Y)
		{
			for(int32 X = 0; X < Image.SizeX; ++X)
			{
				FColor& SourceColor = *(RectData + Y * Image.SizeX + X);

				const float NormalX = SourceColor.R / 255.0f * 2 - 1;
				const float NormalY = SourceColor.G / 255.0f * 2 - 1;
				const float NormalZ = FMath::Sqrt(FMath::Clamp<float>(1 - (NormalX * NormalX + NormalY * NormalY), 0, 1));
				SourceColor.B = FMath::TruncToInt((NormalZ + 1) / 2.0f * 255.0f);
			}
		}
	}
}

/**
 * Checks if the passed image is a proper power-of-2 image
 *
 * @param Image The Image to evaluate
 * @return true if the image is a power of 2, else false
 */
static bool ValidateImagePower(const FImage& Image)
{
    // Image must already have power of 2 dimensions
    bool bDimensionsValid = true;
    int DimX = Image.SizeX;
    int DimY = Image.SizeY;
    while(DimX >= 2)
    {
        if(DimX % 2 == 1)
        {
            bDimensionsValid = false;
            break;
        }
        DimX /= 2;
    }
    while(DimY >= 2 && bDimensionsValid)
    {
        if(DimY % 2 == 1)
        {
            bDimensionsValid = false;
            break;
        }
        DimY /= 2;
    }
    
    if(!bDimensionsValid)
    {
        return false;
    }

	return true;
}

/**
 * Fills the output structure with the original uncompressed mip information
 *
 * @param InImage The mip to compress
 * @param OutCompressImage The output image (uncompressed in this case)
 */
static void UseOriginal(const FImage& InImage, FCompressedImage2D& OutCompressedImage, EPixelFormat CompressedPixelFormat, bool bSRGB)
{   
	// Get Raw Data
	FImage Image;
	InImage.CopyTo(Image, ERawImageFormat::BGRA8, bSRGB);

    // Fill out the output information
    OutCompressedImage.SizeX = Image.SizeX;
    OutCompressedImage.SizeY = Image.SizeY;
    OutCompressedImage.PixelFormat = CompressedPixelFormat;
	
	// Output Data
	OutCompressedImage.RawData.Init(Image.SizeX * Image.SizeY * 4);
    void* MipData = (void*)Image.RawData.GetData();
    FMemory::Memcpy(MipData, Image.RawData.GetData(), Image.SizeX * Image.SizeY * 4);
}

/**
 * PVR texture format handler.
 */
class FTextureFormatPVR : public ITextureFormat
{
	virtual bool AllowParallelBuild() const OVERRIDE
	{
		return true;
	}

	virtual uint16 GetVersion(FName Format) const OVERRIDE
	{
		return 6;
	}

	virtual void GetSupportedFormats(TArray<FName>& OutFormats) const OVERRIDE
	{
		for (int32 i = 0; i < ARRAY_COUNT(GSupportedTextureFormatNames); ++i)
		{
			OutFormats.Add(GSupportedTextureFormatNames[i]);
		}
	}

	virtual FTextureFormatCompressorCaps GetFormatCapabilities() const OVERRIDE
	{
		FTextureFormatCompressorCaps RetCaps;
		// PVR compressor is limited to <=4096 in any direction.
		RetCaps.MaxTextureDimension = 4096;
		return RetCaps;
	}

	virtual bool CompressImage(
		const FImage& InImage,
		const struct FTextureBuildSettings& BuildSettings,
		bool bImageHasAlphaChannel,
		FCompressedImage2D& OutCompressedImage
		) const OVERRIDE
	{
        // Get Raw Image Data from passed in FImage
        FImage Image;
        InImage.CopyTo(Image, ERawImageFormat::BGRA8, BuildSettings.bSRGB);
        
        // Get the compressed format
        EPixelFormat CompressedPixelFormat = PF_Unknown;
		if (BuildSettings.TextureFormatName == GTextureFormatNamePVRTC2)
		{
			CompressedPixelFormat = PF_PVRTC2;
		}
		else if (BuildSettings.TextureFormatName == GTextureFormatNamePVRTC4 || BuildSettings.TextureFormatName == GTextureFormatNamePVRTCN)
		{
			CompressedPixelFormat = PF_PVRTC4;
		}
		else if (BuildSettings.TextureFormatName == GTextureFormatNameAutoPVRTC)
		{
			CompressedPixelFormat = bImageHasAlphaChannel ? PF_PVRTC4 : PF_PVRTC2;
		}

		// Verify Power of 2
		if ( !ValidateImagePower(Image) )
		{
            UE_LOG(LogTextureFormatPVR, Warning, TEXT("Mip size (%d,%d) does not have power-of-two dimensions and cannot be compressed to PVRTC%d"), 
				Image.SizeX, Image.SizeY, CompressedPixelFormat == PF_PVRTC2 ? 2 : 4);
            return false;
		}

		// Squarify image
		int32 FinalSquareSize = FGenericPlatformMath::Max(Image.SizeX, Image.SizeY);
		SquarifyImage(Image, (CompressedPixelFormat == PF_PVRTC2) ? 16 : 8);
		check(Image.SizeX == Image.SizeY);

		if ( BuildSettings.TextureFormatName == GTextureFormatNamePVRTCN )
		{
			// Derive Z from X and Y to be consistent with BC5 normal maps used on PC (toss the texture's actual Z)
			DeriveNormalZ(Image);
		}

		bool bCompressionSucceeded = true;
		int32 SliceSize = Image.SizeX * Image.SizeY;
		for (int32 SliceIndex = 0; SliceIndex < Image.NumSlices && bCompressionSucceeded; ++SliceIndex)
		{
			TArray<uint8> CompressedSliceData;
			bCompressionSucceeded = CompressImageUsingPVRTexTool(
				Image.AsBGRA8() + SliceIndex * SliceSize,
				CompressedPixelFormat,
				Image.SizeX,
				Image.SizeY,
				Image.bSRGB,
				FinalSquareSize,
				CompressedSliceData
				);
			OutCompressedImage.RawData.Append(CompressedSliceData);
		}

		if ( bCompressionSucceeded )
		{
			OutCompressedImage.SizeX = FinalSquareSize;
			OutCompressedImage.SizeY = FinalSquareSize;
			OutCompressedImage.PixelFormat = CompressedPixelFormat;
		}

        // Return success status
        return bCompressionSucceeded;
	}

	static bool CompressImageUsingPVRTexTool( void* SourceData, EPixelFormat PixelFormat, int32 SizeX, int32 SizeY, bool bSRGB, int32 FinalSquareSize, TArray<uint8>& OutCompressedData )
	{
		// Figure out whether to use 2 bits or 4 bits per pixel (PVRTC2/PVRTC4)
		bool bIsPVRTC2 = (PixelFormat == PF_PVRTC2);

		const int32 BlockSizeX = bIsPVRTC2 ? 8 : 4;		// PVRTC2 uses 8x4 blocks, PVRTC4 uses 4x4 blocks
		const int32 BlockSizeY = 4;
		const int32 BlockBytes = 8;						// Both PVRTC2 and PVRTC4 are 8 bytes per block
		const uint32 DestSizeX = FinalSquareSize;
		const uint32 DestSizeY = FinalSquareSize;

		// min 2x2 blocks per mip
		const uint32 DestBlocksX = FGenericPlatformMath::Max<uint32>(DestSizeX / BlockSizeX, 2);
		const uint32 DestBlocksY = FGenericPlatformMath::Max<uint32>(DestSizeY / BlockSizeY, 2);
		const uint32 DestNumBytes = DestBlocksX * DestBlocksY * BlockBytes;

		// If using an image that's too small, compressor needs to generate mips for us with an upscaled image
		check(SizeX == SizeY);
		const int32 SourceSquareSize = SizeX;
		bool bGenerateMips = (FinalSquareSize < SourceSquareSize) ? true : false;

		// Allocate space to store compressed data.
		OutCompressedData.Empty(DestNumBytes);
		OutCompressedData.AddUninitialized(DestNumBytes);
		void* MipData = OutCompressedData.GetData();

		// Write SourceData into PVR file on disk

		// Init Header
		FPVRHeader PVRHeader;
		PVRHeader.Version = 0x03525650; // endianess does not match
		PVRHeader.Flags = 0;
		PVRHeader.PixelFormat = 0x0808080861726762; // Format of the UTexture 8bpp and BGRA ordered
		PVRHeader.ColorSpace = 0; // Setting to 1 indicates SRGB, but causes PVRTexTool to unpack to linear. We want the image to remain in sRGB space as we do the conversion in the shader.
		PVRHeader.ChannelType = 0;
		PVRHeader.Height = SizeY;
		PVRHeader.Width = SizeX;
		PVRHeader.Depth = 1;
		PVRHeader.NumSurfaces = 1;
		PVRHeader.NumFaces = 1;
		PVRHeader.NumMipmaps = 1;
		PVRHeader.MetaDataSize = 0;

		// Get file paths for intermediates. Unique path to avoid filename collision
		FGuid Guid;
		FPlatformMisc::CreateGuid(Guid);
		FString InputFilePath = FString::Printf(TEXT("Cache/%x%x%x%xRGBToPVRIn.pvr"), Guid.A, Guid.B, Guid.C, Guid.D);
		InputFilePath = FPaths::GameIntermediateDir() + InputFilePath;
		FString OutputFilePath = FString::Printf(TEXT("Cache/%x%x%x%xRGBToPVROut.pvr"), Guid.A, Guid.B, Guid.C, Guid.D);
		OutputFilePath = FPaths::GameIntermediateDir() + OutputFilePath;

		FArchive* PVRFile = NULL;
		while(!PVRFile)
		{
			PVRFile = IFileManager::Get().CreateFileWriter(*InputFilePath);	// Occasionally returns NULL due to error code ERROR_SHARING_VIOLATION
			FPlatformProcess::Sleep(0.01f);									// ... no choice but to wait for the file to become free to access
		}

		// Write out header
		uint32 HeaderSize = sizeof( PVRHeader );
		check(HeaderSize==52);
		PVRFile->Serialize(&PVRHeader, HeaderSize);

		// Write out uncompressed data
		PVRFile->Serialize(SourceData, SizeX * SizeY * sizeof(uint32));

		// Finished writing file
		PVRFile->Close();
		delete PVRFile;

		// Compress PVR file to PVRTC
#if PLATFORM_MAC
		// Compress PVR file to PVRTC (using apple's commandline tool)
//		FString Params = FString::Printf(TEXT("-sdk iphoneos texturetool %s-e PVRTC --bits-per-pixel-%d -o %s%s -f PVR %s%s"),
//			bGenerateMips ? TEXT("-m ") : TEXT(""),
//			bIsPVRTC2 ? 2 : 4,
//			FPlatformProcess::BaseDir(), *OutputFilePath,
//			FPlatformProcess::BaseDir(), *InputFilePath);

		// Use PowerVR's tool
		FString Params = FString::Printf(TEXT("-i \"%s\" -o \"%s\" %s -legacypvr -q pvrtcbest -f PVRTC1_%d"),
			*InputFilePath, *OutputFilePath,
			bGenerateMips ? TEXT("-m") : TEXT(""),
			bIsPVRTC2 ? 2 : 4);
//		FString CompressorPath(TEXT("/usr/bin/xcrun"));
		FString CompressorPath(FPaths::EngineDir() + TEXT("Binaries/ThirdParty/ImgTec/PVRTexToolCL"));
		UE_LOG(LogTemp, Log, TEXT("Running texturetool with '%s'"), *Params);
#else
		// Use PowerVR's tool
		FString Params = FString::Printf(TEXT("%s%s -fOGLPVRTC%d -yflip0 -i\"%s\" -o\"%s\""),
			bGenerateMips ? TEXT("-m ") : TEXT(""),                             
			TEXT("-pvrtciterations8"),
			bIsPVRTC2 ? 2 : 4, *InputFilePath, *OutputFilePath);
		FString CompressorPath(FPaths::EngineDir() + TEXT("Binaries/ThirdParty/ImgTec/PVRTexTool.exe"));
#endif

		// Give a debug message about the process
		if (IsRunningCommandlet())
		{
			//UE_LOG(LogTextureFormatPVR, Display, TEXT("Compressing mip (%dx%d) to PVRTC%d for mobile devices..."), ImageSizeX, ImageSizeY, bIsPVRTC2 ? 2 : 4);
		}

		// Start Compressor
		FProcHandle Proc = FPlatformProcess::CreateProc(*CompressorPath, *Params, true, false, false, NULL, -1, NULL, NULL);
		bool bConversionWasSuccessful = true;

		// Failed to start the compressor process?
		if (!Proc.IsValid())
		{
			UE_LOG(LogTextureFormatPVR, Error, TEXT("Failed to start PVR compressor tool. (Path:%s)"), *CompressorPath);
			bConversionWasSuccessful = false;
		}

		if ( bConversionWasSuccessful )
		{
			// Wait for the process to complete
			int ReturnCode;
			while (!FPlatformProcess::GetProcReturnCode(Proc, &ReturnCode))
			{
				FPlatformProcess::Sleep(0.01f);
			}
			Proc.Close();

			// Did it fail?
			if ( ReturnCode != 0 )
			{
				UE_LOG(LogTextureFormatPVR, Error, TEXT("PVR tool Failed with Return Code %d Mip Size (%d,%d)"), ReturnCode, SizeX, SizeY);
				bConversionWasSuccessful = false;
			}
		}

		// Open compressed file and put the data in OutCompressedImage
		if ( bConversionWasSuccessful )
		{
			// Calculate which mip to pull from compressed image
			int32 MipLevel = 0;
			{
				int32 i = SizeX;
				while ( i > FinalSquareSize )
				{
					i /= 2;
					++MipLevel;
				}
			}

			// Get Raw File Data
			TArray<uint8> PVRData;
			FFileHelper::LoadFileToArray(PVRData, *OutputFilePath);

			// Process It
			FPVRHeader* Header = (FPVRHeader*)PVRData.GetData();

			// Calculate the offset to get to the mip data
			int FileOffset = HeaderSize;
			for(int32 i = 0; i < MipLevel; ++i)
			{
				// Get the mip size for each image before the mip we want
				uint32 LocalMipSizeX = FGenericPlatformMath::Max<uint32>(SizeX >> i, 1);
				uint32 LocalMipSizeY = LocalMipSizeX;
				uint32 LocalBlocksX = FGenericPlatformMath::Max<uint32>(LocalMipSizeX / BlockSizeX, 2);
				uint32 LocalBlocksY = FGenericPlatformMath::Max<uint32>(LocalMipSizeY / BlockSizeY, 2);
				uint32 LocalMipSize = LocalBlocksX * LocalBlocksY * BlockBytes;

				// Add that mip's size to the offset
				FileOffset += LocalMipSize;
			}

			// Copy compressed data
			FMemory::Memcpy(MipData, PVRData.GetTypedData() + FileOffset, DestNumBytes);
		}

		// Delete intermediate files
		IFileManager::Get().Delete(*InputFilePath);
		IFileManager::Get().Delete(*OutputFilePath);

		return bConversionWasSuccessful;
	}
};

/**
 * Module for PVR texture compression.
 */
static ITextureFormat* Singleton = NULL;

class FTextureFormatPVRModule : public ITextureFormatModule
{
public:
	virtual ~FTextureFormatPVRModule()
	{
		delete Singleton;
		Singleton = NULL;
	}
	virtual ITextureFormat* GetTextureFormat()
	{
		if (!Singleton)
		{
			Singleton = new FTextureFormatPVR();
		}
		return Singleton;
	}
};

IMPLEMENT_MODULE(FTextureFormatPVRModule, TextureFormatPVR);