#include <Windows.h>

#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <regex>

typedef UINT uint;

#define LogMessage(...) do { fprintf(stderr, "[%s:%d]", __FILE__, __LINE__); fprintf(stderr, __VA_ARGS__); } while(0);

#define SOURCE_LINE_FLAG_NONE 0
#define SOURCE_LINE_FLAG_DISCARD 1
#define SOURCE_LINE_FLAG_FORCE_NEWLINE 2

struct SourceLine {
	SourceLine(uint l, const std::string& s) : lineNumber(l), string(s), flag(SOURCE_LINE_FLAG_NONE) {}
	uint lineNumber;
	std::string string;
	uint flag;
};

struct EnumBlock {
	std::string enumName;
	std::vector<std::pair<std::string,std::string>> enumEntries;
};

struct InterfaceBlock {
	std::string interfaceName;
	std::string interfaceGUID;
	std::vector<std::string> interfaceMethods;

};

struct StructBlock {
	std::string structName;
	std::vector<std::pair<std::string, std::string>> structEntries;
};

struct DefineBlock {
	std::vector<std::pair<std::string,std::string>> defines;
};

static size_t getFileSize(const char* filename) {
	FILE* f = fopen(filename, "rb");
	if (f == NULL) return 0;
	fseek(f, 0, SEEK_END);
	size_t size = ftell(f);
	fclose(f);
	return size;
}

static bool readLinesFromFile(const char* filename, std::vector<SourceLine>& lines) {
	std::ifstream ifs;
	ifs.open(filename, std::ios::in);
	if (!ifs.good()) {
		LogMessage("readLinesFromFile() - could not open file %s", filename);
		return false;
	}

	std::string str;
	uint lineNumber = 0;
	while (ifs.good() && getline(ifs, str))
		lines.push_back(SourceLine(++lineNumber, str));
	ifs.close();
	return true;
}

std::string trim(std::string& s) {
	s.erase(0, s.find_first_not_of(" \t\n\r\v\f"));
	s.erase(s.find_last_not_of(" \t\n\r\v\f") + 1);
	return s;
}

void inplaceFindAndReplace(std::string& src, const std::string& find, const std::string& replace) {
	std::string::size_type j, s = 0;
	while ((j = src.find(find, s)) != std::string::npos) {
		src.replace(j, find.length(), replace);
		s = j + replace.length();
	}
}

std::string regexReplace(const std::string& src, const std::string& rgx, const std::string& fmt) {
	return std::tr1::regex_replace(src, std::tr1::regex(rgx), fmt);
}

static void stripLines(std::vector<SourceLine>& lines) {
	bool removeComments = true;
	bool removeFiles = true;
	bool removeStrings = true;
	bool removeSpaces = true;
	bool removeLocations = true;
	bool removeTabulators = true;
	bool replaceTemp = true;
	const char* tempString = "%temp";
	const char* tempReplacementString = "%_";

	uint N = lines.size();
	for (uint i = 0; i < N; ++i) {
		SourceLine& sl = lines[i];

		// Trim left and right whitespaces
		trim(sl.string);

		if (sl.string.empty()) {
			// Remove empty lines
			sl.flag |= SOURCE_LINE_FLAG_DISCARD;
		} else if (sl.string.find("//") == 0 && removeComments) {
			// Remove comments
			sl.flag |= SOURCE_LINE_FLAG_DISCARD;
		} else if (sl.string.find(".file") == 0 && removeFiles) {
			// Remove .file
			sl.flag |= SOURCE_LINE_FLAG_DISCARD;
		} else if (sl.string.find(".b8 $str") != std::string::npos && removeStrings) {
			// Remove strings
			sl.flag |= SOURCE_LINE_FLAG_DISCARD;
		} else if (sl.string.find(".loc") == 0 && removeLocations) {
			// Remove locations
			sl.flag |= SOURCE_LINE_FLAG_DISCARD;
		} else {
			// Remove all tabs
			if (removeTabulators)
				inplaceFindAndReplace(sl.string, "\t", "");

			if (removeSpaces) {
				// Remove possible duplicate spaces
				inplaceFindAndReplace(sl.string, "  ", " ");
				inplaceFindAndReplace(sl.string, "  ", " ");

				// Remove space after ','
				inplaceFindAndReplace(sl.string, ", ", ",");

				// Remove space before '%'
				inplaceFindAndReplace(sl.string, " %", "%");

				// Remove space before '['
				inplaceFindAndReplace(sl.string, " [", "[");

				// Remove space before '.'
				//if (sl.string.find(".align") == std::string::npos)
				//	inplaceFindAndReplace(sl.string, " .", ".");
				sl.string = regexReplace(sl.string, "([A-Za-z_]+)(\\s+)(\\.)", "$1$3");

			}

			if (replaceTemp) {
				inplaceFindAndReplace(sl.string, tempString, tempReplacementString);
			}

			// If the string is empty, discard the line
			if (sl.string.empty()) {
				sl.flag |= SOURCE_LINE_FLAG_DISCARD;
			}
		}

		if ((sl.flag & SOURCE_LINE_FLAG_DISCARD) == 0) {
				if (sl.string.find(".version") == 0 ||
				sl.string.find(".target") == 0 ||
				sl.string.find(".address_size") == 0) {
					sl.flag |= SOURCE_LINE_FLAG_FORCE_NEWLINE;
			}
		}
	}
}

