#include "ANSODTraininingEngine.h"
#include "ANSYOLOTrainingEngine.h"
#include "ANSYOLOCLTrainingEngine.h"
#include "ANSYOLOSEGTrainingEngine.h"
#include "ANSTrainingEngine.h"
#include <fstream>
#include <sstream>
#include <string>
#include <iostream>
#include <filesystem>
#include <iostream>
#include <thread>
#include <queue>
#include <mutex>
#include <condition_variable>
#include <string>
#include <cstdio>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/json_parser.hpp>
#include <chrono>
#include <ctime>
#include <json.hpp>
#ifdef _WIN32
#define popen _popen
#define pclose _pclose
#endif

using namespace boost::property_tree;
using json = nlohmann::ordered_json;


// Thread-safe queue
std::queue<std::string> outputQueue;
std::mutex queueMutex;
std::condition_variable cv;
bool done = false;
// Thread 1: Run command and push output lines
void runCommand(const std::string& command) {
	// Redirect both stdout and stderr to ensure nothing leaks directly
	std::string cmdWithRedirect = command + " 2>&1";

	FILE* pipe = popen(cmdWithRedirect.c_str(), "r");
	if (!pipe) {
		std::cerr << "Failed to run command: " << command << "\n";
		{
			std::lock_guard<std::mutex> lock(queueMutex);
			done = true;
		}
		cv.notify_all();
		return;
	}

	char buffer[256];
	while (fgets(buffer, sizeof(buffer), pipe)) {
		{
			std::lock_guard<std::mutex> lock(queueMutex);
			outputQueue.push(std::string(buffer));
		}
		cv.notify_one();
	}

	pclose(pipe);
	{
		std::lock_guard<std::mutex> lock(queueMutex);
		done = true;
	}
	cv.notify_all();
}
void parseOutput(const std::string& message) {
	const std::string requiredSubstring = "Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size";
	const std::vector<std::string> excludedSubstrings = {
		"C:\\ProgramData\\Sh7O7nUe7vJ",
		"train: Scanning",
		"val: Scanning","ultralytics","YOLO","python","traceback","Ultralytics","Python"
	};

	while (true) {
		std::unique_lock<std::mutex> lock(queueMutex);
		cv.wait(lock, [] { return !outputQueue.empty() || done; });

		while (!outputQueue.empty()) {
			std::string line = outputQueue.front();
			outputQueue.pop();
			lock.unlock();

			bool shouldExclude = false;
			for (const auto& exclude : excludedSubstrings) {
				if (line.find(exclude) != std::string::npos) {
					shouldExclude = true;
					break;
				}
			}
			//if (!shouldExclude && line.find(requiredSubstring) != std::string::npos) {
			if (!shouldExclude) {
				if (!line.empty())std::cout << message << line;
			}

			lock.lock();
		}

		if (done && outputQueue.empty()) break;
	}
}
std::string getCurrentUtcIso8601() {
	using namespace std::chrono;

	auto now = system_clock::now();
	auto now_time_t = system_clock::to_time_t(now);
	auto now_ms = duration_cast<milliseconds>(now.time_since_epoch()) % 1000;

	std::tm utc_tm = *std::gmtime(&now_time_t);

	std::ostringstream oss;
	oss << std::put_time(&utc_tm, "%Y-%m-%dT%H:%M:%S");
	oss << '.' << std::setw(3) << std::setfill('0') << now_ms.count() << "Z";

	return oss.str();
}

