#include "ANSOCRBase.h"
#include "Utility.h"
#include <opencv2/highgui.hpp>
#include <omp.h>
#include <json.hpp>
#include "ANSLibsLoader.h"

static bool ansocrLicenceValid = false;
// Global once_flag to protect license checking
static std::once_flag ansocrLicenseOnceFlag;
template <typename T>
T GetData(const boost::property_tree::ptree& pt, const std::string& key)
{
	T ret;
	if (boost::optional<T> data = pt.get_optional<T>(key))
	{
		ret = data.get();
	}
	return ret;
}

namespace ANSCENTER {
	/// <summary>
	/// Base class
	/// </summary>
	/// 
	/// 
	static void VerifyGlobalANSOCRLicense(const std::string& licenseKey) {
		try {
			ansocrLicenceValid = ANSCENTER::ANSLicenseHelper::LicenseVerification(licenseKey, 1005, "ANSOCR");//Default productId=1006
			if (!ansocrLicenceValid) { // we also support ANSTS license
				ansocrLicenceValid = ANSCENTER::ANSLicenseHelper::LicenseVerification(licenseKey, 1003, "ANSVIS");//Default productId=1003 (ANSVIS)
			}
		}
		catch (std::exception& e) {
			ansocrLicenceValid = false;
		}
	}
	void ANSOCRBase::CheckLicense() {
		try {
			// Check once globally
			std::call_once(ansocrLicenseOnceFlag, [this]() {
				VerifyGlobalANSOCRLicense(_licenseKey);
				});

			// Update this instance's local license flag
			_licenseValid = ansocrLicenceValid;
		}
		catch (const std::exception& e) {
			this->_logger.LogFatal("ANSOCRBase::CheckLicense. Error:", e.what(), __FILE__, __LINE__);
		}
	}

	bool ANSOCRBase::Init(const std::string& licenseKey, OCRModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, int engineMode) {
		try {
			ANSCENTER::ANSLibsLoader::Initialize();

			_licenseKey = licenseKey;
			_engineMode = engineMode;
			_licenseValid = false;
			_modelFolder = "";
			_modelConfigFile = "";
			_modelConfig = modelConfig;
			_modelFolder.clear();
			_modelConfigFile.clear();

			CheckLicense();
			if (!_licenseValid) {
				this->_logger.LogError("ANSOCRBase::Initialize", "Invalid License", __FILE__, __LINE__);
				return false;
			}
			_licenseValid = true;


			// 0. Check if the modelZipFilePath exist?
			if (!FileExist(modelZipFilePath)) {
				this->_logger.LogFatal("ANSOCRBase::Initialize", "Model zip file is not exist", __FILE__, __LINE__);
				return false;
			}
			// 1. Unzip model zip file to a special location with folder name as model file (and version)
			std::string outputFolder;
			std::vector<std::string> passwordArray;
			if (!modelZipPassword.empty()) passwordArray.push_back(modelZipPassword);
			passwordArray.push_back("AnsDemoModels20@!");
			passwordArray.push_back("Sh7O7nUe7vJ/417W0gWX+dSdfcP9hUqtf/fEqJGqxYL3PedvHubJag==");
			passwordArray.push_back("3LHxGrjQ7kKDJBD9MX86H96mtKLJaZcTYXrYRdQgW8BKGt7enZHYMg==");
			std::string modelName = GetFileNameWithoutExtension(modelZipFilePath);
			//this->_logger.LogDebug("ANSOCRBase::Initialize. Model name", modelName, __FILE__, __LINE__);
			size_t vectorSize = passwordArray.size();
			for (size_t i = 0; i < vectorSize; i++) {
				if (ExtractPasswordProtectedZip(modelZipFilePath, passwordArray[i], modelName, _modelFolder, false))
					break; // Break the loop when the condition is met. 
			}
			// 2. Check if the outputFolder exist
			if (!std::filesystem::exists(_modelFolder)) {
				this->_logger.LogError("ANSOCRBase::Initialize. Output model folder is not exist", modelName, __FILE__, __LINE__);
				return false; // That means the model file is not exist or the password is not correct
			}

			return true;
		}
		catch (std::exception& e) {
			this->_logger.LogFatal("ANSOCRBase::Initialize", e.what(), __FILE__, __LINE__);
			return false;
		}
	}