static bool isValueDefine(const std::string& s, std::string& name, std::string& value) {
	size_t definePos = s.find("#define");

	if (definePos != std::string::npos &&
		s.find("__") == std::string::npos && s.find("\\") == std::string::npos) {

		std::string tmp = s.substr(definePos+std::string("#define").size());
		trim(tmp);
		size_t whiteSpacePos = tmp.find_first_of(" \t\v");
		if (whiteSpacePos == std::string::npos) { // || tmp.find("MAKE_D3D11") != std::string::npos || tmp.find("MAKE_DXGI") != std::string::npos) {
			return false;
		}

		name = tmp.substr(0, whiteSpacePos);
		tmp = tmp.substr(whiteSpacePos);
		trim(tmp);

		inplaceFindAndReplace(tmp, "  ", " ");
		inplaceFindAndReplace(tmp, "  ", " ");
		inplaceFindAndReplace(tmp, " (", "(");
		inplaceFindAndReplace(tmp, "( ", "(");
		inplaceFindAndReplace(tmp, " )", ")");
		inplaceFindAndReplace(tmp, ") ", ")");
		inplaceFindAndReplace(tmp, "| ", "|");
		inplaceFindAndReplace(tmp, " |", "|");
		inplaceFindAndReplace(tmp, " <<", "<<");
		inplaceFindAndReplace(tmp, "<< ", "<<");
		inplaceFindAndReplace(tmp, " >>", ">>");
		inplaceFindAndReplace(tmp, ">> ", ">>");
		tmp = regexReplace(tmp, "(\\d+)(UL)", "$1");
		tmp = regexReplace(tmp, "(\\d+)(L)", "$1");
		tmp = regexReplace(tmp, "(\\d+)(\\.*f)", "$1");
		value = tmp;

		return true;
	}
	return false;
}

static bool isInterfaceDeclaration(const std::string& s, std::string& interfaceName) {
	size_t typedefPos = s.find("typedef");
	size_t structPos = s.find("struct");
	size_t vtblPos = s.find("Vtbl");

	if (typedefPos != std::string::npos && structPos != std::string::npos && vtblPos != std::string::npos) {
		structPos += std::string("struct").size();
		std::string tmp = s.substr(structPos, vtblPos - structPos);
		trim(tmp);
		//std::cerr << "#" << tmp << "#" << std::endl;
		interfaceName = tmp;
		return true;
	}

	return false;
}

static bool isBeginInterface(const std::string& s) {
	size_t beginInterfacePos = s.find("BEGIN_INTERFACE");
	if (beginInterfacePos != std::string::npos) {
		return true;
	}
	return false;
}

static bool isFunctionDeclaration(const std::string& s, std::string& functionName) {
	std::string stringSTDMETHODCALLTYPE("STDMETHODCALLTYPE");
	size_t stdcallPos = s.find(stringSTDMETHODCALLTYPE);
	if (stdcallPos != std::string::npos) {
		std::string tmp = s.substr(stdcallPos + stringSTDMETHODCALLTYPE.size());
		size_t ptrPos = tmp.find("*")+1;
		size_t bracketPos = tmp.find(")");
		tmp = tmp.substr(ptrPos, bracketPos - ptrPos);
		functionName = trim(tmp);
		//std::cerr << "#" << tmp << "#" << std::endl;

		return true;
	}
	return false;

}

static bool isEndInterface(const std::string& s) {
	size_t beginInterfacePos = s.find("END_INTERFACE");
	if (beginInterfacePos != std::string::npos) {
		return true;
	}
	return false;
}