int CreateANSTSProjectJsonInfo(std::string projectName, int trainingEngine,
	std::string projectDirectory, std::string projectStatus,
	int experimentNumber, int extractorType,
	int numberOfSteps, int batchSize, double learningRate,
	std::string workingDirectory, std::string classNameFilePath)
{
	json j;

	j["Name"] = projectName;
	j["Created"] = getCurrentUtcIso8601();

	if (trainingEngine == 0) {
		j["Training Engine Type"] = "ODHUB";

	}
	else if (trainingEngine == 1) {
		j["Training Engine Type"] = "Object Detection";
	}
	else if (trainingEngine == 2) {
		j["Training Engine Type"] = "Classification";
	}

	else if (trainingEngine == 3) {
		j["Training Engine Type"] = "Segmentation";
	}
	else if (trainingEngine == 4) {
		j["Training Engine Type"] = "Oriented Object Detection";
	}
	else if (trainingEngine == 5) {
		j["Training Engine Type"] = "Abnomal Detection";
	}
	else {
		j["Training Engine Type"] = "Unknown";
	}
	j["Project Directory"] = projectDirectory;
	j["Status"] = projectStatus;


	std::ifstream file(classNameFilePath);
	if (!file.is_open()) {
		std::cerr << "Failed to open classes.names\n";
		return 0;
	}

	std::string line;
	int numberOfClasses = 0;
	while (std::getline(file, line)) {
		if (!line.empty()) {
			++numberOfClasses;
		}
	}

	file.close();
	json trainingParam;
	if (extractorType == 0) {
		trainingParam["Model Type"] = "FAST";
	}
	else if (extractorType == 1) {
		trainingParam["Model Type"] = "ACCURATE";
	}
	else if (extractorType == 2) {
		trainingParam["Model Type"] = "VERY ACCURATE";
	}
	else {
		trainingParam["Model Type"] = "Unknown";
	}
	trainingParam["Number Step"] = numberOfSteps;  // integer
	trainingParam["Batch Size"] = batchSize;  // integer
	trainingParam["Learning Rate"] = learningRate;  // double

	json experiment;
	experiment["Id"] = experimentNumber;  // integer
	experiment["Number Class"] = numberOfClasses;  // integer
	experiment["Training Parameter"] = trainingParam;
	experiment["Evaluation Loss"] = json::array();  // empty array
	experiment["Training Loss"] = json::array();  // empty array
	experiment["Status"] = projectStatus;
	experiment["GPU Model Path"] = "";
	experiment["CPU Model Path"] = "";
	experiment["Training Time(s)"] = 0;  // integer
	j["Experiments"] = json::array({ experiment });
	j["Working Directory"] = workingDirectory;

	json labelMaps;

	std::ifstream infile(classNameFilePath);
	if (!infile.is_open()) {
		std::cerr << "Failed to open classes.names\n";
		return 0;
	}

	int id = 1;
	while (std::getline(infile, line)) {
		if (line.empty()) continue;
		labelMaps.push_back(json::object({
			{"Id", id++},
			{"Name", line}
			}));
	}
	infile.close();
	j["Label Maps"] = labelMaps;

	//	// Write to file
	try {
		std::string ANSTSRoot = "C:\\ProgramData\\ANSCENTER\\ANSTIS\\";
		std::string ANSTSProjectDir;
		if (trainingEngine == 0) {
			ANSTSProjectDir = ANSTSRoot + "odhub\\";
		}
		else if (trainingEngine == 1) {
			ANSTSProjectDir = ANSTSRoot + "ansdl\\";
		}
		else if (trainingEngine == 2) {
			ANSTSProjectDir = ANSTSRoot + "anscl\\";
		}
		else if (trainingEngine == 3) {
			ANSTSProjectDir = ANSTSRoot + "ansseg\\";
		}
		else if (trainingEngine == 4) {
			ANSTSProjectDir = ANSTSRoot + "ansobb\\";
		}
		else {
			ANSTSProjectDir = ANSTSRoot + "ansdl\\";
		}
		std::string projectJson = ANSTSProjectDir + projectName + ".json";

		// Write to file
		std::ofstream outFile(projectJson);
		outFile << std::setw(4) << j << std::endl;
		outFile.close();
	}
	catch (const std::exception& e) {
		std::cerr << "Error writing JSON: " << e.what() << std::endl;
		return -1;
	}

	return 1;

}

