#include "ANSCustomRVA.h"

struct ExtractedInfo {
	std::string revision;
	std::string model;
	std::string serialNumber;
};
ExtractedInfo extractDeviceInfo(const std::string& jsonText) {
	ExtractedInfo info;
	std::smatch match;

	std::regex revPattern("\"class_name\":\"Rev:([^\"]+)\"");
	if (std::regex_search(jsonText, match, revPattern)) {
		info.revision = match[1].str();
	}

	std::regex modelPattern("\"class_name\":\"Model:([^\"]+)\"");
	if (std::regex_search(jsonText, match, modelPattern)) {
		info.model = match[1].str();
	}

	std::regex snPattern("\"class_name\":\"S[\\\\\/]*N:([^\"]+)\"");
	if (std::regex_search(jsonText, match, snPattern)) {
		info.serialNumber = match[1].str();
	}

	return info;
}
std::string extractQRCodeValue(const std::string& jsonText) {
	std::smatch match;

	// Step 1: Find the entire QRCode object
	std::regex objectPattern("\\{[^{}]*\"extra_info\":\"QRCode\"[^{}]*\\}");

	if (std::regex_search(jsonText, match, objectPattern)) {
		std::string qrObject = match[0].str();

		// Step 2: Extract class_name from that object
		std::regex namePattern("\"class_name\":\"([^\"]+)\"");
		std::smatch nameMatch;

		if (std::regex_search(qrObject, nameMatch, namePattern)) {
			return nameMatch[1].str();
		}
	}

	return "";
}

std::string toUpperCase(const std::string& input) {
	std::string result = input;
	std::transform(result.begin(), result.end(), result.begin(),
		[](unsigned char c) { return std::toupper(c); });
	return result;
}
ANSCustomRVA::ANSCustomRVA()
{
	_isInitialized = false;
	_readROIs = false;
	// Default values
	_numberOfScrews = 4;
	_numberOfScrewsPCB = 1;
	_numberOfPowerCable = 1;
	_numberOfBMPRO = 1;
	_numberOfLabel = 1;
	_numberOfPlasticPackge = 1;
}
ANSCustomRVA::~ANSCustomRVA()
{
	Destroy();
}

bool ANSCustomRVA::Destroy()
{
	try {
		_detector.reset();
		if (ansOCRHandle != nullptr) { ReleaseANSOCRHandle(&ansOCRHandle); ansOCRHandle = nullptr; }
		if (ansOpenCVHandle != nullptr) { ReleaseANSCVHandle(&ansOpenCVHandle); ansOpenCVHandle = nullptr; }
		return true;
	}
	catch (const std::exception& e) {
		return false;
	}
	catch (...) {
		return false;
	}
}
bool ANSCustomRVA::OptimizeModel(bool fp16)
{
	try {
		if (!_detector) return false;
		if (_detector->Optimize(fp16) != 0) return false;
		return true;
	}
	catch (const std::exception& e) {
		return false;
	}
	catch (...) {
		return false;
	}
}

