// Copyright Epic Games, Inc. All Rights Reserved.

#include "view_cmd.h"

#include <zencore/compactbinary.h>
#include <zencore/compactbinaryutil.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/scopeguard.h>
#include <zenremotestore/builds/updatefoldereventdiskstream.h>

#include "../imgui_ui.h"
#include "imgui.h"

#include <memory>
#include <numeric>

namespace zen {

namespace {
	class ViewUI
	{
	public:
		virtual ~ViewUI() {}

		virtual std::string Title()										 = 0;
		virtual void		Read(const std::filesystem::path& InputPath) = 0;
		virtual void		Render(int display_w, int display_h)		 = 0;
		virtual uint64_t	GetEndTimeUs()								 = 0;
		virtual void		SetCurrentTimeUs(uint64_t CurrentUs)		 = 0;
	};

	class UpdateFolderEventDiskStreamUI : public ViewUI
	{
	public:
		virtual std::string Title() override { return "Update Folder"; }
		virtual void		Read(const std::filesystem::path& InputPath) override
		{
			{
				RwLock::ExclusiveLockScope _(m_Lock);
				m_InputReady = false;
				m_InputPath	 = InputPath;
			}
			UpdateFolderEventDiskStream::Read(InputPath,
											  m_RemoteContent,
											  m_BlockDescriptions,
											  m_DownloadChunks,
											  m_DownloadBlockRanges,
											  m_DownloadBlocks,
											  m_SequencesToWrite,
											  m_Events);

			if (!m_Events.empty())
			{
				std::sort(m_Events.begin(),
						  m_Events.end(),
						  [](const UpdateFolderEventDiskStream::Event& Lhs, const UpdateFolderEventDiskStream::Event& Rhs) {
							  const uint64_t TimeStampLshUs = Lhs.Header.TimestampUs();
							  const uint64_t TimeStampRshUs = Rhs.Header.TimestampUs();
							  if (TimeStampLshUs < TimeStampRshUs)
							  {
								  return true;
							  }
							  else if (TimeStampLshUs > TimeStampRshUs)
							  {
								  return false;
							  }
							  return uint8_t(Lhs.Header.EventType) < uint8_t(Rhs.Header.EventType);
						  });
			}

			ResetSequenceState();

			m_RealtimePlaybackStartUs.Reset();

			{
				RwLock::ExclusiveLockScope _(m_Lock);
				m_InputReady = true;
			}
		}

		void ResetSequenceState()
		{
			m_RenderSequenceStatues.clear();
			m_RenderSequenceStatues.reserve(m_RemoteContent.ChunkedContent.SequenceRawHashes.size());

			tsl::robin_map<IoHash, uint64_t> RawHashToRawSize;
			RawHashToRawSize.reserve(m_RemoteContent.Paths.size());
			for (uint32_t PathIndex = 0; PathIndex < m_RemoteContent.Paths.size(); PathIndex++)
			{
				RawHashToRawSize.insert({m_RemoteContent.RawHashes[PathIndex], m_RemoteContent.RawSizes[PathIndex]});
			}

			for (uint32_t SequenceIndex = 0; SequenceIndex < m_RemoteContent.ChunkedContent.SequenceRawHashes.size(); SequenceIndex++)
			{
				const uint64_t Size = RawHashToRawSize.at(m_RemoteContent.ChunkedContent.SequenceRawHashes[SequenceIndex]);
				m_RenderSequenceStatues.push_back(ChunkSequenceStatus{.Size = Size, .Completed = true});
			}
			for (const UpdateFolderEventStream::Sequence& SequenceToWrite : m_SequencesToWrite)
			{
				m_RenderSequenceStatues[SequenceToWrite.RemoteSequenceIndex].Completed = false;
			}
		}

