/*
===========================================================================
Copyright (C) 2022-2024 Gian 'myT' Schellenbaum

This file is part of Challenge Quake 3 (CNQ3).

Challenge Quake 3 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.

Challenge Quake 3 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 Challenge Quake 3. If not, see <https://www.gnu.org/licenses/>.
===========================================================================
*/
// compiles core shaders as byte code to be embedded into the CNQ3 client


#include <Windows.h>
#include <stdarg.h>
#include <stdio.h>
#include <assert.h>
#include <shlwapi.h>
#include <sal.h>

#include "../renderer/grp_uber_shaders.h"


char repoPath[MAX_PATH];
char outputPath[MAX_PATH];
char bin2headerPath[MAX_PATH];
char dxcPath[MAX_PATH];

const char* targetVS = "vs_6_0";
const char* targetPS = "ps_6_0";
const char* targetCS = "cs_6_0";


#define PS(Data) #Data,
const char* uberShaderPixelStates[] =
{
	UBER_SHADER_PS_LIST(PS)
};
#undef PS


// -Zi embeds debug info
// -Qembed_debug embeds debug info in shader container
// -Vn header variable name
// -WX warnings as errors
// -O3 or -O0 optimization level
// -Wno-warning disables the warning


const char* va(_Printf_format_string_ const char* format, ...)
{
	static char string[64][32000];
	static int index = 0;
	char* buf = string[index++ & 63];
	va_list argptr;

	va_start(argptr, format);
	vsprintf(buf, format, argptr);
	va_end(argptr);

	return buf;
}


struct ShaderArgs
{
	const char* headerPath;
	const char* shaderPath;
	const char* entryPoint;
	const char* targetProfile;
};

const char* OutputPath(const char* path)
{
	return va("%s\\%s", outputPath, path);
}

const char* HeaderVariable(const char* name)
{
	// the output variable isn't declared static, polluting the global namespace...
	return va("\"static %s\"", name);
}

void CompileShader(const ShaderArgs& args, int extraCount = 0, const char** extras = NULL)
{
	static char temp[4096];

	const char* headerPath = va("%s\\%s", outputPath, args.headerPath);

	// -Ges: Enable strict mode
	// -Gis: Force IEEE strictness
	// -Zi: Embed debug info
	// -Qembed_debug: Embed debug info in shader container
	strcpy(temp, va("%s -Fh %s -E %s -T %s -WX -Ges -Gis -Zi -Qembed_debug",
		dxcPath, headerPath, args.entryPoint, args.targetProfile));

	for(int i = 0; i < extraCount; ++i)
	{
		strcat(temp, " ");
		strcat(temp, extras[i]);
	}

	strcat(temp, " ");
	strcat(temp, args.shaderPath);

	printf("%s\n", temp);
	system(temp);
}

struct SMAAArgs
{
	const char* headerPath;
	const char* shaderPath;
	const char* presetMacro;
	const char* variableName;
	bool vertexShader;
};

void CompileSMAAShader(const SMAAArgs& smaaArgs)
{
	const char* extras[] =
	{
		"-Vn", HeaderVariable(smaaArgs.variableName),
		smaaArgs.presetMacro,
		smaaArgs.vertexShader ? "-D SMAA_INCLUDE_VS=1" : "-D SMAA_INCLUDE_PS=1",
		"-D SMAA_HLSL_5_1=1",
		"-D SMAA_RT_METRICS=rtMetrics"
	};

	ShaderArgs args;
	args.entryPoint = smaaArgs.vertexShader ? "vs" : "ps";
	args.headerPath = smaaArgs.headerPath;
	args.shaderPath = smaaArgs.shaderPath;
	args.targetProfile = smaaArgs.vertexShader ? targetVS : targetPS;
	CompileShader(args, _countof(extras), extras);
}

void ProcessSMAAShadersForPreset(const char* presetName, const char* presetMacro)
{
	SMAAArgs args;
	args.presetMacro = presetMacro;
	for(int pass = 0; pass < 3; ++pass)
	{
		for(int ps = 0; ps < 2; ++ps)
		{
			args.headerPath = va("smaa_%s_%d_%s.h", presetName, pass + 1, ps ? "ps" : "vs");
			args.shaderPath = va("smaa_%d.hlsl", pass + 1);
			args.variableName = va("%s_%d_%s", presetName, pass + 1, ps ? "ps" : "vs");
			args.vertexShader = ps == 0;
			CompileSMAAShader(args);
		}
	}
}