int ANSTSCLI(int trainingEngineType, std::string projectWorkingDir, std::string projectName,
	int experimentNumber, int extractorType, int numberOfSteps, int batchSize,
	double learningRate, std::string saveModelPath)
{// 
	ANSCENTER::ANYLTRE ansEngine;
	std::string projectDirectory = "";
	std::string engineFolderDir = "";
	std::string modelTemplateDir = "";
	ANSCENTER::ANSTRE* infHandle;
	//1. Create the Engine handle
	bool initResult = CreateANSTREHandle(&infHandle, "", projectDirectory.c_str(), engineFolderDir.c_str(), modelTemplateDir.c_str(), "", trainingEngineType, false);
	if (initResult) {
		std::cout << "ANSTS Engine instantiation completed successfully." << std::endl;
		std::string strProjects;
		ANSTRE_GetProjects_CPP(&infHandle, strProjects);
		// Set the working directory
		bool setWorkingDirectory = ANSTRE_SetWorkingDirectory(&infHandle, projectWorkingDir.c_str());
		if (!setWorkingDirectory) {
			std::cout << "Error: Project directory path is invalid or inaccessible." << std::endl;
			return -1;
		}
	}
	else {
		std::cout << "Error: Unable to initialize ANSTS Engine." << std::endl;

		int engineInstalled = ANSTRE_CheckEngine(&infHandle);
		if (engineInstalled == 1) {
			std::cout << "Engine Installed:" << engineInstalled << std::endl;
		}
		else if (engineInstalled == 0) {
			std::cout << "Engine Not Installed:" << engineInstalled << std::endl;
		}
		else
			std::cout << "Unknown Error" << std::endl;
		return -1;
	}

	////2. Upload data to ANSTS Training Engine
	ANSTRE_GetProjectDirectory_CPP(&infHandle, projectDirectory);
	std::string projectLocation = projectDirectory + "\\" + projectName;
	if (!std::filesystem::exists(projectLocation)) {
		bool createdProject = ANSTRE_CreateProject(&infHandle, projectName.c_str());
		if (!createdProject) {
			std::cout << "Error: Unable to create project due to invalid configuration or missing parameters." << std::endl;
			return -1;
		}
	}
	std::string projectData = projectLocation + "\\" + "data";
	if (std::filesystem::exists(projectData)) {
		// Delete the existing data directory
		std::filesystem::remove_all(projectData);
	}

	std::cout << "Uploading project data to ANSTS Training Engine...." << std::endl;
	bool dataUploaded = ANSTRE_UploadTrainingData(&infHandle, projectName.c_str());
	if (!dataUploaded || (!std::filesystem::exists(projectData)))
	{
		std::cout << "Error: Unable to upload data. Please check your connection or try again." << std::endl;
		return -1;
	}

	std::cout << "Project data uploaded successfully." << std::endl;
	// Create the JSON file
	std::string projectStatus = "UPLOADED";
	std::string labels = projectData + "\\classes.names";
	CreateANSTSProjectJsonInfo(projectName, trainingEngineType, projectLocation, projectStatus,
		experimentNumber, extractorType, numberOfSteps, batchSize,
		learningRate, projectWorkingDir, labels);
	//////3. Train the model
	bool createEngine = ANSTRE_CreateTrainingEngine(&infHandle, projectName.c_str(), experimentNumber,
		extractorType, numberOfSteps, batchSize, learningRate);
	if (!createEngine) {
		std::cout << "Error: Training engine instance creation failed." << std::endl;
		return -1;
	}

	std::string projectExperimentStatus;
	ANSTRE_GetProjectExperimentStatus_CPP(&infHandle, projectName.c_str(), experimentNumber, projectExperimentStatus);
	//std::cout << projectExperimentStatus << std::endl;

	std::string trainingCommand;
	ANSTRE_GenerateTrainingCommand_CPP(&infHandle, projectName.c_str(), experimentNumber, extractorType, trainingCommand);

	//// Execute the training command
	std::cout << "Initializing training process. This may take a few minutes, please wait..." << std::endl;
	std::thread cmdThread(runCommand, trainingCommand);
	std::thread parserThread(parseOutput, "ANSTS Training: ");
	cmdThread.join();
	parserThread.join();
	
	//// 4. Evaluate the model
	std::string mAPResult;
	ANSTRE_EvaluateModel_CPP(&infHandle, projectName.c_str(), experimentNumber, mAPResult);
	std::cout << "Model evaluation in progress. This may take a moment.. ...." << std::endl;
	std::thread cmdEvalThread(runCommand, mAPResult);
	std::thread parserEvalThread(parseOutput, "");
	cmdEvalThread.join();
	parserEvalThread.join();

	//// 5. Download the model
	// Check if the save model path exists, if not create it
	if (!std::filesystem::exists(saveModelPath)) {
		std::filesystem::create_directory(saveModelPath);
	}
	std::cout << "Exporting CPU model, please wait.." << std::endl;
	int cpuExportResult = ANSTRE_DownloadModel(&infHandle, projectName.c_str(), experimentNumber, saveModelPath.c_str(), 1);
	if (cpuExportResult == 1) {
		std::cout << "CPU model exported successfully." << std::endl;
		std::cout << "Exported model path: " << saveModelPath << std::endl;
	}
	else {
		std::cout << "Error exporting CPU model." << std::endl;
	}
	std::cout << "Exporting GPU model, please wait.." << std::endl;
	int gpuExportResullt = ANSTRE_DownloadModel(&infHandle, projectName.c_str(), experimentNumber, saveModelPath.c_str(), 0);
	if (gpuExportResullt == 1) {
		std::cout << "GPU model exported successfully." << std::endl;
		std::cout << "Exported model path: " << saveModelPath << std::endl;
	}
	else {
		std::cout << "Error exporting GPU model." << std::endl;
	}

	ReleaseANSTREHandle(&infHandle);
	return 0;
}