static bool isStructDeclaration(const std::string& s, std::string& structName) {
	size_t typedefPos = s.find("typedef");
	size_t structPos = s.find("struct");
	size_t vtblPos = s.find("Vtbl");

	if (typedefPos != std::string::npos && structPos != std::string::npos && vtblPos == std::string::npos) {
		structPos += std::string("struct").size();
		std::string tmp = s.substr(structPos);
		trim(tmp);
		//std::cerr << "#" << tmp << "#" << std::endl;
		structName = tmp;
		return true;
	}

	return false;
}


static bool isEnum(const std::string& s, std::string& enumName) {
	size_t enumPos = s.find("enum");
	if (enumPos != std::string::npos) {
		std::string tmp = s.substr(enumPos + std::string("enum").size());
		trim(tmp);
		enumName = tmp;
		return true;
	}
	return false;
}

static bool isValidEnumEntry(const std::string& s, std::string& name, std::string& value) {
	size_t equalSignPos = s.find("=");
	if (equalSignPos == std::string::npos)
		return false;
	
	std::string tmp = s.substr(0, equalSignPos);
	trim(tmp);
	name = tmp;

	tmp = s.substr(equalSignPos+1);

	inplaceFindAndReplace(tmp, ",", " ");
	inplaceFindAndReplace(tmp, "  ", " ");
	inplaceFindAndReplace(tmp, "  ", " ");
	inplaceFindAndReplace(tmp, " (", "(");
	inplaceFindAndReplace(tmp, "( ", "(");
	inplaceFindAndReplace(tmp, " )", ")");
	inplaceFindAndReplace(tmp, ") ", ")");
	inplaceFindAndReplace(tmp, "| ", "|");
	inplaceFindAndReplace(tmp, " |", "|");
	inplaceFindAndReplace(tmp, " <<", "<<");
	inplaceFindAndReplace(tmp, "<< ", "<<");
	inplaceFindAndReplace(tmp, " >>", ">>");
	inplaceFindAndReplace(tmp, ">> ", ">>");
	tmp = regexReplace(tmp, "(\\d+)(UL)", "$1");
	tmp = regexReplace(tmp, "(\\d+)(L)", "$1");
	tmp = regexReplace(tmp, "(\\d+)(\\.*f)", "$1");
	
	trim(tmp);

	value = tmp;

	//size_t commaPos = s.find(",");
	

	return true;
}

static bool isEnumBegin(const std::string& s) {
	if (s.find("{") != std::string::npos)
		return true;
	return false;
}

static bool isEnumEnd(const std::string& s) {
	if (s.find("}") != std::string::npos)
		return true;
	return false;
}




//static 

static void parseInterfaces(std::vector<SourceLine>& lines, std::vector<InterfaceBlock>& interfaces, std::vector<EnumBlock>& enums, DefineBlock& defineBlock) {
	bool foundInterfaceBlock = false;
	bool foundBlockBegin = false;
	bool foundEnum = false;
	bool foundEnumBlockBegin = false;


	bool end = false;
	int currentLineIndex = 0;

	std::string interfaceName;
	std::string methodName;
	std::string enumName;
	std::string structName;

	size_t N = lines.size();
	int currentInterfaceIndex = -1;
	int currentEnumIndex = -1;

	for (size_t i = 0; i < N; ++i) {
		SourceLine& sl = lines[i];
		// Trim left and right whitespaces
		trim(sl.string);
		{ // Interfaces
			if (isInterfaceDeclaration(sl.string, interfaceName)) {
				//std::cout << "interface " << interfaceName << std::endl;
				foundInterfaceBlock = true;
			}

			if (isBeginInterface(sl.string) && foundInterfaceBlock) {
				foundBlockBegin = true;

				InterfaceBlock ib;
				ib.interfaceMethods.clear();
				ib.interfaceName = interfaceName;
				interfaces.push_back(ib);

				currentInterfaceIndex++;
			}

			if (foundBlockBegin) {
				if (isFunctionDeclaration(sl.string, methodName)) {
					//std::cout << "\t" << methodName << std::endl;
					if (foundBlockBegin) {
						InterfaceBlock& ib = interfaces.at(currentInterfaceIndex);
						ib.interfaceMethods.push_back(methodName);
					}
				}
			}

			if (isEndInterface(sl.string)) {
				foundInterfaceBlock = false;
				foundBlockBegin = false;
				//std::cout << std::endl;
			}
		}

		{ // Structs
			if (isStructDeclaration(sl.string, structName)) {
				std::cout << structName << std::endl;
			}
		}

		{ // Defines
			std::string defineName, defineValue;
			if (isValueDefine(sl.string, defineName, defineValue)) {
				//std::cout << defineName<< " " << defineValue << std::endl;
				defineBlock.defines.push_back(std::make_pair(defineName, defineValue));
			}
		}

		{
			
			if (isEnum(sl.string, enumName)) {
				foundEnum = true;
				//std::cout << enumName << std::endl;
			}

			if (foundEnum) {
				if (isEnumBegin(sl.string)) {
					foundEnumBlockBegin = true;
					EnumBlock eb;
					eb.enumEntries.clear();
					eb.enumName = enumName;
					enums.push_back(eb);
					currentEnumIndex++;
				}
			}
			
			if (foundEnumBlockBegin) {
				std::string enumElementName, enumElementValue;
				if (isValidEnumEntry(sl.string, enumElementName, enumElementValue)) {
					//std::cout << enumElementName << " " << enumElementValue << std::endl;
					EnumBlock& eb = enums.at(currentEnumIndex);
					eb.enumEntries.push_back(make_pair(enumElementName, enumElementValue));
				}
			}

			if (isEnumEnd(sl.string)) {
				foundEnum = false;
				foundEnumBlockBegin = false;
			}

		}

	}

	size_t numInterfaces = interfaces.size();

}