void CompileSMAAShaders()
{
	ProcessSMAAShadersForPreset("low", "-D SMAA_PRESET_LOW=1");
	ProcessSMAAShadersForPreset("medium", "-D SMAA_PRESET_MEDIUM=1");
	ProcessSMAAShadersForPreset("high", "-D SMAA_PRESET_HIGH=1");
	ProcessSMAAShadersForPreset("ultra", "-D SMAA_PRESET_ULTRA=1");
}

void CompileGraphics(const char* headerPath, const char* shaderPath, const char* varName)
{
	const char* vsHeaderRelPath = va("%s.vs.h", shaderPath);
	const char* psHeaderRelPath = va("%s.ps.h", shaderPath);
	const char* vsHeaderPath = OutputPath(vsHeaderRelPath);
	const char* psHeaderPath = OutputPath(psHeaderRelPath);
	const char* vsExtras[] =
	{
		"-D", "VERTEX_SHADER=1",
		"-Vn", HeaderVariable(va("g_%s_vs", varName))
	};
	const char* psExtras[] =
	{
		"-D", "PIXEL_SHADER=1",
		"-Vn", HeaderVariable(va("g_%s_ps", varName))
	};

	ShaderArgs args;
	args.entryPoint = "vs";
	args.headerPath = vsHeaderRelPath;
	args.shaderPath = shaderPath;
	args.targetProfile = targetVS;
	CompileShader(args, _countof(vsExtras), vsExtras);

	args.entryPoint = "ps";
	args.headerPath = psHeaderRelPath;
	args.shaderPath = shaderPath;
	args.targetProfile = targetPS;
	CompileShader(args, _countof(psExtras), psExtras);
	
	const char* outHeaderPath = OutputPath(headerPath);
	system(va("type %s %s > %s", vsHeaderPath, psHeaderPath, outHeaderPath));
	system(va("del %s", vsHeaderPath));
	system(va("del %s", psHeaderPath));
}

void CompileVertexShader(const char* headerPath, const char* shaderPath, const char* varName)
{
	const char* extras[] =
	{
		"-D", "VERTEX_SHADER=1",
		"-Vn", HeaderVariable(va("g_%s_vs", varName))
	};

	ShaderArgs args;
	args.entryPoint = "vs";
	args.headerPath = headerPath;
	args.shaderPath = shaderPath;
	args.targetProfile = targetVS;
	CompileShader(args, _countof(extras), extras);
}

void CompilePixelShader(const char* headerPath, const char* shaderPath, const char* varName)
{
	const char* extras[] =
	{
		"-D", "PIXEL_SHADER=1",
		"-Vn", HeaderVariable(va("g_%s_ps", varName))
	};

	ShaderArgs args;
	args.entryPoint = "ps";
	args.headerPath = headerPath;
	args.shaderPath = shaderPath;
	args.targetProfile = targetPS;
	CompileShader(args, _countof(extras), extras);
}

void CompileCompute(const char* headerPath, const char* shaderPath, const char* varName)
{
	const char* extras[] =
	{
		"-D", "COMPUTE_SHADER=1",
		"-Vn", HeaderVariable(va("g_%s_cs", varName))
	};

	ShaderArgs args;
	args.entryPoint = "cs";
	args.headerPath = headerPath;
	args.shaderPath = shaderPath;
	args.targetProfile = targetCS;
	CompileShader(args, _countof(extras), extras);
}

void CompileUberVS(const char* headerPath, const char* shaderPath, int stageCount)
{
	const char* extras[] =
	{
		"-D", "VERTEX_SHADER=1",
		"-D", va("STAGE_COUNT=%d", stageCount),
		"-Vn", HeaderVariable(va("g_vs_%d", stageCount))
	};

	ShaderArgs args;
	args.entryPoint = "vs";
	args.headerPath = headerPath;
	args.shaderPath = shaderPath;
	args.targetProfile = targetVS;
	CompileShader(args, _countof(extras), extras);
}