	bool ANSOCRBase::Initialize(const std::string& licenseKey, OCRModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, int engineMode) {
		try {

			_licenseKey = licenseKey;
			_engineMode = engineMode;
			_licenseValid = false;
			_modelFolder = "";
			_modelConfigFile = "";
			_modelConfig = modelConfig;
			_modelFolder.clear();
			_modelConfigFile.clear();

			CheckLicense();
			if (!_licenseValid) {
				this->_logger.LogError("ANSOCRBase::Initialize", "Invalid License", __FILE__, __LINE__);
				return false;
			}
			_licenseValid = true;


			// 0. Check if the modelZipFilePath exist?
			if (!FileExist(modelZipFilePath)) {
				this->_logger.LogFatal("ANSOCRBase::Initialize", "Model zip file is not exist", __FILE__, __LINE__);
				return false;
			}
			// 1. Unzip model zip file to a special location with folder name as model file (and version)
			std::string outputFolder;
			std::vector<std::string> passwordArray;
			if (!modelZipPassword.empty()) passwordArray.push_back(modelZipPassword);
			passwordArray.push_back("AnsDemoModels20@!");
			passwordArray.push_back("Sh7O7nUe7vJ/417W0gWX+dSdfcP9hUqtf/fEqJGqxYL3PedvHubJag==");
			passwordArray.push_back("3LHxGrjQ7kKDJBD9MX86H96mtKLJaZcTYXrYRdQgW8BKGt7enZHYMg==");
			std::string modelName = GetFileNameWithoutExtension(modelZipFilePath);
			//this->_logger.LogDebug("ANSOCRBase::Initialize. Model name", modelName, __FILE__, __LINE__);
			size_t vectorSize = passwordArray.size();
			for (size_t i = 0; i < vectorSize; i++) {
				if (ExtractPasswordProtectedZip(modelZipFilePath, passwordArray[i], modelName, _modelFolder, false))
					break; // Break the loop when the condition is met. 
			}
			// 2. Check if the outputFolder exist
			if (!std::filesystem::exists(_modelFolder)) {
				this->_logger.LogError("ANSOCRBase::Initialize. Output model folder is not exist", modelName, __FILE__, __LINE__);
				return false; // That means the model file is not exist or the password is not correct
			}
			// 3. Check if the model has the configuration file
			std::string modelConfigName = "model_config.json";
			_modelConfigFile = CreateFilePath(_modelFolder, modelConfigName);

			//4. For now we do have the model folder so we will assign paths to OCR models
			_modelConfig.detectionModelDir = _modelFolder;
			_modelConfig.recognizerModelDir = _modelFolder;
			_modelConfig.clsModelDir = _modelFolder;
			_modelConfig.layoutModelDir = _modelFolder;
			_modelConfig.layourDictionaryPath = _modelFolder;
			_modelConfig.tableModelDir = _modelFolder;
			_modelConfig.tableCharDictionaryPath = _modelFolder;
			_modelConfig.recogizerCharDictionaryPath = CreateFilePath(_modelFolder, "dict_ch.txt");

			_modelConfig.detectionModelFile = CreateFilePath(_modelFolder, "ansocrdec.onnx");
			_modelConfig.detectionModelParam = CreateFilePath(_modelFolder, "ansocrdec.onnx");
			_modelConfig.clsModelFile = CreateFilePath(_modelFolder, "ansocrcls.onnx");
			_modelConfig.clsModelParam = CreateFilePath(_modelFolder, "ansocrcls.onnx");
			_modelConfig.recognizerModelFile = CreateFilePath(_modelFolder, "ansocrrec.onnx");
			_modelConfig.recognizerModelParam = CreateFilePath(_modelFolder, "ansocrrec.onnx");
			// For now we do have _modelConfig and _modelFolder
			return true;
		}
		catch (std::exception& e) {
			this->_logger.LogFatal("ANSOCRBase::Initialize", e.what(), __FILE__, __LINE__);
			return false;
		}
	}
	std::string ANSCENTER::ANSOCRUtility::OCRDetectionToJsonString(const std::vector<OCRObject>& dets)
	{
		if (dets.empty()) {
			return R"({"results":[]})";
		}

		try {
			nlohmann::json root;
			auto& results = root["results"] = nlohmann::json::array();

			for (const auto& det : dets) {
				results.push_back({
					{"class_id", std::to_string(det.classId)},
					{"track_id", std::to_string(det.trackId)},
					{"class_name", det.className},
					{"prob", std::to_string(det.confidence)},
					{"x", std::to_string(det.box.x)},
					{"y", std::to_string(det.box.y)},
					{"width", std::to_string(det.box.width)},
					{"height", std::to_string(det.box.height)},
					{"mask", ""},  // TODO: convert masks to comma separated string
					{"extra_info", det.extraInfo},
					{"camera_id", det.cameraId},
					{"polygon", PolygonToString(det.polygon)},
					{"kps", KeypointsToString(det.kps)}
					});
			}

			return root.dump();
		}
		catch (const std::exception& e) {
			// Add your error logging here if needed
			return R"({"results":[],"error":"Serialization failed"})";
		}
	}