		virtual void Render(int display_w, int display_h) override
		{
			ZEN_UNUSED(display_w, display_h);

			const ImGuiViewport* main_viewport = ImGui::GetMainViewport();
			ImGui::SetNextWindowPos(ImVec2(main_viewport->WorkPos.x, main_viewport->WorkPos.y), ImGuiCond_Always);
			ImGui::SetNextWindowSize(ImVec2(main_viewport->WorkSize.x, main_viewport->WorkSize.y), ImGuiCond_Always);

			// Main body of the Demo window starts here.
			if (!ImGui::Begin("Zen Visualizer", &VisualizerOpen, 0))
			{
				// Early out if the window is collapsed, as an optimization.
				ImGui::End();
				return;
			}

			RwLock::SharedLockScope Lock(m_Lock);
			if (m_InputReady)
			{
				Lock.ReleaseNow();

				uint64_t CurrentTimeUs = m_RealtimePlaybackStartUs.GetElapsedTimeUs();

				SetCurrentTimeUs(CurrentTimeUs);

				UpdateState();

				uint64_t EndTimeUs = GetEndTimeUs();

				float Progress =
					(EndTimeUs == 0) || (CurrentTimeUs >= EndTimeUs) ? 1.f : static_cast<float>(double(CurrentTimeUs) / double(EndTimeUs));
				ImGui::ProgressBar(Progress, ImVec2(-FLT_MIN, 0.f));
			}
			else if (std::filesystem::path InputPath = m_InputPath; !InputPath.empty())
			{
				Lock.ReleaseNow();

				ImGui::Text(fmt::format("Loading '{}'", InputPath.generic_string()).c_str());
			}
			else
			{
				Lock.ReleaseNow();
			}
			ImGui::End();
		}

		bool VisualizerOpen = false;

		virtual uint64_t GetEndTimeUs()
		{
			RwLock::SharedLockScope Lock(m_Lock);
			if (m_InputReady)
			{
				return m_Events.empty() ? 0 : m_Events.back().Header.TimestampUs() + 1;
			}
			else
			{
				return 0;
			}
		}

		virtual void SetCurrentTimeUs(uint64_t CurrentUs) { m_CurrentTimeUs.store(CurrentUs); }

	private:
		void UpdateState()
		{
			ZEN_ASSERT(m_InputReady);

			const uint64_t CurrentTimeUs = m_CurrentTimeUs.load();
			if (CurrentTimeUs != m_LastRenderTimeUs)
			{
				while (CurrentTimeUs < m_LastRenderTimeUs && m_LastRenderEventIndex > 0)
				{
					if (m_Events[m_LastRenderEventIndex - 1].Header.TimestampUs() >= CurrentTimeUs)
					{
						RemoveEventFromState(m_Events[m_LastRenderEventIndex]);
						m_LastRenderEventIndex--;
					}
					else
					{
						break;
					}
				}
				while (m_LastRenderTimeUs < CurrentTimeUs && m_LastRenderEventIndex < m_Events.size())
				{
					if (m_Events[m_LastRenderEventIndex].Header.TimestampUs() <= CurrentTimeUs)
					{
						AddEventToState(m_Events[m_LastRenderEventIndex]);
						m_LastRenderEventIndex++;
					}
					else
					{
						break;
					}
				}
				m_LastRenderTimeUs = CurrentTimeUs;
			}
		}

		void RemoveEventFromState(const UpdateFolderEventDiskStream::Event& Event)
		{
			switch (Event.Header.EventType)
			{
				case UpdateFolderEventDiskStream::Event::EventType::WroteSequence:
					{
						ChunkSequenceStatus& Sequence = m_RenderSequenceStatues[Event.Payload.WroteSequence.RemoteSequenceIndex];
						ZEN_ASSERT(!Sequence.Completed);
						Sequence.RemoveCompletedRange(
							ChunkSequenceStatus::Range{Event.Payload.WroteSequence.Offset, Event.Payload.WroteSequence.Length});
					}
					break;
				case UpdateFolderEventDiskStream::Event::EventType::CompletedSequence:
					{
						ChunkSequenceStatus& Sequence = m_RenderSequenceStatues[Event.Payload.CompletedSequence.RemoteSequenceIndex];
						ZEN_ASSERT(Sequence.Completed);
						Sequence.Completed = false;
					}
					break;
				case UpdateFolderEventDiskStream::Event::EventType::WriteComplete:
					break;
				default:
					// TODO
					break;
			}
		}