bool ANSCustomRVA::Initialize(const std::string& modelDirectory, float detectionScoreThreshold, std::string& labelMap)
{
	try {
		_modelDirectory = modelDirectory;
		_detectionScoreThreshold = detectionScoreThreshold;

		// Initialize the OCR engine
		if (ansOCRHandle != nullptr) ReleaseANSOCRHandle(&ansOCRHandle);
		std::string licenseKey = "";
		std::string ocrModelFilePath = "C:\\ProgramData\\ANSCENTER\\Shared\\ANS_GenericOCR_v1.0.zip";
		int language = 0; // English
		int engine = 0;
		int ocrCreateResult = CreateANSOCRHandle(&ansOCRHandle, licenseKey.c_str(), ocrModelFilePath.c_str(), "", language, engine);

		// Initialize ANSCV
		if (ansOpenCVHandle != nullptr) ReleaseANSCVHandle(&ansOpenCVHandle);
		int qrCreateResult = CreateANSCVHandle(&ansOpenCVHandle, licenseKey.c_str());

		// Create detector through ANSLIB factory
		_detector = ANSLIBPtr(ANSCENTER::ANSLIB::Create(), &ANSCENTER::ANSLIB::Destroy);
		if (!_detector) {
			_isInitialized = false;
			return false;
		}

		// Check hardware type
		engineType = _detector->GetEngineType();
		if (engineType == 0) {
			_detectorModelType = 4;     // TENSORRT
			_detectorDetectionType = 1; // DETECTION
		}
		else {
			_detectorModelType = 5;     // OPENVINO
			_detectorDetectionType = 1; // DETECTION
		}

		if (_detectionScoreThreshold < 0.25f) _detectionScoreThreshold = 0.25f;

		float detConfThreshold = 0.5f;
		float detNMSThreshold = 0.5f;
		bool autoEngineDetection = true;
		bool loadEngineOnCreation = _loadEngineOnCreate;
#ifdef FNS_DEBUG
		loadEngineOnCreation = true;
#endif

		std::string detLabelMap;
		int detResult = _detector->LoadModelFromFolder(
			licenseKey.c_str(),
			"rva",
			"rva.names",
			_detectionScoreThreshold,
			detConfThreshold,
			detNMSThreshold,
			autoEngineDetection,
			_detectorModelType,
			_detectorDetectionType,
			loadEngineOnCreation,
			modelDirectory.c_str(),
			detLabelMap);

		if (detResult != 1) {
			_isInitialized = false;
			return false;
		}

		labelMap = detLabelMap;

		if (ocrCreateResult * qrCreateResult > 0) {
			_isInitialized = true;
		}
		else {
			_isInitialized = false;
			return false;
		}
		return true;
	}
	catch (const std::exception& e) {
		return false;
	}
	catch (...) {
		return false;
	}
}

bool ANSCustomRVA::ConfigureParameters(CustomParams& param) {
	// Configure parameters for the model, if needed
	param.ROI_Config.clear();
	param.ROI_Options.clear();
	param.Parameters.clear();
	param.ROI_Values.clear();

	CustomParameter numberOfScrews;
	numberOfScrews.Name = "No_Screws";
	numberOfScrews.DataType = "Integer";
	numberOfScrews.NoOfDecimals = 0;
	numberOfScrews.MaxValue = 100;
	numberOfScrews.MinValue = 0;
	numberOfScrews.StartValue = "4";
	numberOfScrews.ListItems.clear();
	numberOfScrews.DefaultValue = "4";
	param.Parameters.push_back(numberOfScrews);

	CustomParameter numberOfScrewsPCB;
	numberOfScrewsPCB.Name = "No_Screws_PCB";
	numberOfScrewsPCB.DataType = "Integer";
	numberOfScrewsPCB.NoOfDecimals = 0;
	numberOfScrewsPCB.MaxValue = 100;
	numberOfScrewsPCB.MinValue = 0;
	numberOfScrewsPCB.StartValue = "1";
	numberOfScrewsPCB.ListItems.clear();
	numberOfScrewsPCB.DefaultValue = "1";
	param.Parameters.push_back(numberOfScrewsPCB);

	CustomParameter numberOfPowerCable;
	numberOfPowerCable.Name = "No_Power_Cable";
	numberOfPowerCable.DataType = "Integer";
	numberOfPowerCable.NoOfDecimals = 0;
	numberOfPowerCable.MaxValue = 100;
	numberOfPowerCable.MinValue = 0;
	numberOfPowerCable.StartValue = "1";
	numberOfPowerCable.ListItems.clear();
	numberOfPowerCable.DefaultValue = "1";
	param.Parameters.push_back(numberOfPowerCable);

	CustomParameter numberOfBMPPRO;
	numberOfBMPPRO.Name = "No_BMPRO";
	numberOfBMPPRO.DataType = "Integer";
	numberOfBMPPRO.NoOfDecimals = 0;
	numberOfBMPPRO.MaxValue = 100;
	numberOfBMPPRO.MinValue = 0;
	numberOfBMPPRO.StartValue = "1";
	numberOfBMPPRO.ListItems.clear();
	numberOfBMPPRO.DefaultValue = "1";
	param.Parameters.push_back(numberOfBMPPRO);

	CustomParameter numberOfLabel;
	numberOfLabel.Name = "No_Label";
	numberOfLabel.DataType = "Integer";
	numberOfLabel.NoOfDecimals = 0;
	numberOfLabel.MaxValue = 100;
	numberOfLabel.MinValue = 0;
	numberOfLabel.StartValue = "1";
	numberOfLabel.ListItems.clear();
	numberOfLabel.DefaultValue = "1";
	param.Parameters.push_back(numberOfLabel);

	CustomParameter numberOfPlasticPackage;
	numberOfPlasticPackage.Name = "No_PlasticPackage";
	numberOfPlasticPackage.DataType = "Integer";
	numberOfPlasticPackage.NoOfDecimals = 0;
	numberOfPlasticPackage.MaxValue = 100;
	numberOfPlasticPackage.MinValue = 0;
	numberOfPlasticPackage.StartValue = "1";
	numberOfPlasticPackage.ListItems.clear();
	numberOfPlasticPackage.DefaultValue = "1";
	param.Parameters.push_back(numberOfPlasticPackage);

	return true;
}

