// Copyright (c) 2014- PPSSPP Project.

// This program 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, version 2.0 or later versions.

// This program 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 2.0 for more details.

// A copy of the GPL 2.0 should have been included with the program.
// If not, see http://www.gnu.org/licenses/

// Official git repository and contact information can be found at
// https://github.com/hrydgard/ppsspp and http://www.ppsspp.org/.

#include <algorithm>
#include <map>
#include <vector>
#include <string>

#include "Common/StringUtils.h"
#include "GPU/GPUState.h"
#include "Common/GPU/Shader.h"
#include "Common/GPU/ShaderWriter.h"
#include "Common/Data/Convert/ColorConv.h"
#include "GPU/Common/Draw2D.h"
#include "GPU/Common/TextureShaderCommon.h"
#include "GPU/Common/DepalettizeShaderCommon.h"

static const VaryingDef varyings[1] = {
	{ "vec2", "v_texcoord", Draw::SEM_TEXCOORD0, 0, "highp" },
};

static const SamplerDef samplers[2] = {
	{ 0, "tex", SamplerFlags::ARRAY_ON_VULKAN },
	{ 1, "pal" },
};

TextureShaderCache::TextureShaderCache(Draw::DrawContext *draw, Draw2D *draw2D) : draw_(draw), draw2D_(draw2D) { }

TextureShaderCache::~TextureShaderCache() {
	DeviceLost();
}

void TextureShaderCache::DeviceRestore(Draw::DrawContext *draw) {
	draw_ = draw;
}

void TextureShaderCache::DeviceLost() {
	Clear();
	draw_ = nullptr;
}

Draw::SamplerState *TextureShaderCache::GetSampler(bool linearFilter) {
	if (linearFilter) {
		if (!linearSampler_) {
			Draw::SamplerStateDesc desc{};
			desc.magFilter = Draw::TextureFilter::LINEAR;
			desc.minFilter = Draw::TextureFilter::LINEAR;
			desc.wrapU = Draw::TextureAddressMode::CLAMP_TO_EDGE;
			desc.wrapV = Draw::TextureAddressMode::CLAMP_TO_EDGE;
			desc.wrapW = Draw::TextureAddressMode::CLAMP_TO_EDGE;
			linearSampler_ = draw_->CreateSamplerState(desc);
		}
		return linearSampler_;
	} else {
		if (!nearestSampler_) {
			Draw::SamplerStateDesc desc{};
			desc.wrapU = Draw::TextureAddressMode::CLAMP_TO_EDGE;
			desc.wrapV = Draw::TextureAddressMode::CLAMP_TO_EDGE;
			desc.wrapW = Draw::TextureAddressMode::CLAMP_TO_EDGE;
			nearestSampler_ = draw_->CreateSamplerState(desc);
		}
		return nearestSampler_;
	}
}

void TextureShaderCache::Clear() {
	for (auto shader = pipelineCache_.begin(); shader != pipelineCache_.end(); ++shader) {
		shader->second->Release();
	}
	pipelineCache_.clear();
	if (nearestSampler_) {
		nearestSampler_->Release();
		nearestSampler_ = nullptr;
	}
	if (linearSampler_) {
		linearSampler_->Release();
		linearSampler_ = nullptr;
	}
}

Draw2DPipeline *TextureShaderCache::GetDepalettizeShader(uint32_t clutMode, GETextureFormat textureFormat, GEBufferFormat bufferFormat, bool smoothedDepal, u32 depthUpperBits, bool depthIntoHalves) {
	using namespace Draw;

	// Generate an ID for depal shaders.
	u64 id = ((u64)smoothedDepal << 63) | ((u64)depthIntoHalves << 62) | ((u64)depthUpperBits << 32) | (clutMode & 0xFFFFFF) | (textureFormat << 24) | (bufferFormat << 28);

	auto shader = pipelineCache_.find(id);
	if (shader != pipelineCache_.end()) {
		return shader->second;
	}

	// TODO: Parse these out of clutMode some nice way, to become a bit more stateless.
	DepalConfig config;
	config.clutFormat = gstate.getClutPaletteFormat();
	config.startPos = gstate.getClutIndexStartPos();
	config.shift = gstate.getClutIndexShift();
	config.mask = gstate.getClutIndexMask();
	config.bufferFormat = bufferFormat;
	config.textureFormat = textureFormat;
	config.smoothedDepal = smoothedDepal;
	config.depthUpperBits = depthUpperBits;
	config.depthIntoHalves = depthIntoHalves;

	Draw2DPipeline *ts = draw2D_->Create2DPipeline([config](ShaderWriter &writer) -> Draw2DPipelineInfo {
		GenerateDepalFs(writer, config);
		return Draw2DPipelineInfo{
			"depal",
			config.bufferFormat == GE_FORMAT_DEPTH16 ? RASTER_DEPTH : RASTER_COLOR,
			RASTER_COLOR,
			samplers
		};
	});

	pipelineCache_[id] = ts;

	return ts->pipeline ? ts : nullptr;
}