	std::vector<cv::Rect> ANSCENTER::ANSOCRUtility::GetBoundingBoxes(const std::string& strBBoxes) {
		std::vector<cv::Rect> bBoxes;
		bBoxes.clear();
		std::stringstream ss;
		ss << strBBoxes;
		boost::property_tree::ptree pt;
		boost::property_tree::read_json(ss, pt);
		BOOST_FOREACH(const boost::property_tree::ptree::value_type & child, pt.get_child("results"))
		{
			const boost::property_tree::ptree& result = child.second;
			const auto x = GetData<float>(result, "x");
			const auto y = GetData<float>(result, "y");
			const auto width = GetData<float>(result, "width");
			const auto height = GetData<float>(result, "height");
			cv::Rect rectTemp;
			rectTemp.x = x;
			rectTemp.y = y;
			rectTemp.width = width;
			rectTemp.height = height;
			bBoxes.push_back(rectTemp);
		}
		return bBoxes;
	}

	std::string ANSCENTER::ANSOCRUtility::PolygonToString(const std::vector<cv::Point2f>& polygon) {
		if (polygon.empty()) {
			return "";
		}

		std::string result;
		result.reserve(polygon.size() * 20);

		char buffer[64];
		for (size_t i = 0; i < polygon.size(); ++i) {
			if (i > 0) {
				snprintf(buffer, sizeof(buffer), ";%.3f;%.3f", polygon[i].x, polygon[i].y);
			}
			else {
				snprintf(buffer, sizeof(buffer), "%.3f;%.3f", polygon[i].x, polygon[i].y);
			}
			result += buffer;
		}

		return result;
	}
	std::string ANSCENTER::ANSOCRUtility::KeypointsToString(const std::vector<float>& kps) {
		if (kps.empty()) {
			return "";
		}

		std::string result;
		result.reserve(kps.size() * 10);

		char buffer[32];
		for (size_t i = 0; i < kps.size(); ++i) {
			if (i > 0) result += ';';
			snprintf(buffer, sizeof(buffer), "%.3f", kps[i]);
			result += buffer;
		}

		return result;
	}
	std::vector<cv::Point2f> ANSCENTER::ANSOCRUtility::RectToNormalizedPolygon(const cv::Rect& rect, float imageWidth, float imageHeight) {
		// Ensure imageWidth and imageHeight are non-zero to avoid division by zero
		if (imageWidth <= 0 || imageHeight <= 0) {
			std::vector<cv::Point2f> emptyPolygon;
			return emptyPolygon;
		}

		// Calculate normalized points for each corner of the rectangle
		std::vector<cv::Point2f> polygon = {
			{ rect.x / imageWidth, rect.y / imageHeight },                     // Top-left
			{ (rect.x + rect.width) / imageWidth, rect.y / imageHeight },      // Top-right
			{ (rect.x + rect.width) / imageWidth, (rect.y + rect.height) / imageHeight }, // Bottom-right
			{ rect.x / imageWidth, (rect.y + rect.height) / imageHeight }      // Bottom-left
		};

		return polygon;
	}

};