		void AddEventToState(const UpdateFolderEventDiskStream::Event& Event)
		{
			switch (Event.Header.EventType)
			{
				case UpdateFolderEventDiskStream::Event::EventType::WroteSequence:
					{
						ChunkSequenceStatus& Sequence = m_RenderSequenceStatues[Event.Payload.WroteSequence.RemoteSequenceIndex];
						ZEN_ASSERT(!Sequence.Completed);
						Sequence.AddCompletedRange(
							ChunkSequenceStatus::Range{Event.Payload.WroteSequence.Offset, Event.Payload.WroteSequence.Length});
					}
					break;
				case UpdateFolderEventDiskStream::Event::EventType::CompletedSequence:
					{
						ChunkSequenceStatus& Sequence = m_RenderSequenceStatues[Event.Payload.CompletedSequence.RemoteSequenceIndex];
						ZEN_ASSERT(!Sequence.Completed);
						uint64_t CompletedSize = Sequence.CompletedSize();
						if (Sequence.CompletedRanges.size() != 1)
						{
							ZEN_ASSERT(CompletedSize < Sequence.Size);
							ZEN_WARN("Missing {} bytes from sequence {}",
									 Sequence.Size - CompletedSize,
									 Event.Payload.CompletedSequence.RemoteSequenceIndex);
						}
						else
						{
							if (Sequence.CompletedRanges.front().Offset != 0)
							{
								ZEN_WARN("Missing {} bytes from sequence {} (chunk is not from start)",
										 Sequence.Size - CompletedSize,
										 Event.Payload.CompletedSequence.RemoteSequenceIndex);
							}
							if (Sequence.CompletedRanges.front().Length != Sequence.Size)
							{
								ZEN_WARN("Missing {} bytes from sequence {} (chunk is not to end)",
										 Sequence.Size - CompletedSize,
										 Event.Payload.CompletedSequence.RemoteSequenceIndex);
							}
						}
						Sequence.Completed = true;
					}
					break;
				case UpdateFolderEventDiskStream::Event::EventType::WriteComplete:
					break;
				default:
					// TODO
					break;
			}
		}

		RwLock				  m_Lock;
		std::filesystem::path m_InputPath;
		bool				  m_InputReady			 = false;
		std::atomic<uint64_t> m_CurrentTimeUs		 = 0;
		uint64_t			  m_LastRenderTimeUs	 = 0;
		uint64_t			  m_LastRenderEventIndex = 0;

		Stopwatch m_RealtimePlaybackStartUs;

		ChunkedFolderContent							 m_RemoteContent;
		std::vector<ChunkBlockDescription>				 m_BlockDescriptions;
		std::vector<UpdateFolderEventStream::Chunk>		 m_DownloadChunks;
		std::vector<UpdateFolderEventStream::BlockRange> m_DownloadBlockRanges;
		std::vector<UpdateFolderEventStream::Block>		 m_DownloadBlocks;
		std::vector<UpdateFolderEventStream::Sequence>	 m_SequencesToWrite;
		std::vector<UpdateFolderEventDiskStream::Event>	 m_Events;

		std::vector<ChunkSequenceStatus> m_RenderSequenceStatues;
	};
}  // namespace

ViewCommand::ViewCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "", "input-stream", "Path to a recorded stream file", cxxopts::value(m_InputStream), "<path>");

	m_Options.parse_positional({"input-stream"});
	m_Options.positional_help("input-stream");
}

ViewCommand::~ViewCommand() = default;

void
ViewCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	using namespace std::literals;

	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return;
	}

	//	SequenceStatus TestStatus{ .Size = 722580630, .NormalizedSize = 1000, .Completed = false };
	//	TestStatus.AddCompletedRange(SequenceStatus::Range(262144, 241664));
	//	TestStatus.AddCompletedRange(SequenceStatus::Range(262144 + 241664, 1360));
	//	TestStatus.AddCompletedRange(SequenceStatus::Range(0, 262144));
	//	TestStatus.AddCompletedRange(SequenceStatus::Range(262144 + 241664 + 1360, 722580630 - (262144 + 241664 + 1360)));
	//
	//	if (true)
	//	{
	//		return;
	//	}

	std::unique_ptr<ViewUI> UI;

	std::filesystem::path AbsInputStreamPath = MakeSafeAbsolutePath(m_InputStream);
	if (ToLower(AbsInputStreamPath.extension().string()) == ".zrb")
	{
		UI = std::make_unique<UpdateFolderEventDiskStreamUI>();
	}

	if (!UI)
	{
		throw OptionParseException(fmt::format("Unknown input format for '{}'", m_InputStream), m_Options.help());
	}

	std::atomic<bool> CloseUIFlag = false;

	Event UiClosed;

	std::thread UIThread = OpenImGuiWindow(
		1280,
		720,
		UI->Title(),
		[&UI](int display_w, int display_h) { UI->Render(display_w, display_h); },
		[&UiClosed]() { UiClosed.Set(); },
		[&CloseUIFlag]() { return CloseUIFlag.load(); });

	auto CloseUI = MakeGuard([&]() {
		UiClosed.Wait();
		UIThread.join();
	});

	try
	{
		UI->Read(m_InputStream);
	}
	catch (const std::exception&)
	{
		CloseUIFlag.store(true);
		throw;
	}
}

}  // namespace zen