std::vector<std::string> TextureShaderCache::DebugGetShaderIDs(DebugShaderType type) const {
	std::vector<std::string> ids;
	for (auto &entry : pipelineCache_) {
		ids.push_back(StringFromFormat("%016llx", (unsigned long long)entry.first));
	}
	return ids;
}

std::string TextureShaderCache::DebugGetShaderString(const std::string &idstr, DebugShaderType type, DebugShaderStringType stringType) const {
	unsigned long long id = 0;
	if (sscanf(idstr.c_str(), "%llx", &id) != 1) {
		return "";
	}
	auto iter = pipelineCache_.find(id);
	if (iter == pipelineCache_.end())
		return "";
	switch (stringType) {
	case SHADER_STRING_SHORT_DESC:
		return idstr;
	case SHADER_STRING_SOURCE_CODE:
		return iter->second->code;
	default:
		return "";
	}
}

void ClutTextureCache::DeviceRestore(Draw::DrawContext *draw) {
	draw_ = draw;
}

void ClutTextureCache::DeviceLost() {
	Clear();
	draw_ = nullptr;
}

// TODO: We could cache many CLUTs in one big texture, might be more efficient.
ClutTexture ClutTextureCache::GetClutTexture(GEPaletteFormat clutFormat, const u32 clutHash, const u32 *rawClut) {
	// Simplistic, but works well enough.
	u32 clutId = clutHash ^ (uint32_t)clutFormat;

	auto oldtex = texCache_.find(clutId);
	if (oldtex != texCache_.end()) {
		oldtex->second->lastFrame = gpuStats.totals.numFlips;
		return *oldtex->second;
	}

	int maxClutEntries = clutFormat == GE_CMODE_32BIT_ABGR8888 ? 256 : 512;

	ClutTexture *tex = new ClutTexture();

	Draw::TextureDesc desc{};
	desc.width = 512;  // We always use 512-sized textures here for simplicity, though the most common is that only up to 256 entries are used.
	desc.height = 1;
	desc.depth = 1;
	desc.mipLevels = 1;
	desc.tag = "clut";
	desc.type = Draw::TextureType::LINEAR2D;  // TODO: Try LINEAR1D?
	desc.format = Draw::DataFormat::R8G8B8A8_UNORM;  // TODO: Also support an BGR format. We won't bother with the 16-bit formats here.

	uint8_t convTemp[2048]{};

	switch (clutFormat) {
	case GEPaletteFormat::GE_CMODE_32BIT_ABGR8888:
		desc.initData.push_back((const uint8_t *)rawClut);
		break;
	case GEPaletteFormat::GE_CMODE_16BIT_BGR5650:
		ConvertRGB565ToRGBA8888((u32 *)convTemp, (const u16 *)rawClut, maxClutEntries);
		desc.initData.push_back(convTemp);
		break;
	case GEPaletteFormat::GE_CMODE_16BIT_ABGR5551:
		ConvertRGBA5551ToRGBA8888((u32 *)convTemp, (const u16 *)rawClut, maxClutEntries);
		desc.initData.push_back(convTemp);
		break;
	case GEPaletteFormat::GE_CMODE_16BIT_ABGR4444:
		ConvertRGBA4444ToRGBA8888((u32 *)convTemp, (const u16 *)rawClut, maxClutEntries);
		desc.initData.push_back(convTemp);
		break;
	}


	for (int i = 0; i < 3; i++) {
		tex->rampLengths[i] = 0;
		tex->rampStarts[i] = 0;
	}
	// Quick check for how many continuously growing entries we have at the start.
	// Bilinearly filtering CLUTs only really makes sense for this kind of ramp.
	int i = 0;
	for (int j = 0; j < ClutTexture::MAX_RAMPS; j++) {
		tex->rampStarts[j] = i;
		int lastR = 0;
		int lastG = 0;
		int lastB = 0;
		int lastA = 0;
		for (; i < maxClutEntries; i++) {
			int r = desc.initData[0][i * 4];
			int g = desc.initData[0][i * 4 + 1];
			int b = desc.initData[0][i * 4 + 2];
			int a = desc.initData[0][i * 4 + 3];
			if (r < lastR || g < lastG || b < lastB || a < lastA) {
				lastR = r; lastG = g; lastB = b; lastA = a;
				break;
			} else {
				lastR = r;
				lastG = g;
				lastB = b;
				lastA = a;
			}
		}
		tex->rampLengths[j] = i - tex->rampStarts[j];
		if (i >= maxClutEntries) {
			break;
		}
	}

	tex->texture = draw_->CreateTexture(desc);
	tex->lastFrame = gpuStats.totals.numFlips;

	texCache_[clutId] = tex;
	return *tex;
}

