#include "sgu_stdafx.h"
#include "sgu_arg_opt.h"
#include "sgu_utils.h"
#include "sgu_mode_gen_rand_file.h"
#include "sgu_mode_rand_modify.h"
#include "sgu_mode_rand_remove.h"
#include "sgu_mode_rand_insert.h"
#include "sgu_mode_make_diff.h"
#include "sgu_mode_apply_patch.h"
#include "sgu_mode_hash_file.h"
#include "sgu_mode_npk_dump_list.h"
#include "sgu_mode_npk_extract_files.h"
#include "sgu_mode_download.h"
#include "sgu_mode_run_lua_file.h"

#include "libEncrypt_AES.h"

//int TestFunc(lua_State* L)
//{
//	printf("# TestFunc\n");
//	return lua_yield(L, 0);
//}
//
//int TestFunc2(lua_State* L, float a)
//{
//	printf("# TestFunc2(L,%f)\n", a);
//	return lua_yield(L, 0);
//}
//
//class TestClass
//{
//public:
//
//	int TestFunc(lua_State* L, int a)
//	{
//		m_data = a;
//		printf("# TestClass::TestFunc\n");
//		return lua_yield(L, 0);
//	}
//
//	int TestFunc2(lua_State* L, float a)
//	{
//		printf("# TestClass::TestFunc2(L,%d, %f)\n", m_data, a);
//		return lua_yield(L, 0);
//	}
//
//private:
//	int m_data;
//};
//
//int foo()
//{
//	lua_State* L = lua_open();
//	luaopen_base(L);
//	luaopen_string(L);
//
//	lua_tinker::def(L, "TestFunc", &TestFunc);
//	lua_tinker::def(L, "TestFunc2", &TestFunc2);
//
//	lua_tinker::class_add<TestClass>(L, "TestClass");
//	lua_tinker::class_def<TestClass>(L, "TestFunc", &TestClass::TestFunc);
//	lua_tinker::class_def<TestClass>(L, "TestFunc2", &TestClass::TestFunc2);
//
//	TestClass g_test;
//	lua_tinker::set(L, "g_test", &g_test);
//
//
//	
//	lua_tinker::dostring(L, 
//		"function ThreadTest()\n"
//		"	print(\"ThreadTest\")\n"
//
//		"	print(\"TestFunc Begin\")"
//		"	TestFunc()\n"
//		"	TestFunc2(1.2)\n"
//		"	print(\"TestFunc End\")\n"
//
//		"	print(\"g_test::TestFunc()\")\n"
//		"	g_test:TestFunc(123)\n"
//		"	g_test:TestFunc2(2.3)\n"
//		"	print(\"g_test::TestFunc()\")\n"
//		"end\n"
//	);
//
//	lua_newthread(L);
//	lua_pushstring(L, "ThreadTest");
//	lua_gettable(L, LUA_GLOBALSINDEX);
//
//	printf("* lua_resume() 1\n");
//	lua_resume(L, 0);
//
//	printf("* lua_resume() 2\n");
//	lua_resume(L, 0);
//
//	printf("* lua_resume() 3\n");
//	lua_resume(L, 0);
//
//	printf("* lua_resume() 4\n");
//	lua_resume(L, 0);
//
//	printf("* lua_resume() 5\n");
//	lua_resume(L, 0);
//
//	lua_close(L);
//
//	return 0;
//}
////////////////////////////////////////////////////////////////////////////////////////////
enum { OPT_MODE, 
	OPT_INPUT, OPT_INPUT2, OPT_OUTPUT, 
	OPT_MIN_SIZE, OPT_MAX_SIZE, 
	OPT_COUNTS, 
	OPT_TYPE,
	OPT_VERBOSE_MODE, 
	OPT_HELP 
};

CSimpleOptA::SOption g_rgOptions[] = {
	{ OPT_MODE, "-m",     SO_REQ_SEP },  // "-m mode"
	{ OPT_INPUT, "-i",  SO_REQ_SEP }, // "-i input.bin"
	{ OPT_INPUT2, "-i2",  SO_REQ_SEP }, // "-i2 input2.bin"
	{ OPT_OUTPUT, "-o",  SO_REQ_SEP }, // "-o output.bin"
	{ OPT_MIN_SIZE, "-min",  SO_REQ_SEP }, // "-min 100|201MB|2.5GB"
	{ OPT_MAX_SIZE, "-max",  SO_REQ_SEP }, // "-max 100|201MB|2.5GB"
	{ OPT_COUNTS, "-c",  SO_REQ_SEP }, // "-c 5"
	{ OPT_TYPE, "-t",  SO_REQ_SEP }, // "-t md5|sha256"
	{ OPT_VERBOSE_MODE, "-v",  SO_NONE },	// "-v"

	{ OPT_HELP, "-?",     SO_NONE },	// "-?"
	{ OPT_HELP, "--help", SO_NONE },	 // "--help"
	SO_END_OF_OPTIONS                   // END
};

