/*
 * 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/.
 */

#pragma once

#include "Math/ByteSize.hpp"
#include "Base/Event.hpp"
#include "Base/Timer.hpp"
#include "Base/Wrapper.hpp"
#include <atomic>
#include <condition_variable>
#include <filesystem>
#include <mutex>
#include <memory>
#include <string>
#include <thread>
#include <queue>

namespace OpenVulkano
{
	class IEventHandler;
	class IArchiveWriter;
	class MultiPartArchiveWriter;
	class JpegWithTagsWriter;
}

namespace OpenVulkano::AR
{
	class ArSession;
	class ArFrame;

	enum class RecordingMode
	{
		/**
		 * Frames will only be saved when manually requested through ArFrame::Save()
		 */
		MANUAL,
		/**
		 * Frames will only be saved when the frame is manually requested through ArSession::GetFrame()
		 */
		FRAME_REQUEST,
		/**
		 * Every new frame will be stored automatically
		 */
		NEW_FRAME
	};

	struct RecordingSettings
	{
		RecordingMode recordingMode = RecordingMode::MANUAL;
		std::filesystem::path path;
		size_t archiveSize = 2_GiB;
		bool downsampleColor = true;
		bool highResFramesInSeparateArchive = true;
		bool saveHighResFrames = true;
		bool asyncRecording = true;
	};

	class ArRecorder final
	{
		struct SaveToFileRequest final
		{
			Ptr<ArFrame> frame;
			std::filesystem::path path;
			bool downsample, addAux;
		};
		
		struct AsyncProcessor final
		{
			ArRecorder* recorder;
			std::mutex queueMutex;
			std::thread processingThread;
			std::queue<Ptr<ArFrame>> frameQueue, highResFrameQueue;
			std::queue<SaveToFileRequest> toFile;
			std::condition_variable newDataAvailable;
			std::atomic_bool requestExit{};

			explicit AsyncProcessor(ArRecorder* recorder);
			~AsyncProcessor();
			void Close();
			void Queue(const Ptr<ArFrame>& frame, bool highRes);
			void Queue(const Ptr<ArFrame>& frame, const std::filesystem::path& path, bool downsample, bool aux);
			void Handler();
			bool Empty() { return frameQueue.empty() && highResFrameQueue.empty() && toFile.empty(); }
		};

		ArSession* m_session;
		std::unique_ptr<MultiPartArchiveWriter> m_colorWriter, m_depthWriter, m_confidenceWriter, m_metadataWriter, m_highResWriter;
		RecordingSettings m_settings;
		uint64_t m_frameCount = 0;
		uint32_t m_skippedFrames = 0;
		bool m_recording = false, m_persistent = false;
		Timer m_timer;

		IEventHandler* m_newFrameHandler = nullptr;

		AsyncProcessor m_asyncProcessor;

		void Write(ArFrame* frame, bool highRes = false);
		void WriteMetadata(ArFrame* frame, IArchiveWriter* metaWriter);
		void WriteMetadata(ArFrame* frame, IArchiveWriter* metaWriter, const char* name);
		void WriteColorImage(ArFrame* arFrame, IArchiveWriter* colorWriter, JpegWithTagsWriter* jpgWriter, bool highRes) const;
		void WriteDepthImage(ArFrame *arFrame, IArchiveWriter* depthWriter, IArchiveWriter* confWriter);
		void WriteDepthImage(ArFrame *arFrame, IArchiveWriter* depthWriter, IArchiveWriter* confWriter, const char* depthName, const char* confidenceName);
		void WriteToFile(const std::shared_ptr<ArFrame>& frame, const std::filesystem::path& path, bool downsample, bool saveAux);

		void SplitWriters();
		void WriteMetadataFile();

	public:
		ArRecorder(ArSession* session);

		~ArRecorder();

		void Save(const std::shared_ptr<ArFrame>& frame);

		void SaveHighResolution(const std::shared_ptr<ArFrame>& frame);

		void SaveToFile(const std::shared_ptr<ArFrame>& frame, const std::filesystem::path& path, bool downsample = false, bool saveAux = true);

		/**
		 * Starts the recording of the owning AR session
		 */
		void Start();

		/**
		 * Stops the recording of the owning AR session
		 */
		void Stop();

		/**
		 * Sets the directory into which the AR recording should be stored.
		 * The path needs to be set to make the recording persistent.
		 * If path is changed after starting the recording, the already running recording will be moved to the new path.
		 * @param path The path to be used to store the recording
		 */
		[[deprecated]] void SetRecordingPath(const std::string& path)
		{
			std::filesystem::path p(path);
			SetRecordingPath(p);
		}

		/**
		 * Sets the directory into which the AR recording should be stored.
		 * The path needs to be set to make the recording persistent.
		 * If path is changed after starting the recording, the already running recording will be moved to the new path.
		 * @param path The path to be used to store the recording
		 */
		void SetRecordingPath(const std::filesystem::path& path);

		/**
		 * Gets the current recording path
		 * @return Current recording dir
		 */
		[[nodiscard]] const std::filesystem::path& GetRecordingPath() const { return m_settings.path; }

		/**
		 * Checks if a path to be used for the recording has been set.
		 * If no path has been set the recording will be stored in a temporary location and will be deleted once the recording is stopped.
		 * @return True if a path has been set and the recording will persist after stopping the recording.
		 */
		[[nodiscard]] bool IsPersistent() const { return m_persistent; }

		/**
		 * Checks if the recording of the AR session is running
		 * @return True if recording is started
		 */
		[[nodiscard]] bool IsRecording() const { return m_recording; }

		/**
		 * Sets the used recording mode
		 * @param mode The mode that should be used to record the current session
		 */
		void SetRecordingMode(RecordingMode mode);

		/**
		 * Checks the currently used recording mode for the AR session
		 * @return The currently used recording mode
		 */
		[[nodiscard]] RecordingMode GetRecordingMode() const { return m_settings.recordingMode; }

		/**
		 * If enabled color images will be reduced to 1/4 of the original resolution (both axes will have half the resolution).
		 * This option can be used to reduce the used storage space and storage bandwidth during recording.
		 *
		 * @param downsample true = downsample images; false = no downsampling
		 */
		void SetDownsampleColorImages(bool downsample = true) { m_settings.downsampleColor = downsample; }

		/**
		 * Checks if color image downsampling is enabled.
		 * @return true = downsample images; false = no downsampling
		 */
		bool GetDownsampleColorImages() const { return m_settings.downsampleColor; }

		/**
		 * If enabled the requested high resolution images will be stored in a separate archive ("highres.tar") instead of the normal ar recording archives.
		 * @param separateArchive true = using the separate archive file; false = using the normal archive files
		 */
		void SetHighResFramesInSeparateArchive(bool separateArchive = true) { m_settings.highResFramesInSeparateArchive = separateArchive; }

		/**
		 * Checks if a separate archive is used for the high res frames.
		 * @return true = using the separate archive file; false = using the normal archive files
		 */
		bool GetHighResFramesInSeparateArchive() const { return m_settings.highResFramesInSeparateArchive; }

		void SetArchivePartMaxFileSize(size_t maxPartSize = 2_GiB) { m_settings.archiveSize = maxPartSize; }

		size_t GetArchivePartMaxFileSize() const { return m_settings.archiveSize; }

		const RecordingSettings& GetRecordingSettings() const { return m_settings; }

		void SetRecordHighResImages(bool recHighRes = true) { m_settings.saveHighResFrames = recHighRes; }

		Event<ArRecorder*, bool> OnRecordingStateChanged;
	};
}