/* Copyright (C) 2024 Wildfire Games.
 * This file is part of 0 A.D.
 *
 * 0 A.D. is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 2 of the License, or
 * (at your option) any later version.
 *
 * 0 A.D. is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with 0 A.D.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "precompiled.h"

#include "TextureConverter.h"

#include "lib/allocators/shared_ptr.h"
#include "lib/bits.h"
#include "lib/regex.h"
#include "lib/tex/tex.h"
#include "lib/timer.h"
#include "maths/MD5.h"
#include "ps/CLogger.h"
#include "ps/CStr.h"
#include "ps/Profiler2.h"
#include "ps/TaskManager.h"
#include "ps/Util.h"
#include "ps/XML/Xeromyces.h"

#if CONFIG2_NVTT

#include "nvtt/nvtt.h"

// We assume NVTT is recent enough to support the alpha flag in the DXT1a format. If users try to use an
// old version of NVTT, the game will crash when trying to decode dds files generated by NVTT.
//
// The support was added upstream in https://github.com/castano/nvidia-texture-tools/commit/782a127071895f538c1ae49925a6e15687e3c966
// so, in theory, 2.0.7 and newer should be enough, but all 2.0.x releases define NVTT_VERSION as 200, so
// we can't distinguish them. NVTT_VERSION is 201 in all development versions of the 2.0.x era, so we also
// have to exclude that value.
#if !defined NVTT_VERSION || NVTT_VERSION == 200 || NVTT_VERSION == 201
#error Please use NVTT 2.1.0 or newer. \
If your system does not provide it, you should use the bundled version by NOT passing --with-system-nvtt to premake.
#endif

namespace
{

// Completely arbitrary constant - there is some main-thread cost to loading textures and the textures
// use a lot of memory, so probably should not be too high.
// Note that some results in the result queue may already be ready.
constexpr size_t MAX_QUEUE_SIZE_FOR_OPTIMAL_UTILIZATION{12};

/**
 * Output handler to collect NVTT's output into a simplistic buffer.
 */
struct BufferOutputHandler : public nvtt::OutputHandler
{
	std::vector<u8> buffer;

	virtual void beginImage(int UNUSED(size), int UNUSED(width), int UNUSED(height), int UNUSED(depth), int UNUSED(face), int UNUSED(miplevel))
	{
	}

	virtual bool writeData(const void* data, int size)
	{
		size_t off = buffer.size();
		buffer.resize(off + size);
		memcpy(&buffer[off], data, size);
		return true;
	}

	virtual void endImage()
	{
	}
};

/**
 * Arguments to the asynchronous task.
 */
struct ConversionRequest
{
	VfsPath dest;
	CTexturePtr texture;
	nvtt::InputOptions inputOptions;
	nvtt::CompressionOptions compressionOptions;
	nvtt::OutputOptions outputOptions;
};

} // anonymous namespace

/**
 * Response from the asynchronous task.
 */
struct CTextureConverter::ConversionResult
{
	VfsPath dest;
	CTexturePtr texture;
	BufferOutputHandler output;
	bool ret; // true if the conversion succeeded
};

#endif // CONFIG2_NVTT

void CTextureConverter::Settings::Hash(MD5& hash)
{
	hash.Update((const u8*)&format, sizeof(format));
	hash.Update((const u8*)&mipmap, sizeof(mipmap));
	hash.Update((const u8*)&normal, sizeof(normal));
	hash.Update((const u8*)&alpha, sizeof(alpha));
	hash.Update((const u8*)&filter, sizeof(filter));
	hash.Update((const u8*)&kaiserWidth, sizeof(kaiserWidth));
	hash.Update((const u8*)&kaiserAlpha, sizeof(kaiserAlpha));
	hash.Update((const u8*)&kaiserStretch, sizeof(kaiserStretch));
}

