/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#include "ArFrameMetadata.hpp"
#include "Extensions/RymlConverters.hpp"
#include <fmt/format.h>
#include <pugixml.hpp>
#include <sstream>

namespace OpenVulkano::AR
{
	namespace
	{
		template<typename T, int S>
		T ReadMat(const pugi::xml_node& matrixNode)
		{
			T mat(0);
			auto rowNode = matrixNode.first_child();
			for (int row = 0; row < S && rowNode; row++, rowNode = rowNode.next_sibling())
			{
				auto columnNode = rowNode.first_child();
				for (int column = 0; column < S && columnNode; column++, columnNode = columnNode.next_sibling())
				{
					mat[column][row] = columnNode.text().as_float();
				}
			}
			return mat;
		}

		template<glm::length_t SIZE, typename T, glm::qualifier Q>
		void MatToXML(const glm::mat<SIZE, SIZE, T, Q>& mat, std::ostream& stream, const std::string& nl = "\n")
		{
			auto nlRow = nl + "\t";
			stream <<  nl;
			for(int r = 0 ; r < SIZE; r++)
			{
				stream << "<row>" << nlRow;
				for (int c = 0; c < SIZE; c++)
				{
					stream << "<column>" << mat[c][r] << "</column>";
				}
				stream << nl << "</row>";
			}
		}

		template<typename T, int S>
		void ReadMat(const ryml::NodeRef& node, T& mat)
		{
			if (!node.is_seq()) return;
			int row = 0;
			for(auto iter = node.begin(); iter != node.end() && row < S; ++iter)
			{
				const auto rowNode = *iter;
				int col = 0;
				for(auto colIter = rowNode.begin(); colIter != rowNode.end() && col < S; ++colIter)
				{
					(*colIter) >> mat[col][row];
					col++;
				}
				row++;
			}
		}
	}

	ArFrameMetadata ArFrameMetadata::FromXML(const char* xml, size_t length)
	{
		ArFrameMetadata frameData;

		pugi::xml_document doc;
		pugi::xml_parse_result result = doc.load_buffer(xml, length);
		if (!result) throw std::runtime_error("Failed to parse frame metadata");
		pugi::xml_node nodeFrame = doc.child("arframe");

		pugi::xml_node nodeCamera = nodeFrame.child("camera");
		frameData.transformation = ReadMat<Math::Matrix4f, 4>(nodeCamera.child("transform"));
		frameData.projection = ReadMat<Math::Matrix4f, 4>(nodeCamera.child("projection"));
		pugi::xml_node nodeRes = nodeCamera.child("resolution");
		Math::Vector2i resolution = { nodeRes.child("width").text().as_int(), nodeRes.child("height").text().as_int() };
		frameData.intrinsic = { ReadMat<Math::Matrix3f, 3>(nodeCamera.child("intrinsics")), resolution };
		frameData.exposureTime = nodeCamera.child("exposureDuration").text().as_float();
		frameData.exposureOffset = nodeCamera.child("exposureOffset").text().as_float();

		pugi::xml_node nodeLight = nodeFrame.child("light");
		frameData.lightIntensity = nodeLight.child("ambientIntensity").text().as_float();
		frameData.lightColorTemp = nodeLight.child("ambientColorTemp").text().as_float();

		frameData.timestamp = nodeFrame.child("timestamp").text().as_double();
		frameData.trackingState = ArTrackingState::GetFromName(nodeFrame.child("trackingState").child("camera").text().as_string());

		return frameData;
	}

	ArFrameMetadata ArFrameMetadata::FromYaml(const char* yaml, size_t length)
	{
		ArFrameMetadata frameData;
		ryml::Tree tree;
		ryml::parse_in_arena(c4::csubstr(yaml, length), &tree);
		ryml::NodeRef root = tree.rootref();
		ryml::NodeRef camNode = root["Camera"];
		ReadMat<Math::Matrix4f, 4>(camNode["Transform"], frameData.transformation);
		ReadMat<Math::Matrix4f, 4>(camNode["Projection"], frameData.projection);
		Math::Matrix3f intrinsic;
		Math::Vector2i res;
		camNode["Resolution"] >> res;
		ReadMat<Math::Matrix3f, 3>(camNode["Intrinsics"], intrinsic);
		frameData.intrinsic = {intrinsic, res};
		camNode["ExposureDuration"] >> frameData.exposureTime;
		camNode["ExposureOffset"] >> frameData.exposureOffset;
		if (camNode["FocalLength"].readable()) camNode["FocalLength"] >> frameData.focalLength;
		if (camNode["FNumber"].readable()) camNode["FNumber"] >> frameData.fNumber;
		uint64_t nanos;
		root["Timestamp"] >> nanos;
		frameData.timestamp = nanos;
		root["TimestampDepth"] >> nanos;
		frameData.timestampDepth = nanos;
		std::string tracking;
		ryml::NodeRef trackingNode = root["TrackingState"];
		trackingNode["Camera"] >> tracking;
		frameData.trackingState = ArTrackingState::GetFromName(tracking);
		ryml::NodeRef lightNode = root["Light"];
		lightNode["Intensity"] >> frameData.lightIntensity;
		lightNode["ColorTemp"] >> frameData.lightColorTemp;
		return frameData;
	}

