#define _CRT_SECURE_NO_DEPRECATE

// Builds all possible shader files from ps2hw.fx and stores them in
// a preprocessed database
#include <stdio.h>
#include <assert.h>

#include <zlib.h>
#include "zpipe.h"

#include <Cg/cg.h>
#include <Cg/cgGL.h>

#define SAFE_RELEASE(x) { if( (x) != NULL ) { (x)->Release(); x = NULL; } }

#include <map>
#include <vector>

using namespace std;

#include "zerogsshaders.h"

char* srcfilename = "ps2hw.fx";
char* dstfilename = "ps2hw.dat";

#ifndef ARRAYSIZE
#define ARRAYSIZE(x) (sizeof(x) / sizeof((x)[0]))
#endif

struct SHADERINFO
{
	int type;
	vector<char> buf;
};

map<int, SHADERINFO> mapShaders;
CGcontext g_cgcontext;

void LoadShader(int index, const char* pshader, CGprofile prof, vector<const char*>& vargs, int context)
{
	vector<const char*> realargs;
	realargs.reserve(16);
	realargs.resize(vargs.size());
	if( vargs.size() > 0 )
		memcpy(&realargs[0], &vargs[0], realargs.size() * sizeof(realargs[0]));
	realargs.push_back(context ? "-Ictx1" : "-Ictx0");
	realargs.push_back(NULL);

	CGprogram prog = cgCreateProgramFromFile(g_cgcontext, CG_SOURCE, srcfilename, prof, pshader, &realargs[0]);
	if( !cgIsProgram(prog) ) {
		printf("Failed to load shader %s: \n%s\n", pshader, cgGetLastListing(g_cgcontext));
		return;
	}

	if( mapShaders.find(index) != mapShaders.end() ) {
		printf("error: two shaders share the same index %d\n", index);
		exit(0);
	}

	if( !cgIsProgramCompiled(prog) )
		cgCompileProgram(prog);
	
	const char* pstr = cgGetProgramString(prog, CG_COMPILED_PROGRAM);

	const char* pprog = strstr(pstr, "#program");
	if( pprog == NULL ) {
		printf("program field not found!\n");
		return;
	}
	pprog += 9;
	const char* progend = strchr(pprog, '\r');
	if( progend == NULL ) progend = strchr(pprog, '\n');

	if( progend == NULL ) {
		printf("prog end not found!\n");
		return;
	}

	const char* defname = "main";

	SHADERINFO info;
	info.type = 0;
	info.buf.resize(strlen(pstr)+1);
	
	// change the program name to main
	memset(&info.buf[0], 0, info.buf.size());
	memcpy(&info.buf[0], pstr, pprog-pstr);
	memcpy(&info.buf[pprog-pstr], defname, 4);
	memcpy(&info.buf[pprog-pstr+4], progend, strlen(pstr)-(progend-pstr));

	if( mapShaders.find(index) != mapShaders.end() )
		printf("same shader\n");
	assert( mapShaders.find(index) == mapShaders.end() );
	mapShaders[index] = info;

	cgDestroyProgram(prog);
}