CTextureConverter::SettingsFile* CTextureConverter::LoadSettings(const VfsPath& path) const
{
	CXeromyces XeroFile;
	if (XeroFile.Load(m_VFS, path, "texture") != PSRETURN_OK)
		return NULL;

	// Define all the elements used in the XML file
	#define EL(x) int el_##x = XeroFile.GetElementID(#x)
	#define AT(x) int at_##x = XeroFile.GetAttributeID(#x)
	EL(textures);
	EL(file);
	AT(pattern);
	AT(format);
	AT(mipmap);
	AT(normal);
	AT(alpha);
	AT(filter);
	AT(kaiserwidth);
	AT(kaiseralpha);
	AT(kaiserstretch);
	#undef AT
	#undef EL

	XMBElement root = XeroFile.GetRoot();

	if (root.GetNodeName() != el_textures)
	{
		LOGERROR("Invalid texture settings file \"%s\" (unrecognised root element)", path.string8());
		return NULL;
	}

	std::unique_ptr<SettingsFile> settings = std::make_unique<SettingsFile>();

	XERO_ITER_EL(root, child)
	{
		if (child.GetNodeName() == el_file)
		{
			Match p;

			XERO_ITER_ATTR(child, attr)
			{
				if (attr.Name == at_pattern)
				{
					p.pattern = attr.Value.FromUTF8();
				}
				else if (attr.Name == at_format)
				{
					CStr v(attr.Value);
					if (v == "dxt1")
						p.settings.format = FMT_DXT1;
					else if (v == "dxt3")
						p.settings.format = FMT_DXT3;
					else if (v == "dxt5")
						p.settings.format = FMT_DXT5;
					else if (v == "rgba")
						p.settings.format = FMT_RGBA;
					else if (v == "alpha")
						p.settings.format = FMT_ALPHA;
					else
						LOGERROR("Invalid attribute value <file format='%s'>", v.c_str());
				}
				else if (attr.Name == at_mipmap)
				{
					CStr v(attr.Value);
					if (v == "true")
						p.settings.mipmap = MIP_TRUE;
					else if (v == "false")
						p.settings.mipmap = MIP_FALSE;
					else
						LOGERROR("Invalid attribute value <file mipmap='%s'>", v.c_str());
				}
				else if (attr.Name == at_normal)
				{
					CStr v(attr.Value);
					if (v == "true")
						p.settings.normal = NORMAL_TRUE;
					else if (v == "false")
						p.settings.normal = NORMAL_FALSE;
					else
						LOGERROR("Invalid attribute value <file normal='%s'>", v.c_str());
				}
				else if (attr.Name == at_alpha)
				{
					CStr v(attr.Value);
					if (v == "none")
						p.settings.alpha = ALPHA_NONE;
					else if (v == "player")
						p.settings.alpha = ALPHA_PLAYER;
					else if (v == "transparency")
						p.settings.alpha = ALPHA_TRANSPARENCY;
					else
						LOGERROR("Invalid attribute value <file alpha='%s'>", v.c_str());
				}
				else if (attr.Name == at_filter)
				{
					CStr v(attr.Value);
					if (v == "box")
						p.settings.filter = FILTER_BOX;
					else if (v == "triangle")
						p.settings.filter = FILTER_TRIANGLE;
					else if (v == "kaiser")
						p.settings.filter = FILTER_KAISER;
					else
						LOGERROR("Invalid attribute value <file filter='%s'>", v.c_str());
				}
				else if (attr.Name == at_kaiserwidth)
				{
					p.settings.kaiserWidth = CStr(attr.Value).ToFloat();
				}
				else if (attr.Name == at_kaiseralpha)
				{
					p.settings.kaiserAlpha = CStr(attr.Value).ToFloat();
				}
				else if (attr.Name == at_kaiserstretch)
				{
					p.settings.kaiserStretch = CStr(attr.Value).ToFloat();
				}
				else
				{
					LOGERROR("Invalid attribute name <file %s='...'>", XeroFile.GetAttributeString(attr.Name));
				}
			}

			settings->patterns.push_back(p);
		}
	}

	return settings.release();
}

CTextureConverter::Settings CTextureConverter::ComputeSettings(const std::wstring& filename, const std::vector<SettingsFile*>& settingsFiles) const
{
	// Set sensible defaults
	Settings settings;
	settings.format = FMT_DXT1;
	settings.mipmap = MIP_TRUE;
	settings.normal = NORMAL_FALSE;
	settings.alpha = ALPHA_NONE;
	settings.filter = FILTER_BOX;
	settings.kaiserWidth = 3.f;
	settings.kaiserAlpha = 4.f;
	settings.kaiserStretch = 1.f;

	for (size_t i = 0; i < settingsFiles.size(); ++i)
	{
		for (size_t j = 0; j < settingsFiles[i]->patterns.size(); ++j)
		{
			Match p = settingsFiles[i]->patterns[j];

			// Check that the pattern matches the texture file
			if (!match_wildcard(filename.c_str(), p.pattern.c_str()))
				continue;

			if (p.settings.format != FMT_UNSPECIFIED)
				settings.format = p.settings.format;

			if (p.settings.mipmap != MIP_UNSPECIFIED)
				settings.mipmap = p.settings.mipmap;

			if (p.settings.normal != NORMAL_UNSPECIFIED)
				settings.normal = p.settings.normal;

			if (p.settings.alpha != ALPHA_UNSPECIFIED)
				settings.alpha = p.settings.alpha;

			if (p.settings.filter != FILTER_UNSPECIFIED)
				settings.filter = p.settings.filter;

			if (p.settings.kaiserWidth != -1.f)
				settings.kaiserWidth = p.settings.kaiserWidth;

			if (p.settings.kaiserAlpha != -1.f)
				settings.kaiserAlpha = p.settings.kaiserAlpha;

			if (p.settings.kaiserStretch != -1.f)
				settings.kaiserStretch = p.settings.kaiserStretch;
		}
	}

	return settings;
}