static bool writeLinesToFile(const char* filename, const std::vector<SourceLine>& lines) {
	std::ofstream ofs;
	ofs.open(filename, std::ios::out);
	if (!ofs.good()) {
		LogMessage("writeLinesToFile() - could not open file %s\n", filename);
		return false;
	}

	uint N = lines.size();
	for (uint i = 0; i < N; ++i) {
		const SourceLine& sl = lines[i];
		if (sl.flag != SOURCE_LINE_FLAG_DISCARD) {
			ofs << sl.string;
			if (sl.flag & SOURCE_LINE_FLAG_FORCE_NEWLINE)
				ofs << std::endl;
		}
	}
	ofs.close();
	return true;
}

static bool writeInterfacesToFileASM(const char* filename, const std::vector<InterfaceBlock>& interfaces, const std::vector<EnumBlock>& enums, const DefineBlock& defineBlock) {
	std::ofstream ofs;
	ofs.open(filename, std::ios::out);
	if (!ofs.good()) {
		LogMessage("writeLinesToFile() - could not open file %s\n", filename);
		return false;
	}
	ofs << "; " << filename << " by las/mercury" << std::endl << std::endl;
	if (std::string(filename).find("d3d11") != std::string::npos) {
		//ofs << "; " << filename << " by las/mercury" << std::endl;
		ofs << "extern _D3D11CreateDeviceAndSwapChain@48" << std::endl;
		ofs << "extern _D3DCompile@44" << std::endl;
		ofs << std::endl;
		/*
		typedef struct D3D11_TEXTURE2D_DESC
    {
    UINT Width;
    UINT Height;
    UINT MipLevels;
    UINT ArraySize;
    DXGI_FORMAT Format;
    DXGI_SAMPLE_DESC SampleDesc;
    D3D11_USAGE Usage;
    UINT BindFlags;
    UINT CPUAccessFlags;
    UINT MiscFlags;
    } 	D3D11_TEXTURE2D_DESC;
		*/
		std::string struc = std::string("");

		struc += "struc D3D11_BUFFER_DESC\n"
			"\t.ByteWidth: resd 1\n"
			"\t.Usage: resd 1\n"
			"\t.BindFlags: resd 1\n"
			"\t.CPUAccessFlags: resd 1\n"
			"\t.MiscFlags: resd 1\n"
			"\t.StructureByteStride: resd 1\n"
			"endstruc\n";
		struc += "\n";

		struc += "struc D3D11_SAMPLER_DESC\n"
			"\t.Filter: resd 1\n"
			"\t.AddressU: resd 1\n"
			"\t.AddressV: resd 1\n"
			"\t.AddressW: resd 1\n"
			"\t.MipLODBias: resd 1\n"
			"\t.MaxAnisotropy: resd 1\n"
			"\t.ComparisonFunc: resd 1\n"
			"\t.BorderColor: resd 4\n"
			"\t.MinLOD: resd 1\n"
			"\t.MaxLOD: resd 1\n"
			"endstruc\n";
		struc += "\n";

		struc += "struc D3D11_TEXTURE2D_DESC\n"
			"\t.Width: resd 1\n"
			"\t.Height: resd 1\n"
			"\t.MipLevels: resd 1\n"
			"\t.ArraySize: resd 1\n"
			"\t.Format: resd 1\n"
			"\n"
			"\t; DXGI_SAMPLE_DESC SampleDesc\n"
			"\t.Count: resd 1\n"
			"\t.Quality: resd 1\n"
			"\n"
			"\t.Usage: resd 1\n"
			"\t.BindFlags: resd 1\n"
			"\t.CPUAccessFlags: resd 1\n"
			"\t.MiscFlags: resd 1\n"
			"endstruc\n";
		struc += "\n";

		ofs << struc << std::endl;

	} else if (std::string(filename).find("dxgi") != std::string::npos) {
		std::string struc = std::string(
		"struc DXGI_SWAP_CHAIN_DESC\n"
			"\t; DXGI_MODE_DESC BufferDesc\n"
			"\t.Width: resd 1\n"
			"\t.Height: resd 1\n"
			"\n"
			"\t; DXGI_RATIONAL RefreshRate\n"
			"\t.Numerator: resd 1\n"
			"\t.Denominator: resd 1\n"
			"\n"
			"\t.Format: resd 1\n"
			"\t.ScanlineOrdering: resd 1\n"
			"\t.Scaling: resd 1\n"
			"\n"
			"\t; DXGI_SAMPLE_DESC SampleDesc\n"
			"\t.Count: resd 1\n"
			"\t.Quality: resd 1\n"
			"\n"
			"\t.BufferUsage: resd 1\n"
			"\t.BufferCount: resd 1\n"
			"\t.OutputWindow: resd 1\n"
			"\t.Windowed: resd 1\n"
			"\t.SwapEffect: resd 1\n"
			"\t.Flags: resd 1\n"
		"endstruc\n");
		ofs << struc << std::endl;
	}

	{
		uint N = enums.size();
		for (uint i = 0; i < N; ++i) {
			const EnumBlock& eb = enums[i];
			uint M = eb.enumEntries.size();
			ofs << "; " << eb.enumName << std::endl;
			for (auto &enumElement : eb.enumEntries) {
				ofs << "%define " << enumElement.first << " " << enumElement.second << std::endl;
			}
			ofs << std::endl;
		}
		ofs << std::endl;

	}


	{
		uint N = defineBlock.defines.size();
		for (uint i = 0; i < N; ++i) {
			const auto& def = defineBlock.defines[i];
			ofs << "%define " << def.first << " " << def.second << std::endl;
		}
		ofs << std::endl;

	}

	{
		uint N = interfaces.size();
		for (uint i = 0; i < N; ++i) {
			const InterfaceBlock& ib = interfaces[i];
			ofs << std::endl <<  "struc " << ib.interfaceName << std::endl;
			uint numInterfaces = ib.interfaceMethods.size();
			for (uint j = 0; j < numInterfaces; ++j) {
				ofs << "\t." << ib.interfaceMethods[j] << ": " << "resb 4" << std::endl;
			}
			ofs << "endstruc" << std::endl;
		}
	}
	ofs.close();
	return true;
}