void CompileUberPS(const char* stateString)
{
	UberPixelShaderState state;
	if(!ParseUberPixelShaderState(state, stateString))
	{
		fprintf(stderr, "ParseUberPixelShaderState failed!\n");
		exit(666);
	}

	const char* extras[16];
	int extraCount = 0;
	extras[extraCount++] = va("-Vn %s", HeaderVariable(va("g_ps_%s", stateString)));
	extras[extraCount++] = "-D USE_INCLUDES=1";
	extras[extraCount++] = "-D PIXEL_SHADER=1";
	if(state.globalState & UBERPS_DITHER_BIT)
	{
		extras[extraCount++] = "-D DITHER=1";
	}
	if(state.globalState & UBERPS_DEPTHFADE_BIT)
	{
		extras[extraCount++] = "-D DEPTH_FADE=1";
	}
	extras[extraCount++] = va("-D STAGE_COUNT=%d", state.stageCount);
	for(int s = 0; s < state.stageCount; ++s)
	{
		extras[extraCount++] = va("-D STAGE%d_BITS=0x%X", s, state.stageStates[s]);
	}

	ShaderArgs args;
	args.entryPoint = "ps";
	args.headerPath = va("uber_shader_ps_%s.h", stateString);
	args.shaderPath = "uber_shader.hlsl";
	args.targetProfile = targetPS;
	CompileShader(args, extraCount, extras);
}

const char* Canonicalize(const char* path)
{
	static char canonPath[MAX_PATH];

	PathCanonicalizeA(canonPath, path);

	return canonPath;
}

void InitDirectory(const char* dirName)
{
	const char* rendererPath = va("%s\\code\\renderer", repoPath);
	const char* cd = Canonicalize(va("%s\\shaders\\%s", rendererPath, dirName));
	SetCurrentDirectoryA(cd);
	const char* out = Canonicalize(va("%s\\compshaders\\%s", rendererPath, dirName));
	strcpy(outputPath, out);
	CreateDirectoryA(outputPath, NULL);

	system(va("del %s\\*.h", outputPath));
	system(va("del %s\\*.temp", outputPath));
}

void ProcessGRP()
{
	InitDirectory("grp");
	targetVS = "vs_6_0";
	targetPS = "ps_6_0";
	targetCS = "cs_6_0";

	CompileGraphics("post_gamma.h", "post_gamma.hlsl", "post");
	CompileGraphics("post_inverse_gamma.h", "post_inverse_gamma.hlsl", "post_inverse");
	CompileGraphics("imgui.h", "imgui.hlsl", "imgui");
	CompileGraphics("nuklear.h", "nuklear.hlsl", "nuklear");
	CompileGraphics("ui.h", "ui.hlsl", "ui");
	CompileGraphics("depth_pre_pass.h", "depth_pre_pass.hlsl", "zpp");
	CompileGraphics("dynamic_light.h", "dynamic_light.hlsl", "dl");
	CompileVertexShader("fog.h", "fog_inside.hlsl", "fog");
	CompilePixelShader("fog_inside.h", "fog_inside.hlsl", "fog_inside");
	CompilePixelShader("fog_outside.h", "fog_outside.hlsl", "fog_outside");
	CompileCompute("mip_1.h", "mip_1.hlsl", "mip_1");
	CompileCompute("mip_2.h", "mip_2.hlsl", "mip_2");
	CompileCompute("mip_3.h", "mip_3.hlsl", "mip_3");

	CompileSMAAShaders();
	system(va("type %s\\smaa*.h > %s\\complete_smaa.h", outputPath, outputPath));

	// type combines all files into one
	system(va("type ..\\common\\state_bits.h.hlsli ..\\common\\blend.hlsli shared.hlsli uber_shader.hlsl > %s\\uber_shader.temp", outputPath));
	system(va("%s --output %s\\uber_shader.h --hname uber_shader_string %s\\uber_shader.temp", bin2headerPath, outputPath, outputPath));
	system(va("del %s\\uber_shader.temp", outputPath));

	for(int i = 0; i < 8; ++i)
	{
		CompileUberVS(va("uber_shader_vs_%i.h", i + 1), "uber_shader.hlsl", i + 1);
	}
	system(va("type %s\\uber_shader_vs_*.h > %s\\complete_uber_vs.h", outputPath, outputPath));
	system(va("del %s\\uber_shader_vs_*.h", outputPath));

	for(int i = 0; i < _countof(uberShaderPixelStates); ++i)
	{
		CompileUberPS(uberShaderPixelStates[i]);
	}
	system(va("type %s\\uber_shader_ps_*.h > %s\\complete_uber_ps.h", outputPath, outputPath));
	system(va("del %s\\uber_shader_ps_*.h", outputPath));
}