static void printUsage(const char* moduleName)
{
	printf("===== DNF SideGameplay Utils  =====\n");
	printf("Usage: %s -m [MODE] [OPTIONS]\n", moduleName);
	printf("\t -m Mode\tfunctioning mode\n");
	printf("\t -i FileName\tInput file name\n");
	printf("\t -i2 FileName\tSecond Input file name\n");
	printf("\t -o FileName\tOutput file name\n");
	printf("\t -min Size\tMini size[kb|mb|gb]\n");
	printf("\t -max Size\tMax size[kb|mb|gb]\n");
	printf("\t -c Counts\n");
	printf("\t -t Type\n");
	printf("\t -v\t\tVerbose Mode\n");
	printf("\t --help -?\tShow help info\n");
	printf("\nMode List:\n");
	printf(" GenRandFile\tGenerate random file with specified size range\n");
	printf(" RandModifyFile\tModifies an input file by replacing multiple randomly selected segments with random data.\n");
	printf(" RandInsertFile\tInserts random data segments at random positions within an input file.\n");
	printf(" RandRemoveFile\tRemoves multiple randomly selected segments from an input file.\n");
	printf(" MakeDiff\tCreates a binary diff patch file between two input files using BSDIFF43.\n");
	printf(" ApplyPatch\tApplies a BSDIFF43 binary patch to an input file and writes the patched result to an output file.\n");
	printf(" HashFile\tCalculates the hash(MD5|SHA256) of the contents of a specified file.\n");
	printf(" NPKDumpList\tDumps the file list from an NPK archive to a text file.\n");
	printf(" NPKExtractFiles\tExtracts all files from an NPK archive to a specified output directory.\n");
	printf(" HttpDownload\tDownloads a file from a specified URL and saves it to a local file.\n");
	printf(" RunLuaFile\t xxxx \n");
	printf("\nSample:\n");
	printf(" -m GenRandFile -min 500 -max 1k -o rand.bin\n");
	printf(" -m RandModifyFile -i ver01.bin -min 100 -max 1k -c 3 -o ver02.bin\n");
	printf(" -m RandInsertFile -i ver01.bin -min 1g -max 2g -c 5 -o ver02.bin\n");
	printf(" -m RandRemoveFile -i ver01.bin -min 10m -max 20m -c 2 -o ver02.bin\n");
	printf(" -m MakeDiff -i ver01.bin -i2 ver02.bin -o patch_ver01_ver02.diff\n");
	printf(" -m ApplyPatch -i ver01.bin -i2 patch_ver01_ver02.diff -o ver02.bin\n");
	printf(" -m HashFile -t md5 -i ver01.bin\n");
	printf(" -m HashFile -t sha256 -i ver01.bin\n");
	printf(" -m NPKDumpList -i abc.npk -o abc_file_list.txt\n");
	printf(" -m NPKExtractFiles -i abc.npk -o output\n");
	printf(" -m HttpDownload -i http://download.com/file.txt -o file.txt\n");
	printf(" -m RunLuaFile -i script.lua\n");
}

enum WorkMode {
	WM_NONE=0,

	WM_GEN_RANDOM_FILE,
	WM_RANDOM_MODIFY_FILE,
	WM_RANDOM_INSERT_FILE,
	WM_RANDOM_REMOVE_FILE,
	WM_MAKE_DIFF,
	WM_APPLY_PATCH,
	WM_HASH_FILE,
	WM_HTTP_DOWNLOAD,
	WM_NPK_DUMPLIST,
	WM_NPK_EXTRACT_FILES,
	WM_RUN_LUA_FILE,
};

enum HashType {
	HT_NONE = 0,

	HT_MD5,
	HT_SHA256,
	HT_CRC32
};

WorkMode parserWorkMode(const char* modeStr)
{
	if(modeStr== nullptr || *modeStr == '\0') {
		return WM_NONE;
	}

	if (_stricmp(modeStr, "GenRandFile") == 0) {
		return WM_GEN_RANDOM_FILE;
	}
	else if (_stricmp(modeStr, "RandModifyFile") == 0) {
		return WM_RANDOM_MODIFY_FILE;
	}
	else if (_stricmp(modeStr, "RandInsertFile") == 0) {
		return WM_RANDOM_INSERT_FILE;
	}
	else if (_stricmp(modeStr, "RandRemoveFile") == 0) {
		return WM_RANDOM_REMOVE_FILE;
	}
	else if (_stricmp(modeStr, "MakeDiff") == 0) {
		return WM_MAKE_DIFF;
	}
	else if (_stricmp(modeStr, "ApplyPatch") == 0) {
		return WM_APPLY_PATCH;
	}
	else if (_stricmp(modeStr, "HashFile") == 0) {
		return WM_HASH_FILE;
	}
	else if (_stricmp(modeStr, "HttpDownload") == 0) {
		return WM_HTTP_DOWNLOAD;
	}
	else if (_stricmp(modeStr, "NPKDumpList") == 0) {
		return WM_NPK_DUMPLIST;
	}
	else if (_stricmp(modeStr, "NPKExtractFiles") == 0) {
		return WM_NPK_EXTRACT_FILES;
	}
	else if (_stricmp(modeStr, "RunLuaFile") == 0) {
		return WM_RUN_LUA_FILE;
	}
	return WM_NONE;
}