int main(int argc, char** argv) {
	if (argc != 3 && argc != 2) {
		fprintf(stderr, "Usage: ptxStrip in.ptx [out.ptx]\n");
		return EXIT_FAILURE;
	}

	const char* inputFile = argv[1];
	const char* outputFile = argv[argc == 2 ? 1 : 2];

	// Read in the source
	std::vector<SourceLine> source;
	if (!readLinesFromFile(inputFile, source)) {
		return EXIT_FAILURE;
	}
	if (std::string(inputFile).find("dxgi.h") != std::string::npos) {
		if (!readLinesFromFile("dxgiformat.h", source)) {
			return EXIT_FAILURE;
		}
		if (!readLinesFromFile("dxgitype.h", source)) {
			return EXIT_FAILURE;
		}		
	}
	if (std::string(inputFile).find("d3d11.h") != std::string::npos) {
		if (!readLinesFromFile("d3dcommon.h", source)) {
			return EXIT_FAILURE;
		}
		if (!readLinesFromFile("d3dcompiler.h", source)) {
			return EXIT_FAILURE;
		}		
	}


	size_t inputSize = getFileSize(inputFile);
	fprintf(stdout, "Input file size: %d bytes\n", inputSize);

	// Find interfaces in the header file
	std::vector<InterfaceBlock> interfaces;
	std::vector<EnumBlock> enums;
	DefineBlock defineBlock;
	
	parseInterfaces(source, interfaces, enums, defineBlock);

	// Write the new *.inc
	if (!writeInterfacesToFileASM(outputFile, interfaces, enums, defineBlock)) {
		return EXIT_FAILURE;
	}

	size_t outputSize = getFileSize(outputFile);
	fprintf(stdout, "Output file size: %d bytes\n", outputSize);

	return EXIT_SUCCESS;
}