void ProcessCRP()
{
	InitDirectory("crp");
	targetVS = "vs_6_6";
	targetPS = "ps_6_6";
	targetCS = "cs_6_6";

	CompileVertexShader("fullscreen.h", "fullscreen.hlsl", "fullscreen");
	CompilePixelShader("blit.h", "blit.hlsl", "blit");
	CompileGraphics("ui.h", "ui.hlsl", "ui");
	CompileGraphics("imgui.h", "imgui.hlsl", "imgui");
	CompileGraphics("nuklear.h", "nuklear.hlsl", "nuklear");
	CompileCompute("mip_1.h", "mip_1.hlsl", "mip_1");
	CompileCompute("mip_2.h", "mip_2.hlsl", "mip_2");
	CompileCompute("mip_3.h", "mip_3.hlsl", "mip_3");
	CompileGraphics("prepass.h", "prepass.hlsl", "prepass");
	CompileGraphics("opaque.h", "opaque.hlsl", "opaque");
	CompileGraphics("transp_draw.h", "transp_draw.hlsl", "transp_draw");
	CompilePixelShader("transp_resolve.h", "transp_resolve.hlsl", "transp_resolve");
	CompilePixelShader("tone_map.h", "tone_map.hlsl", "tone_map");
	CompilePixelShader("tone_map_inverse.h", "tone_map_inverse.hlsl", "tone_map_inverse");
	CompilePixelShader("accumdof_accum.h", "accumdof_accum.hlsl", "accum");
	CompilePixelShader("accumdof_norm.h", "accumdof_norm.hlsl", "norm");
	CompilePixelShader("accumdof_debug.h", "accumdof_debug.hlsl", "debug");
	CompileCompute("gatherdof_split.h", "gatherdof_split.hlsl", "split");
	CompileCompute("gatherdof_coc_tile_gen.h", "gatherdof_coc_tile_gen.hlsl", "coc_tile_gen");
	CompileCompute("gatherdof_coc_tile_max.h", "gatherdof_coc_tile_max.hlsl", "coc_tile_max");
	CompileCompute("gatherdof_blur.h", "gatherdof_blur.hlsl", "blur");
	CompileCompute("gatherdof_fill.h", "gatherdof_fill.hlsl", "fill");
	CompilePixelShader("gatherdof_combine.h", "gatherdof_combine.hlsl", "combine");
	CompilePixelShader("gatherdof_debug.h", "gatherdof_debug.hlsl", "debug");
	CompileGraphics("fog_inside.h", "fog_inside.hlsl", "inside");
	CompileGraphics("fog_outside.h", "fog_outside.hlsl", "outside");
	CompilePixelShader("magnifier.h", "magnifier.hlsl", "magnifier");
}

int main(int /*argc*/, const char** argv)
{
	char dirPath[MAX_PATH];
	strcpy(dirPath, argv[0]);
	int l = strlen(dirPath);
	while(l-- > 0)
	{
		if(dirPath[l] == '/' || dirPath[l] == '\\')
		{
			dirPath[l] = '\0';
			break;
		}
	}
	strcpy(repoPath, Canonicalize(va("%s\\..\\..", dirPath)));
	strcpy(bin2headerPath, Canonicalize(va("%s\\tools\\bin2header.exe", repoPath)));

	char* path = getenv("DXCPATH");
	if(path != NULL)
	{
		strcpy(dxcPath, path);
	}
	else
	{
		strcpy(dxcPath, "dxc.exe");
	}
	system(va("%s --version", dxcPath));

	ProcessGRP();
	ProcessCRP();

	return 0;
}