void ClutTextureCache::Clear() {
	for (auto tex = texCache_.begin(); tex != texCache_.end(); ++tex) {
		if (tex->second->texture) {
			tex->second->texture->Release();
		}
		delete tex->second;
	}
	texCache_.clear();
}

void ClutTextureCache::Decimate() {
	for (auto tex = texCache_.begin(); tex != texCache_.end(); ) {
		if (tex->second->lastFrame + DEPAL_TEXTURE_OLD_AGE < gpuStats.totals.numFlips) {
			if (tex->second->texture) {
				tex->second->texture->Release();
			}
			delete tex->second;
			texCache_.erase(tex++);
		} else {
			++tex;
		}
	}
	gpuStats.perFrame.numClutTextures = (int)texCache_.size();
}

const char *const textureScaleUploadShader = R"(
#version 450
#extension GL_ARB_separate_shader_objects : enable

// 8x8 is the most common compute shader workgroup size, and works great on all major
// hardware vendors. TODO: However, we should probably change to 16x16, as Qualcomm now has
// support for bigger groups...
layout (local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

uniform layout(set = 0, binding = 0, rgba8) writeonly image2D img;

layout(std430, set = 0, binding = 1) buffer Buf {
	uint data[];
} buf;

layout(push_constant) uniform Params {
	int width;
	int height;
} params;

// The cbuffer, if present, is self-declared for layout flexibility
#define CBUFFER_SET 0
#define CBUFFER_BINDING 4

uint readColoru(uvec2 p) {
	return buf.data[p.y * params.width + p.x];
}

vec4 readColorf(uvec2 p) {
	// Unpack the color (we could look it up in a CLUT here if we wanted...)
	// The imageStore repack is free.
	return unpackUnorm4x8(readColoru(p));
}

void writeColorf(ivec2 p, vec4 c) {
	imageStore(img, p, c);
}

%s

// Note that main runs once per INPUT pixel, unlike the old model.
void main() {
	uvec2 xy = gl_GlobalInvocationID.xy;
	// Kill off any out-of-image threads to avoid stray writes.
	// Should only happen on the tiniest mipmaps as PSP textures are power-of-2,
	// and we use a 8x8 workgroup size. Probably not really necessary.
	if (xy.x >= params.width || xy.y >= params.height)
		return;
	// applyScaling will write the upscaled pixels, using writeColorf above.
	// It's expected to write a square of scale*scale pixels, at the location xy*scale.
	applyScaling(xy);
}

)";

bool GetTextureScaleMultipassLayout(int scaleFactor, std::vector<TextureScaleScratchDesc> *scratchDescs, std::vector<TextureScaleStageDesc> *stageDescs) {
	if (scaleFactor == 2) {
		*scratchDescs = {
			{ "nnedi3_vertical", 1, 2 },
			{ "nnedi3_horizontal", 2, 2 },
		};
		*stageDescs = {
			{ 0, -1, 0, 1, 1, 1, 2, false },
			{ 1, 0, 1, 1, 2, 2, 2, false },
			{ 2, 1, -1, 2, 2, 2, 2, false },
		};
	} else if (scaleFactor == 4) {
		*scratchDescs = {
			{ "nnedi3_vertical", 1, 2 },
			{ "nnedi3_horizontal", 2, 2 },
			{ "nnedi3_vertical2", 2, 4 },
			{ "nnedi3_horizontal2", 4, 4 },
		};
		*stageDescs = {
			{ 0, -1, 0, 1, 1, 1, 2, false },
			{ 2, 0, 1, 1, 2, 2, 2, false },
			{ 1, 1, 2, 2, 2, 2, 4, false },
			{ 2, 2, 3, 2, 4, 4, 4, false },
			{ 3, 3, -1, 4, 4, 4, 4, true },
		};
	} else {
		return false;
	}
	return true;
}