CTextureConverter::CTextureConverter(PIVFS vfs, bool highQuality) :
	m_VFS(vfs), m_HighQuality(highQuality)
{
#if CONFIG2_NVTT
	// Verify that we are running with at least the version we were compiled with,
	// to avoid bugs caused by ABI changes
	ENSURE(nvtt::version() >= NVTT_VERSION);
#endif // CONFIG2_NVTT
}

CTextureConverter::~CTextureConverter() = default;

bool CTextureConverter::ConvertTexture(const CTexturePtr& texture, const VfsPath& src, const VfsPath& dest, const Settings& settings)
{
	std::shared_ptr<u8> file;
	size_t fileSize;
	if (m_VFS->LoadFile(src, file, fileSize) < 0)
	{
		LOGERROR("Failed to load texture \"%s\"", src.string8());
		return false;
	}

	Tex tex;
	const Status decodeStatus = tex.decode(file, fileSize);
	if (decodeStatus != INFO::OK)
	{
		LOGERROR("Failed to decode texture \"%s\" %s", src.string8(), GetStatusAsString(decodeStatus).c_str());
		return false;
	}

	if (!is_pow2(tex.m_Width) || !is_pow2(tex.m_Height))
	{
		LOGERROR("Texture to convert \"%s\" should have width and height be power of two: %zux%zu",
			src.string8(), tex.m_Width, tex.m_Height);
		return false;
	}

	// Check whether there's any alpha channel
	bool hasAlpha = ((tex.m_Flags & TEX_ALPHA) != 0);

	if (settings.format == FMT_ALPHA)
	{
		// Convert to uncompressed 8-bit with no mipmaps
		if (tex.transform_to((tex.m_Flags | TEX_GREY) & ~(TEX_DXT | TEX_MIPMAPS | TEX_ALPHA)) < 0)
		{
			LOGERROR("Failed to transform texture \"%s\"", src.string8());
			return false;
		}
	}
	else
	{
		// Convert to uncompressed BGRA with no mipmaps
		if (tex.transform_to((tex.m_Flags | TEX_BGR | TEX_ALPHA) & ~(TEX_DXT | TEX_MIPMAPS)) < 0)
		{
			LOGERROR("Failed to transform texture \"%s\"", src.string8());
			return false;
		}
	}

	// Check if the texture has all alpha=255, so we can automatically
	// switch from DXT3/DXT5 to DXT1 with no loss
	if (hasAlpha)
	{
		hasAlpha = false;
		u8* data = tex.get_data();
		for (size_t i = 0; i < tex.m_Width * tex.m_Height; ++i)
		{
			if (data[i*4+3] != 0xFF)
			{
				hasAlpha = true;
				break;
			}
		}
	}

#if CONFIG2_NVTT

	std::unique_ptr<ConversionRequest> request = std::make_unique<ConversionRequest>();
	request->dest = dest;
	request->texture = texture;

	// Apply the chosen settings:

	request->inputOptions.setMipmapGeneration(settings.mipmap == MIP_TRUE);

	if (settings.alpha == ALPHA_TRANSPARENCY)
		request->inputOptions.setAlphaMode(nvtt::AlphaMode_Transparency);
	else
		request->inputOptions.setAlphaMode(nvtt::AlphaMode_None);

	if (settings.format == FMT_RGBA)
	{
		request->compressionOptions.setFormat(nvtt::Format_RGBA);
		// Change the default component order (see tex_dds.cpp decode_pf)
		request->compressionOptions.setPixelFormat(32, 0xFF, 0xFF00, 0xFF0000, 0xFF000000u);
	}
	else if (settings.format == FMT_ALPHA)
	{
		request->compressionOptions.setFormat(nvtt::Format_RGBA);
		request->compressionOptions.setPixelFormat(8, 0x00, 0x00, 0x00, 0xFF);
	}
	else if (!hasAlpha)
	{
		// if no alpha channel then there's no point using DXT3 or DXT5
		request->compressionOptions.setFormat(nvtt::Format_DXT1);
	}
	else if (settings.format == FMT_DXT1)
	{
		request->compressionOptions.setFormat(nvtt::Format_DXT1a);
	}
	else if (settings.format == FMT_DXT3)
	{
		request->compressionOptions.setFormat(nvtt::Format_DXT3);
	}
	else if (settings.format == FMT_DXT5)
	{
		request->compressionOptions.setFormat(nvtt::Format_DXT5);
	}

	if (settings.filter == FILTER_BOX)
		request->inputOptions.setMipmapFilter(nvtt::MipmapFilter_Box);
	else if (settings.filter == FILTER_TRIANGLE)
		request->inputOptions.setMipmapFilter(nvtt::MipmapFilter_Triangle);
	else if (settings.filter == FILTER_KAISER)
		request->inputOptions.setMipmapFilter(nvtt::MipmapFilter_Kaiser);

	if (settings.normal == NORMAL_TRUE)
		request->inputOptions.setNormalMap(true);

	request->inputOptions.setKaiserParameters(settings.kaiserWidth, settings.kaiserAlpha, settings.kaiserStretch);

	request->inputOptions.setWrapMode(nvtt::WrapMode_Mirror); // TODO: should this be configurable?

	request->compressionOptions.setQuality(m_HighQuality ? nvtt::Quality_Production : nvtt::Quality_Fastest);

	// TODO: normal maps, gamma, etc

	// Load the texture data
	request->inputOptions.setTextureLayout(nvtt::TextureType_2D, tex.m_Width, tex.m_Height);
	if (tex.m_Bpp == 32)
	{
		request->inputOptions.setMipmapData(tex.get_data(), tex.m_Width, tex.m_Height);
	}
	else // bpp == 8
	{
		// NVTT requires 32-bit input data, so convert
		const u8* input = tex.get_data();
		u8* rgba = new u8[tex.m_Width * tex.m_Height * 4];
		u8* p = rgba;
		for (size_t i = 0; i < tex.m_Width * tex.m_Height; i++)
		{
			p[0] = p[1] = p[2] = p[3] = *input++;
			p += 4;
		}
		request->inputOptions.setMipmapData(rgba, tex.m_Width, tex.m_Height);
		delete[] rgba;
	}

	m_ResultQueue.push(Threading::TaskManager::Instance().PushTask([request = std::move(request)]
		{
			PROFILE2("compress");
			// Set up the result object
			std::unique_ptr<ConversionResult> result = std::make_unique<ConversionResult>();
			result->dest = request->dest;
			result->texture = request->texture;

			request->outputOptions.setOutputHandler(&result->output);

			// Perform the compression
			nvtt::Compressor compressor;
			result->ret = compressor.process(request->inputOptions, request->compressionOptions,
				request->outputOptions);

			return result;
		}, Threading::TaskPriority::LOW));

	return true;

#else // CONFIG2_NVTT
	LOGERROR("Failed to convert texture \"%s\" (NVTT not available)", src.string8());
	return false;
#endif // !CONFIG2_NVTT
}