	std::string ArFrameMetadata::ToYaml() const
	{
		const auto& camMat = intrinsic.GetMatrix();
		std::string meta = fmt::format(R"(Camera:
  Transform:
    - [ {}, {}, {}, {} ]
    - [ {}, {}, {}, {} ]
    - [ {}, {}, {}, {} ]
    - [ {}, {}, {}, {} ]
  Projection:
    - [ {}, {}, {}, {} ]
    - [ {}, {}, {}, {} ]
    - [ {}, {}, {}, {} ]
    - [ {}, {}, {}, {} ]
  Resolution: '({},{})'
  Intrinsics:
    - [ {}, {}, {} ]
    - [ {}, {}, {} ]
    - [ {}, {}, {} ]
  ExposureDuration: {}
  ExposureOffset: {}
  FocalLength: {}
  FNumber: {}
Timestamp: {}
TimestampDepth: {}
TrackingState:
  Camera: {}
Light:
  Intensity: {}
  ColorTemp: {}
)",
	transformation[0][0], transformation[1][0], transformation[2][0], transformation[3][0],
	transformation[0][1], transformation[1][1], transformation[2][1], transformation[3][1],
	transformation[0][2], transformation[1][2], transformation[2][2], transformation[3][2],
	transformation[0][3], transformation[1][3], transformation[2][3], transformation[3][3],
	projection[0][0], projection[1][0], projection[2][0], projection[3][0],
	projection[0][1], projection[1][1], projection[2][1], projection[3][1],
	projection[0][2], projection[1][2], projection[2][2], projection[3][2],
	projection[0][3], projection[1][3], projection[2][3], projection[3][3],
	intrinsic.GetResolution().x, intrinsic.GetResolution().y,
	camMat[0][0], camMat[1][0], camMat[2][0],
	camMat[0][1], camMat[1][1], camMat[2][1],
	camMat[0][2], camMat[1][2], camMat[2][2],
	exposureTime, exposureOffset, focalLength, fNumber,
	timestamp.GetNanos(), timestampDepth.GetNanos(),
	trackingState.GetName(), lightIntensity, lightColorTemp
);
		return meta;
	}

	std::string ArFrameMetadata::ToXML() const
	{
		std::stringstream meta;
		meta << std::setprecision(std::numeric_limits<double>::digits10);

		meta << "<arframe>\n\t<camera>\n\t\t<transform>";
		MatToXML(transformation, meta, "\n\t\t\t");
		meta << "\n\t\t</transform>\n\t\t<projection>";
		MatToXML(projection, meta, "\n\t\t\t");
		meta << "\n\t\t</projection>\n\t\t<resolution>\n\t\t\t<width>" << intrinsic.GetResolution().x << "</width>\n\t\t\t";
		meta << "<height>" << intrinsic.GetResolution().y << "</height>\n\t\t</resolution>\n\t\t<intrinsics>";
		MatToXML(intrinsic.GetMatrix(), meta, "\n\t\t\t");
		meta << "</intrinsics>\n\t\t<exposureDuration>" << exposureTime << "</exposureDuration>\n\t\t<exposureOffset>" << exposureOffset;
		meta << "</exposureOffset>\n\t</camera>\n\t";
		meta << "<timestamp>" << timestamp.GetSeconds() << "</timestamp>\n\t";
		meta << "<timestampDepth>" << timestampDepth.GetSeconds() << "</timestampDepth>\n\t";
		meta << "<trackingState>" << "\n\t\t<camera>" << trackingState.GetName() << "</camera>" << "\n\t</trackingState>";
		meta << "<light>\n\t\t<ambientIntensity>" << lightIntensity << "</ambientIntensity>\n\t\t";
		meta << "<ambientColorTemp>" << lightColorTemp << "</ambientColorTemp>\n\t</light>";
		meta << "\n</arframe>";

		return meta.str();
	}
}