int main(int argc, char* argv[])
{
	if (argc < 2) {
		std::cerr << "Usage: " << argv[0] << " <rootProjectDir>\n";
		return -1;
	}
	std::string rootProjectDir = argv[1];
	std::string configFilePath = rootProjectDir + "\\ansts_config.json";
	std::string projectName = argv[2];

	// Check if the config file exists
	if (!std::filesystem::exists(configFilePath)) {
		std::cerr << "Config file not found: " << configFilePath << std::endl;
		return -1;
	}
	ptree pt;
	int trainingEngineType = 1; // 0: ANSYOLO, 1: ANSYOLOCL, 2: ANSYOLOSEG
	std::string projectWorkingDir = rootProjectDir;
	int experimentNumber = 1;
	int extractorType = 1; // 0: ANSYOLO, 1: ANSYOLOCL, 2: ANSYOLOSEG
	int numberOfSteps = 1000;
	int batchSize = 1;
	double learningRate = 0.001;
	std::string saveModelPath = rootProjectDir+"\\Models";
	std::cout << "Project Diectory: " << projectWorkingDir << std::endl;
	try {
		read_json(configFilePath, pt);

		std::string trainingEngine = pt.get<std::string>("trainingEngine");
		if(projectWorkingDir.empty())projectWorkingDir = pt.get<std::string>("projectDirectory");
		if(projectName.empty())projectName = pt.get<std::string>("projectName");
		experimentNumber = pt.get<int>("experimentNumber");
		std::string modelType = pt.get<std::string>("modelType");
		numberOfSteps = pt.get<int>("numberOfSteps");
		batchSize = pt.get<int>("batchSize");
		learningRate = pt.get<double>("learningRate");

		// Parse the training engine type
		if (trainingEngine == "ODHUB") { // ODHUB
			trainingEngineType = 0;
		}
		else if (trainingEngine == "ANSOD") { // Object Detection
			trainingEngineType = 1;
		}
		else if (trainingEngine == "ANSCL") { // Classification
			trainingEngineType = 2;
		}
		else if (trainingEngine == "ANSSEG") { // Segmentation
			trainingEngineType = 3;
		}
		else if (trainingEngine == "ANSOBB") { // Oriented Bounding Box
			trainingEngineType = 4;
		}
		else {
			trainingEngineType = 1;
		}

		// Parse the model type
		if (modelType == "FAST") { // N
			extractorType = 0;
		}
		else if (modelType == "ACCURATE") { // M
			extractorType = 1;
		}
		else if (modelType == "VERYACCURATE") { // L
			extractorType = 2;
		}
		else if (modelType == "CUSTOM") { // L
			extractorType = 3;
		}
		else {
			extractorType = 0;
		}
	}
	catch (const std::exception& e) {
		std::cerr << "Error reading config.json: " << e.what() << std::endl;
	}
	// Call the ANSTSCLI function
	ANSTSCLI(trainingEngineType, projectWorkingDir, projectName,
		experimentNumber, extractorType, numberOfSteps, batchSize,
		learningRate, saveModelPath);
	return 0;
}