std::vector<CustomObject> ANSCustomRVA::RunInference(const cv::Mat& input)
{
	return RunInference(input, "CustomCam");
}


std::vector<CustomObject> ANSCustomRVA::RunInference(const cv::Mat& input, const std::string& camera_id) {
	std::lock_guard<std::recursive_mutex> lock(_mutex);
	std::vector<CustomObject> results;

	if (!_isInitialized) {
		return results;
	}
	if (input.empty() || input.cols < 10 || input.rows < 10) {
		return results;
	}

	try {

#ifdef FNS_DEBUG
		// Default values for debug
		_numberOfScrews = 4;
		_numberOfScrewsPCB = 1;
		_numberOfPowerCable = 1;
		_numberOfBMPRO = 1;
		_numberOfLabel = 1;
		_numberOfPlasticPackge = 1;
#else
		//1. Read the ROI parameters only once
		if (!_readROIs) {
			for (auto& param : _params.Parameters) {
				if (param.Name.find("No_Screws") != std::string::npos) {
					_numberOfScrews = std::stoi("4");// std::stoi(param.Value);
				}
				else if (param.Name.find("No_Screws_PCB") != std::string::npos) {
					_numberOfScrewsPCB = std::stoi("1");// std::stoi(param.Value);
				}
				else if (param.Name.find("No_Power_Cable") != std::string::npos) {
					_numberOfPowerCable = std::stoi(param.Value);
				}
				else if (param.Name.find("No_BMPRO") != std::string::npos) {
					_numberOfBMPRO = std::stoi(param.Value);
				}
				else if (param.Name.find("No_Label") != std::string::npos) {
					_numberOfLabel = std::stoi(param.Value);
				}
				else if (param.Name.find("No_PlasticPackage") != std::string::npos) {
					_numberOfPlasticPackge = std::stoi(param.Value);
				}
			}
			_readROIs = true;
		}
#endif
		// 2. Do general detections
		const cv::Mat& frame = input;
		std::vector<ANSCENTER::Object> detectionResults;
		_detector->RunInference(frame, camera_id.c_str(), detectionResults);
		if (detectionResults.empty()) {
			return results;
		}
		int numOfScrews = 0;
		int numOfScrewsPCB = 0;
		int numOfPowerCable = 0;
		int numOfBMPRO = 0;
		int numOfLabel = 0;
		int numOfPlasticPackage = 0;
		std::string qrCodeData;
		std::string ocrText;
		// Count the number of detected objects
		// 3. Process the detection results
		for (const auto& detection : detectionResults) {
			CustomObject customObject;
			customObject.classId = detection.classId;
			customObject.confidence = detection.confidence;
			customObject.box = detection.box;
			customObject.cameraId = camera_id;
			// Add additional information based on the class ID
			switch (detection.classId)
			{
				case 0: // Screw
					numOfScrews++;
					break;
				case 1: // Screw PCB
					numOfScrewsPCB++;
					break;
				case 2: // Power Cable
					numOfPowerCable++;
					break;
				case 3: // BMPRO
					numOfBMPRO++;
					break;
				case 4: // Label
				{
					numOfLabel++;
					cv::Rect ocrBox(customObject.box.x, customObject.box.y, 4 * customObject.box.width / 5, 1.3 * customObject.box.height / 4);
					cv::Rect qrBox(customObject.box.x, customObject.box.y + customObject.box.height / 5, 1 * customObject.box.width / 4, 2 * customObject.box.height / 5);
					cv::Mat ocrImage = frame(ocrBox);
					cv::Mat qrImage = frame(qrBox);
					for (int i = 0; i <= 5; i++) {
						ocrText.clear();  // Bug fix: was ocrText.empty() which doesn't clear
						qrCodeData.clear();
						RunInferenceCV(&ansOCRHandle, ocrImage, ocrText);
						ANSCV_QRDecoderCV(&ansOpenCVHandle, qrImage, qrCodeData);
						ExtractedInfo ocrInfoDetail = extractDeviceInfo(ocrText);
						std::string qrInfoDetail = extractQRCodeValue(qrCodeData);
						if (!ocrText.empty() &&
							!qrCodeData.empty() &&
							!qrInfoDetail.empty() &&
							!ocrInfoDetail.model.empty() &&
							!ocrInfoDetail.revision.empty()) {
							break; // Exit the loop if OCR or QR code data is found
						}
						// sleep for a short duration to allow OCR/QR code detection to stabilize
						std::this_thread::sleep_for(std::chrono::milliseconds(100));
					}
#ifdef FNS_DEBUG
					std::cout << "OCR Text: " << ocrText << std::endl;
					std::cout << "QR Text: " << qrCodeData << std::endl;
					cv::imshow("OCR Image", ocrImage);
					cv::imshow("QR Image", qrImage);
					cv::waitKey(1);
#endif
					ocrImage.release();
					qrImage.release();
					break;
				}
				case 5: // Plastic Package
					numOfPlasticPackage++;
					break;
			}
			results.push_back(customObject);
		}

		// 4. Check if the number of detected objects matches the expected counts
		CustomObject passfailObject;

		passfailObject.confidence = 1;
		passfailObject.box.x = 50;
		passfailObject.box.y = 50;
		passfailObject.box.width = input.cols - 50;
		passfailObject.box.height = input.rows - 50;
		passfailObject.cameraId = camera_id;
		std::string extraInfo;
		ExtractedInfo ocrInfo = extractDeviceInfo(ocrText);
		std::string qrInfo = extractQRCodeValue(qrCodeData);
		// we need to extract SN and Rev instead using ocr information
		std::string sn = qrInfo;
		if (sn.empty()) sn = ocrInfo.serialNumber;
		extraInfo = "SN:" + sn + "; Model:" + ocrInfo.model + "; REV:" + ocrInfo.revision;

		if ((numOfScrews == _numberOfScrews) &&
			(numOfScrewsPCB == _numberOfScrewsPCB) &&
			(numOfPowerCable == _numberOfPowerCable) &&
			(numOfBMPRO == _numberOfBMPRO) &&
			(numOfLabel == _numberOfLabel) &&
			(numOfPlasticPackage == _numberOfPlasticPackge))
		{
			passfailObject.classId = 6; // PASS
			passfailObject.className = "PASS";
		}
		else {
			passfailObject.classId = 7; // FAILED
			passfailObject.className = "FAIL";
			std::string summary = "; Screws:" + std::to_string(numOfScrews) + " out of " + std::to_string(_numberOfScrews) +
				"; Screws_PCB:" + std::to_string(numOfScrewsPCB) + " out of " + std::to_string(_numberOfScrewsPCB) +
				"; Power Cables:" + std::to_string(numOfPowerCable) + " out of " + std::to_string(_numberOfPowerCable) +
				"; BMPROs:" + std::to_string(numOfBMPRO) + " out of " + std::to_string(_numberOfBMPRO) +
				"; Labels:" + std::to_string(numOfLabel) + " out of " + std::to_string(_numberOfLabel) +
				"; Plastic Packages: " + std::to_string(numOfPlasticPackage) + " of " + std::to_string(_numberOfPlasticPackge);
			extraInfo = extraInfo + summary;
		}

		passfailObject.extraInfo = extraInfo;
		results.push_back(passfailObject);
#ifdef FNS_DEBUG
		std::cout << "RVA Attributes:" << extraInfo << std::endl;
#endif

	}
	catch (const std::exception& e) {
	}
	catch (...) {
	}

	return results;
}