bool CTextureConverter::Poll(CTexturePtr& texture, VfsPath& dest, bool& ok)
{
#if CONFIG2_NVTT
	if (m_ResultQueue.empty() || !m_ResultQueue.front().IsReady())
	{
		// no work to do
		return false;
	}

	std::unique_ptr<ConversionResult> result = m_ResultQueue.front().Get();
	m_ResultQueue.pop();

	if (!result->ret)
	{
		// conversion had failed
		ok = false;
		return true;
	}

	// Move output into a correctly-aligned buffer
	size_t size = result->output.buffer.size();
	std::shared_ptr<u8> file;
	AllocateAligned(file, size, maxSectorSize);
	memcpy(file.get(), &result->output.buffer[0], size);
	if (m_VFS->CreateFile(result->dest, file, size) < 0)
	{
		// error writing file
		ok = false;
		return true;
	}

	// Succeeded in converting texture
	texture = result->texture;
	dest = result->dest;
	ok = true;
	return true;

#else // CONFIG2_NVTT
	return false;
#endif // !CONFIG2_NVTT
}

bool CTextureConverter::IsBusy() const
{
#if CONFIG2_NVTT
	return m_ResultQueue.size() >= MAX_QUEUE_SIZE_FOR_OPTIMAL_UTILIZATION;
#else // CONFIG2_NVTT
	return false;
#endif // !CONFIG2_NVTT
}