HashType parserHashType(const char* typeStr)
{
	if(typeStr == nullptr || *typeStr == '\0') {
		return HT_NONE;
	}

	if (_stricmp(typeStr, "md5") == 0) {
		return HT_MD5;
	}
	else if (_stricmp(typeStr, "sha256") == 0) {
		return HT_SHA256;
	}
	else if (_stricmp(typeStr, "crc32") == 0) {
		return HT_CRC32;
	}
	return HT_NONE; // default to HT_NONE
}

int main(int argc, char* argv[]) 
{
	CSimpleOptA args(argc, argv, g_rgOptions);
	
	WorkMode workMode = WM_NONE;
	const char* inputFilename = nullptr;
	const char* inputFilename2 = nullptr;
	const char* outputFilename = nullptr;
	int64_t minSize = 0;
	int64_t maxSize = 0;
	int32_t counts = 0;
	const char* stringType = nullptr;
	bool verboseMode = false;

	while (args.Next()) {
		if (args.LastError() == SO_SUCCESS) {
			if (args.OptionId() == OPT_HELP) {
				printUsage(argv[0]);
				return 0;
			}
			else if (args.OptionId() == OPT_MODE) {
				workMode = parserWorkMode(args.OptionArg());
			}
			else if (args.OptionId() == OPT_INPUT) {
				inputFilename = args.OptionArg();
			}
			else if (args.OptionId() == OPT_INPUT2) {
				inputFilename2 = args.OptionArg();
			}
			else if (args.OptionId() == OPT_OUTPUT) {
				outputFilename = args.OptionArg();
			}
			else if (args.OptionId() == OPT_MIN_SIZE) {
				minSize = sizeFromString(args.OptionArg());
			}
			else if (args.OptionId() == OPT_MAX_SIZE) {
				maxSize = sizeFromString(args.OptionArg());
			}
			else if (args.OptionId() == OPT_COUNTS) {
				counts = atoi(args.OptionArg());
			}
			else if (args.OptionId() == OPT_TYPE) {
				stringType = args.OptionArg();
			}
			else if (args.OptionId() == OPT_VERBOSE_MODE) {
				verboseMode = true;
			}
		}
		else {
			printf("Invalid argument: %s\n", args.OptionText());
			printUsage(::PathFindFileNameA(argv[0]));
			return 1;
		}
	}

	//init random seed
	std::srand((unsigned int)time(NULL));

	switch (workMode)
	{
	case WM_GEN_RANDOM_FILE:
		{
			if(outputFilename==nullptr || minSize==0 || maxSize==0) {
				printf("Error: For 'gen' mode, you must specify output file and size range.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			printf("Generating random file with min size: %s, max size: %s, output: %s\n", 
				sizeToString(minSize).c_str(), sizeToString(maxSize).c_str(), outputFilename);
			if (!generateRandomFile(minSize, maxSize, outputFilename))
			{
				return 1;
			}
			break;
		}
		break;

	case WM_RANDOM_MODIFY_FILE:
		{
			if(inputFilename == nullptr || outputFilename == nullptr || minSize == 0 || maxSize == 0 || counts==0) 
			{
				printf("Error: For 'modify' mode, you must specify input file, output file, size range and counts.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			if(!randomModifyFile(inputFilename, outputFilename, minSize, maxSize, counts))
			{
				printf("Error: Failed to modify file.\n");
				return 1;
			}
		}
		break;

	case WM_RANDOM_REMOVE_FILE:
		{
			if (inputFilename == nullptr || outputFilename == nullptr || minSize == 0 || maxSize == 0 || counts == 0)
			{
				printf("Error: For 'remove' mode, you must specify input file, output file, size range and counts.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			if (!randRemovePartOfFile(inputFilename, outputFilename, minSize, maxSize, counts))
			{
				printf("Error: Failed to remove part of file.\n");
				return 1;
			}
		}
		break;

	case WM_RANDOM_INSERT_FILE:
		{
			if (inputFilename == nullptr || outputFilename == nullptr || minSize == 0 || maxSize == 0 || counts == 0)
			{
				printf("Error: For 'insert' mode, you must specify input file, output file, size range and counts.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			if(!insertRandomDataInFile(inputFilename, outputFilename, minSize, maxSize, counts))
			{
				printf("Error: Failed to insert data in file.\n");
				return 1;
			}
		}
		break;
	case WM_HASH_FILE:
		{
			if(inputFilename==nullptr ) 
			{
				printf("Error: For 'hash' mode, you must specify input file.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			HashType hashType = parserHashType(stringType);
			if(hashType==HT_NONE) 
			{
				printf("Error: Invalid hash type '%s'. Supported types are 'md5' and 'sha256'.\n", stringType);
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			printf("Calculating hash for file[%s]: %s\n", stringType, inputFilename);

			if(hashType == HT_SHA256) 
			{
				Sha256Digest result;
				if (!sha256File(inputFilename, result)) 
				{
					return 1;
				}
				printf("SHA256(%s) = %s\n", inputFilename, sha256ToString(result).c_str());
			}
			else if (hashType == HT_MD5) 
			{
				MD5Digest result;
				if (!md5File(inputFilename, result))
				{
					return 1;
				}
				printf("MD5(%s) = %s\n", inputFilename, md5ToString(result).c_str());
			}
			else if(hashType == HT_CRC32) 
			{
				uint32_t crc32;
				if (!crc32File(inputFilename, crc32))
				{
					return 1;
				}
				printf("CRC32(%s) = %08X\n", inputFilename, crc32);
			}
			else 
			{
				printf("Error: Unsupported hash type.\n");
				return 1;
			}
		}
		break;

	case WM_MAKE_DIFF:
		{
			if (inputFilename == nullptr || inputFilename2 == nullptr || outputFilename == nullptr) 
			{
				printf("Error: For 'diff' mode, you must specify two input files and an output file.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			printf("Creating diff between %s and %s, output: %s\n", inputFilename, inputFilename2, outputFilename);
			if (!createDiff(inputFilename, inputFilename2, outputFilename)) 
			{
				printf("Error: Failed to create diff file.\n");
				return 1;
			}
		}
		break;

	case WM_APPLY_PATCH:
		{
			if (inputFilename == nullptr || inputFilename2 == nullptr || outputFilename == nullptr)
			{
				printf("Error: For 'patch' mode, you must specify two input files and an output file.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			printf("Applying patch from %s with patch %s, output: %s\n", inputFilename, inputFilename2, outputFilename);
			if (!applyPatch(inputFilename, inputFilename2, outputFilename))
			{
				printf("Error: Failed to apply patch.\n");
				return 1;
			}
		}
		break;
	case WM_HTTP_DOWNLOAD:
		{
			if (inputFilename == nullptr || outputFilename == nullptr)
			{
				printf("Error: For 'download' mode, you must specify input URL and output file.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			printf("Downloading from %s to %s\n", inputFilename, outputFilename);
			if (!downloadFile(inputFilename, outputFilename))
			{
				printf("Error: Failed to download file.\n");
				return 1;
			}
		}
		break;

	case WM_NPK_DUMPLIST:
		{
			if (inputFilename == nullptr || outputFilename == nullptr) 
			{
				printf("Error: For 'NPKDumpList' mode, you must specify input NPK file and output file.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			printf("Dumping NPK list from %s to %s\n", inputFilename, outputFilename);
			if (!dumpNPKFileList(inputFilename, outputFilename))
			{
				printf("Error: Failed to dump NPK list.\n");
				return 1;
			}
		}
		break;

	case WM_NPK_EXTRACT_FILES:
		{
			if (inputFilename == nullptr || outputFilename == nullptr)
			{
				printf("Error: For 'NPKExtractFiles' mode, you must specify input NPK file and output directory.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			printf("Extracting files from NPK '%s' to directory '%s'\n", inputFilename, outputFilename);
			if(!extractNPKFiles(inputFilename, outputFilename))
			{
				printf("Error: Failed to extract files from NPK.\n");
				return 1;
			}
		}
		break;
	case WM_RUN_LUA_FILE:
		{
			if (inputFilename == nullptr)
			{
				printf("Error: For 'RunLuaFile' mode, you must specify input Lua script file.\n");
				printUsage(::PathFindFileNameA(argv[0]));
				return 1;
			}
			printf("Running Lua script file: %s\n", inputFilename);
			if (!run_lua_file(inputFilename))
			{
				printf("Error: Failed to run Lua script file.\n");
				return 1;
			}
		}
		break;
	case WM_NONE:
	default:
	{
		printf("Error: Mode is not specified. Use -m option.\n");
		printUsage(::PathFindFileNameA(argv[0]));
		return 1;
	}
	}
	return 0;
}