int main(int argc, char** argv)
{
	printf("usage: [src] [dst] [opts]\n");

	if( argc >= 2 ) srcfilename = argv[1];
	if( argc >= 3 ) dstfilename = argv[2];

	FILE* fsrc = fopen(srcfilename, "r");
	if( fsrc == NULL ) {
		printf("cannot open %s\n", srcfilename);
		return 0;
	}
	fclose(fsrc);
	
	g_cgcontext = cgCreateContext();
	if( !cgIsContext(g_cgcontext) ) {
		printf("failed to create cg context\n");
		return -1;
	}

	CGprofile cgvProf = CG_PROFILE_ARBVP1;
	CGprofile cgfProf = CG_PROFILE_ARBFP1;
	if( !cgGLIsProfileSupported(cgvProf) != CG_TRUE ) {
		printf("arbvp1 not supported\n");
		return 0;
	}
	if( !cgGLIsProfileSupported(cgfProf) != CG_TRUE ) {
		printf("arbfp1 not supported\n");
		return 0;
	}

	cgGLEnableProfile(cgvProf);
	cgGLEnableProfile(cgfProf);
	cgGLSetOptimalOptions(cgvProf);
	cgGLSetOptimalOptions(cgfProf);

	vector<const char*> vmacros;

	LoadShader(SH_BITBLTVS, "BitBltVS", cgvProf, vmacros, 0);
	LoadShader(SH_BITBLTPS, "BitBltPS", cgfProf, vmacros, 0);
	LoadShader(SH_BITBLTDEPTHPS, "BitBltDepthPS", cgfProf, vmacros, 0);
	LoadShader(SH_BITBLTDEPTHMRTPS, "BitBltDepthMRTPS", cgfProf, vmacros, 0);
	LoadShader(SH_CRTCTARGPS, "CRTCTargPS", cgfProf, vmacros, 0);
	LoadShader(SH_CRTCPS, "CRTCPS", cgfProf, vmacros, 0);
	LoadShader(SH_CRTC_NEARESTPS, "CRTCPS_Nearest", cgfProf, vmacros, 0);
	LoadShader(SH_CRTC24PS, "CRTC24PS", cgfProf, vmacros, 0);
	LoadShader(SH_ZEROPS, "ZeroPS", cgfProf, vmacros, 0);
	LoadShader(SH_BASETEXTUREPS, "BaseTexturePS", cgfProf, vmacros, 0);
	LoadShader(SH_BITBLTAAPS, "BitBltPS", cgfProf, vmacros, 0);
	LoadShader(SH_CRTCTARGINTERPS, "CRTCTargInterPS", cgfProf, vmacros, 0);
	LoadShader(SH_CRTCINTERPS, "CRTCInterPS", cgfProf, vmacros, 0);
	LoadShader(SH_CRTCINTER_NEARESTPS, "CRTCInterPS_Nearest", cgfProf, vmacros, 0);
	LoadShader(SH_CRTC24INTERPS, "CRTC24InterPS", cgfProf, vmacros, 0);
	LoadShader(SH_CONVERT16TO32PS, "Convert16to32PS", cgfProf, vmacros, 0);
	LoadShader(SH_CONVERT32TO16PS, "Convert32to16PS", cgfProf, vmacros, 0);

	const int vsshaders[4] = { SH_REGULARVS, SH_TEXTUREVS, SH_REGULARFOGVS, SH_TEXTUREFOGVS };
	const char* pvsshaders[4] = { "RegularVS", "TextureVS", "RegularFogVS", "TextureFogVS" };

	// load the texture shaders
	char str[255], strdir[255];

	strcpy(strdir, srcfilename);
	int i = (int)strlen(strdir);
	while(i > 0) {
		if( strdir[i-1] == '/' || strdir[i-1] == '\\' )
			break;
		--i;
	}

	strdir[i] = 0;

	for(i = 0; i < ARRAYSIZE(vsshaders); ++i) {
		for(int writedepth = 0; writedepth < 2; ++writedepth ) {

			if( writedepth ) vmacros.push_back("-DWRITE_DEPTH");
			LoadShader(vsshaders[i]|(writedepth?SH_WRITEDEPTH:0), pvsshaders[i], cgvProf, vmacros, 0);
			LoadShader(vsshaders[i]|(writedepth?SH_WRITEDEPTH:0)|SH_CONTEXT1, pvsshaders[i], cgvProf, vmacros, 1);
			if( writedepth ) vmacros.pop_back();
		}
	}

	const int psshaders[2] = { SH_REGULARPS, SH_REGULARFOGPS };
	const char* ppsshaders[2] = { "RegularPS", "RegularFogPS" };

	for(i = 0; i < ARRAYSIZE(psshaders); ++i) {
		for(int writedepth = 0; writedepth < 2; ++writedepth ) {
			if( writedepth ) vmacros.push_back("-DWRITE_DEPTH");
			LoadShader(psshaders[i]|(writedepth?SH_WRITEDEPTH:0), ppsshaders[i], cgfProf, vmacros, 0);
			if( writedepth ) vmacros.pop_back();
		}
	}

	printf("creating shaders, note that ctx0/ps2hw_ctx.fx, and ctx1/ps2hw_ctx.fx are required\n");
	vmacros.resize(0);
		
	for(int texwrap = 0; texwrap < NUM_TEXWRAPS; ++texwrap ) {

		if( g_pPsTexWrap[texwrap] != NULL )
			vmacros.push_back(g_pPsTexWrap[texwrap]);

		for(int context = 0; context < 2; ++context) {
			
			for(int texfilter = 0; texfilter < NUM_FILTERS; ++texfilter) {
				for(int fog = 0; fog < 2; ++fog ) {
					for(int writedepth = 0; writedepth < 2; ++writedepth ) {

						if( writedepth )
							vmacros.push_back("-DWRITE_DEPTH");

						for(int testaem = 0; testaem < 2; ++testaem ) {

							if( testaem )
								vmacros.push_back("-DTEST_AEM");

							for(int exactcolor = 0; exactcolor < 2; ++exactcolor ) {
								
								if( exactcolor )
									vmacros.push_back("-DEXACT_COLOR");

								// 32
								sprintf(str, "Texture%s%d_32PS", fog?"Fog":"", texfilter);

								vmacros.push_back("-DACCURATE_DECOMPRESSION");
								LoadShader(GET_SHADER_INDEX(0, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, SHADER_ACCURATE), str, cgfProf, vmacros, context);
								vmacros.pop_back();

								LoadShader(GET_SHADER_INDEX(0, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, 0), str, cgfProf, vmacros, context);

								if( texfilter == 0 ) {
									// tex32
									sprintf(str, "Texture%s%d_tex32PS", fog?"Fog":"", texfilter);

//									vmacros.push_back("-DACCURATE_DECOMPRESSION");
//									LoadShader(GET_SHADER_INDEX(1, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, SHADER_ACCURATE), str, cgfProf, vmacros, context);
//									vmacros.pop_back();

									LoadShader(GET_SHADER_INDEX(1, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, 0), str, cgfProf, vmacros, context);

									// clut32
									sprintf(str, "Texture%s%d_clut32PS", fog?"Fog":"", texfilter);

//									vmacros.push_back("-DACCURATE_DECOMPRESSION");
//									LoadShader(GET_SHADER_INDEX(2, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, SHADER_ACCURATE), str, cgfProf, vmacros, context);
//									vmacros.pop_back();

									LoadShader(GET_SHADER_INDEX(2, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, 0), str, cgfProf, vmacros, context);

									// tex32to16
									sprintf(str, "Texture%s%d_tex32to16PS", fog?"Fog":"", texfilter);

//									vmacros.push_back("-DACCURATE_DECOMPRESSION");
//									LoadShader(GET_SHADER_INDEX(3, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, SHADER_ACCURATE), str, cgfProf, vmacros, context);
//									vmacros.pop_back();

									LoadShader(GET_SHADER_INDEX(3, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, 0), str, cgfProf, vmacros, context);

									// tex16to8h
									sprintf(str, "Texture%s%d_tex16to8hPS", fog?"Fog":"", texfilter);

//									vmacros.push_back("-DACCURATE_DECOMPRESSION");
//									LoadShader(GET_SHADER_INDEX(4, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, SHADER_ACCURATE), str, cgfProf, vmacros, context);
//									vmacros.pop_back();

									LoadShader(GET_SHADER_INDEX(4, texfilter, texwrap, fog, writedepth, testaem, exactcolor, context, 0), str, cgfProf, vmacros, context);
								}

								if( exactcolor )
									vmacros.pop_back();
							}

							if( testaem )
								vmacros.pop_back();
						}

						if( writedepth )
							vmacros.pop_back();
					}
				}
			}
		}

		if( g_pPsTexWrap[texwrap] != NULL )
			vmacros.pop_back();
	}

	if( vmacros.size() != 0 )
		printf("error with macros!\n");

	// create the database

	int num = (int)mapShaders.size();
	
	// first compress
	vector<char> buffer;
	buffer.reserve(10000000); // 10mb
	buffer.resize(sizeof(SHADERHEADER)*num);

	i = 0;	
	for(map<int, SHADERINFO>::iterator it = mapShaders.begin(); it != mapShaders.end(); ++it, ++i) {
		SHADERHEADER h;
		h.index = it->first | it->second.type;
		h.offset = (int)buffer.size();
		h.size = (int)it->second.buf.size();

		memcpy(&buffer[0] + i*sizeof(SHADERHEADER), &h, sizeof(SHADERHEADER));

		size_t cur = buffer.size();
		buffer.resize(cur + it->second.buf.size());
		memcpy(&buffer[cur], &it->second.buf[0], it->second.buf.size());
	}

	int compressed_size;
	int real_size = (int)buffer.size();
	vector<char> dst;
	dst.resize(buffer.size());
	def(&buffer[0], &dst[0], (int)buffer.size(), &compressed_size);

	// write to file
	// fmt: num shaders, size of compressed, compressed data
	FILE* fdst = fopen(dstfilename, "wb");
	if( fdst == NULL ) {
		printf("failed to open %s\n", dstfilename);
		return 0;
	}

	fwrite(&num, 4, 1, fdst);
	fwrite(&compressed_size, 4, 1, fdst);
	fwrite(&real_size, 4, 1, fdst);
	fwrite(&dst[0], compressed_size, 1, fdst);

	fclose(fdst);

	printf("wrote %s\n", dstfilename);

	cgDestroyContext(g_cgcontext);

	return 0;
}