/*
 * 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 "ArRecorder.hpp"
#include "ArSession.hpp"
#include "ArFrame.hpp"
#include "IO/Archive/MultiPartArchiveWriter.hpp"
#include "IO/Files/Pfm.hpp"
#include "IO/Files/Pnm.hpp"
#include "IO/AppFolders.hpp"
#include "Base/BlockProfiler.hpp"
#include "Base/Logger.hpp"
#include "Image/YuvUtils.hpp"
#include <fstream>
#include <iostream>
//#if __has_include("turbojpeg.h")
#include <turbojpeg.h>
#define TURBO_JPEG
//#endif

namespace openVulkanoCpp::AR
{
	namespace
	{
		std::filesystem::path GeneratePath(const std::filesystem::path& baseDir, std::string_view name)
		{
			std::stringstream ss;
			auto t = std::time(nullptr);
			auto localTime = *std::localtime(&t);
			ss << std::put_time(&localTime, "%Y-%m-%d_%H_%M_%S");
			return baseDir / name / ss.str();
		}

		std::string GetFileName(size_t frameId, std::string_view fileExtension)
		{
			std::stringstream fnStream;
			fnStream << std::setw(7) << std::setfill('0') << frameId << '.' << fileExtension;
			return fnStream.str();
		}
	}

	ArRecorder::ArRecorder(ArSession* session)
			: m_session(session), m_asyncProcessor(this)
	{
		m_settings.path = GeneratePath(AppFolders::GetAppDataHomeDir(), "ar_recording");
		session->OnNewFrameHighResolution += EventHandler(this, &ArRecorder::SaveHighResolution);
	}

	ArRecorder::~ArRecorder() = default;

	void ArRecorder::WriteColorImage(ArFrame* arFrame, MultiPartArchiveWriter* colorWriter, bool highRes) const
	{
		BlockProfiler profile("Save AR Frame - Image");
		std::string fileName = GetFileName(arFrame->GetFrameId(), "jpg");
#ifndef TURBO_JPEG
		if (arFrame->GetCameraImageAsJpeg([&fileName, this](const char* data, size_t len){ m_colorWriter->AddFile(fileName.c_str(), data, len); }))
			return;
		//TODO stb???
		Logger::AR->error("Failed to create JPEG! Missing turbojpeg.h");
#else
		//TODO handle non nv12 images
		auto img = arFrame->GetCameraImage();
		auto sizeLum = img.luminescenceOrColor.resolution.x * img.luminescenceOrColor.resolution.y;
		auto sizeUV = img.uv.resolution.x * img.uv.resolution.y;
		auto resX = img.luminescenceOrColor.resolution.x;
		auto resY = img.luminescenceOrColor.resolution.y;

		tjhandle handle = tjInitCompress();
		const uint8_t* buffers[3];
		std::unique_ptr<uint8_t[]> dataBuffer;
		if (m_settings.downsampleColor && !highRes)
		{
			dataBuffer = YuvUtils::PlansFromNV12(static_cast<uint8_t*>(img.luminescenceOrColor.data), static_cast<uint8_t*>(img.uv.data),
												 resX, resY, img.uv.resolution.x, img.uv.resolution.y, 2, 2);
			resX /= 2;
			resY /= 2;
			buffers[0] = dataBuffer.get();
			buffers[1] = buffers[0] + sizeLum / 4;
			buffers[2] = buffers[1] + sizeUV / 4;
		}
		else
		{
			dataBuffer = std::unique_ptr<uint8_t[]>(new uint8_t[sizeUV + sizeUV]);
			YuvUtils::ChromaPlanesFromNV12((uint8_t*)img.uv.data, dataBuffer.get(), sizeUV);
			buffers[0] = static_cast<uint8_t*>(img.luminescenceOrColor.data);
			buffers[1] = dataBuffer.get();
			buffers[2] = buffers[1] + sizeUV;
		}

		uint8_t* outBuffer = nullptr;
		unsigned long size = 0;
		if (tjCompressFromYUVPlanes(handle, buffers, resX, nullptr, resY, TJSAMP_420, &outBuffer, &size, 95, TJFLAG_FASTDCT))
			Logger::AR->error("Failed to create JPEG! {}", tjGetErrorStr());
		else
			colorWriter->AddFile(fileName.c_str(), outBuffer, size);
		tjFree(outBuffer);
#endif
	}

	void ArRecorder::WriteDepthImage(ArFrame* arFrame, MultiPartArchiveWriter* depthWriter, MultiPartArchiveWriter* confWriter)
	{
		//BlockProfiler profile("Save AR Frame - Depth");
		if (!m_depthWriter || !m_confidenceWriter) return;
		auto depthImg = arFrame->GetDepthImage();
		std::vector<std::pair<const void*, size_t>> buffers(2);
		{ // TODO handle alternative depth formats!!!!
			BlockProfiler profile("Save AR Frame - Depth");
			PfmHeader depthHeader(static_cast<uint32_t>(depthImg.depth.resolution.x), static_cast<uint32_t>(depthImg.depth.resolution.y), 5.0f, false);
			std::string header = depthHeader.ToString();
			buffers[0].first = header.c_str();
			buffers[0].second = header.size();
			buffers[1].first = static_cast<const char *>(depthImg.depth.data);
			buffers[1].second = depthImg.depth.resolution.x * depthImg.depth.resolution.y * sizeof(float);

			const std::string fileName = GetFileName(arFrame->GetFrameId(), "pfm");
			depthWriter->AddFile(fileName.c_str(), buffers);
		}

		{
			BlockProfiler profile("Save AR Frame - Confi");
			PnmHeader confidenceHeader(static_cast<uint32_t>(depthImg.confidence.resolution.x), static_cast<uint32_t>(depthImg.confidence.resolution.y), false, 2);
			std::string header = confidenceHeader.ToString();
			buffers[0].first = header.c_str();
			buffers[0].second = header.size();
			buffers[1].first = static_cast<const char *>(static_cast<void*>(depthImg.confidence.data));
			buffers[1].second = static_cast<size_t>(depthImg.confidence.resolution.x * depthImg.confidence.resolution.y);

			const std::string fileName = GetFileName(arFrame->GetFrameId(), "pgm");
			confWriter->AddFile(fileName.c_str(), buffers);
		}
	}

	void ArRecorder::WriteMetadata(ArFrame* frame, MultiPartArchiveWriter* metaWriter)
	{
		BlockProfiler profileMeta("Save AR Frame -  Meta");
		std::string metaContent = frame->GetFrameMetadata().ToXML();
		std::string fileName = GetFileName(frame->GetFrameId(), "meta");
		metaWriter->AddFile(fileName.c_str(), metaContent.c_str(), metaContent.size());
	}

	void ArRecorder::Write(ArFrame* frame, bool highRes)
	{
		if (frame->IsSaved()) return;
		frame->SetSaved();
		bool useHighResWriter = highRes && m_settings.highResFramesInSeparateArchive;
		BlockProfiler profile("Save AR Frame");
		WriteMetadata(frame, useHighResWriter ? m_highResWriter.get() : m_metadataWriter.get());
		WriteColorImage(frame, useHighResWriter ? m_highResWriter.get() : m_colorWriter.get(), highRes);
		WriteDepthImage(frame,
						useHighResWriter ? m_highResWriter.get() : m_depthWriter.get(),
						useHighResWriter ? m_highResWriter.get() : m_confidenceWriter.get());
	}

	void ArRecorder::Start()
	{
		if (!m_colorWriter)
		{
			m_colorWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "color_{:05d}.tar", ArchiveConfig::TAR, m_settings.archiveSize, true);
			m_depthWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "depth_{:05d}.tar", ArchiveConfig::TAR, m_settings.archiveSize, true);
			m_confidenceWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "confidence_{:05d}.tar.gz", ArchiveConfig::TAR_GZ, m_settings.archiveSize, true);
			m_metadataWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "meta_{:05d}.tar.gz", ArchiveConfig::TAR_GZ, m_settings.archiveSize, true);
			m_highResWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "highres_{:05d}.tar", ArchiveConfig::TAR, m_settings.archiveSize, true);

			std::ofstream platformInfoStream(m_settings.path / "ArRecording.xml");
			platformInfoStream << m_session->GetSessionMetadata().ToXML();
			platformInfoStream.close();
		}
		m_recording = true;
	}

	void ArRecorder::Stop()
	{
		if (!m_recording) return;
		m_recording = false;
		for(MultiPartArchiveWriter* writer : { m_colorWriter.get(), m_depthWriter.get(), m_confidenceWriter.get(), m_metadataWriter.get(), m_highResWriter.get() })
		{
			writer->Split();
		}
	}

	void ArRecorder::SetRecordingPath(const std::string& path)
	{
		if (!m_colorWriter)
		{
			for (MultiPartArchiveWriter* writer: { m_colorWriter.get(), m_depthWriter.get(), m_confidenceWriter.get(), m_metadataWriter.get(), m_highResWriter.get() })
			{
				writer->Move(path);
			}
			std::filesystem::rename(m_settings.path / "ArRecording.xml", path + "/ArRecording.xml");
		}

		m_persistent = true;
		m_settings.path = path;
	}

	void ArRecorder::SetRecordingMode(RecordingMode mode)
	{
		if (m_settings.recordingMode == mode) return;
		if (m_settings.recordingMode == RecordingMode::NEW_FRAME && m_newFrameHandler)
		{
			m_session->OnNewFrame -= m_newFrameHandler;
			m_newFrameHandler = nullptr;
		}
		m_settings.recordingMode = mode;
		if (m_settings.recordingMode == RecordingMode::NEW_FRAME)
		{
			m_newFrameHandler = m_session->OnNewFrame += EventHandler(this, &ArRecorder::Save);
		}
	}

	void ArRecorder::Save(const std::shared_ptr<ArFrame>& frame)
	{
		if (!m_recording) return;
		if (m_settings.asyncRecording)
		{
			m_asyncProcessor.Queue(frame, false);
		}
		else Write(frame.get());
	}

	void ArRecorder::SaveHighResolution(const std::shared_ptr<ArFrame>& frame)
	{
		if (!m_recording) return;
		if (m_settings.asyncRecording)
		{
			m_asyncProcessor.Queue(frame, true);
		}
		else Write(frame.get(), true);
	}

	//region AsyncProcessor
	ArRecorder::AsyncProcessor::AsyncProcessor(ArRecorder* recorder)
		: recorder(recorder), processingThread(&ArRecorder::AsyncProcessor::Handler, this)
	{}

	ArRecorder::AsyncProcessor::~AsyncProcessor()
	{
		requestExit = true;
		newDataAvailable.notify_one();
		if (processingThread.joinable()) processingThread.join();
	}

	void ArRecorder::AsyncProcessor::Queue(const std::shared_ptr<ArFrame>& frame, bool highRes)
	{
		if (requestExit) return; // no need to queue up on shutdown
		{
			std::unique_lock lock(queueMutex);
			if (highRes) highResFrameQueue.push(frame);
			else frameQueue.push(frame);
		}
		newDataAvailable.notify_all();
	}

	void ArRecorder::AsyncProcessor::Handler()
	{
		Utils::SetThreadName("ArRecorder");
		std::unique_lock lock(queueMutex);
		do
		{
			newDataAvailable.wait(lock, [this]{ return !frameQueue.empty() || !highResFrameQueue.empty() || requestExit; });
			while(!highResFrameQueue.empty())
			{
				auto frame = std::move(highResFrameQueue.front());
				highResFrameQueue.pop();
				if (frame->IsSaved()) continue;
				lock.unlock();
				recorder->Write(frame.get(), true);
				lock.lock();
			}
			if (requestExit) break;
			while(!frameQueue.empty())
			{
				if (frameQueue.size() > 3)
				{
					Logger::AR->warn("Falling behind saving frames, skipping ...");
					//while(frameQueue.size() > 3) frameQueue.pop();
				}
				auto frame = std::move(frameQueue.front());
				frameQueue.pop();
				if (frame->IsSaved()) continue;
				lock.unlock();
				recorder->Write(frame.get(), false);
				lock.lock();
			}
		}
		while (!requestExit);
	}
	//endregion
}