#include "NdiInput.h"

#include "NdiInputFormat.h"
#include "NdiRuntime.h"
#include "logging/Logger.h"
#include "video/core/VideoIOFormat.h"

#include <Processing.NDI.Lib.h>

#include <chrono>
#include <cstring>
#include <mutex>

namespace RenderCadenceCompositor
{
namespace
{
constexpr std::chrono::milliseconds kDiscoveryRetry(500);
constexpr uint32_t kCaptureTimeoutMilliseconds = 250;

bool IsDefaultDeviceName(const std::string& device)
{
	return device.empty() || device == "default" || device == "auto";
}

}

NdiInput::NdiInput(InputFrameMailbox& mailbox) :
	mMailbox(mailbox)
{
}

NdiInput::~NdiInput()
{
	ReleaseResources();
}

bool NdiInput::Initialize(const NdiInputConfig& config, std::string& error)
{
	ReleaseResources();
	mConfig = config;
	if (!AcquireNdiRuntime())
	{
		error = "NDI runtime initialization failed. The CPU/runtime might not support NDI.";
		return false;
	}

	mInitialized.store(true, std::memory_order_release);
	Log("ndi-input", "NDI input edge initialized for device '" + (IsDefaultDeviceName(config.device) ? std::string("default") : config.device) + "'.");
	return true;
}

bool NdiInput::Start(std::string& error)
{
	if (!mInitialized.load(std::memory_order_acquire))
	{
		error = "NDI input has not been initialized.";
		return false;
	}
	if (mThread.joinable())
		return true;

	mStopping.store(false, std::memory_order_release);
	mRunning.store(true, std::memory_order_release);
	mThread = std::thread([this]() { ThreadMain(); });
	Log("ndi-input", "NDI input capture thread started.");
	return true;
}

void NdiInput::Stop()
{
	mStopping.store(true, std::memory_order_release);
	if (mThread.joinable())
		mThread.join();
	mRunning.store(false, std::memory_order_release);
	DisconnectReceiver();
}

void NdiInput::ReleaseResources()
{
	Stop();
	if (mInitialized.exchange(false, std::memory_order_acq_rel))
		ReleaseNdiRuntime();
}

NdiInputMetrics NdiInput::Metrics() const
{
	NdiInputMetrics metrics;
	metrics.capturedFrames = mCapturedFrames.load(std::memory_order_relaxed);
	metrics.noInputSourceFrames = mNoInputSourceFrames.load(std::memory_order_relaxed);
	metrics.unsupportedFrames = mUnsupportedFrames.load(std::memory_order_relaxed);
	metrics.submitMisses = mSubmitMisses.load(std::memory_order_relaxed);
	metrics.convertMilliseconds = 0.0;
	metrics.submitMilliseconds = mSubmitMilliseconds.load(std::memory_order_relaxed);
	metrics.captureFormat = CaptureFormatName();
	return metrics;
}

void NdiInput::ThreadMain()
{
	while (!mStopping.load(std::memory_order_acquire))
	{
		if (mReceiver == nullptr && !ConnectToConfiguredSource())
		{
			mNoInputSourceFrames.fetch_add(1, std::memory_order_relaxed);
			bool expected = false;
			if (mLoggedNoSource.compare_exchange_strong(expected, true, std::memory_order_relaxed))
				TryLog(LogLevel::Warning, "ndi-input", "No NDI input source available yet; render input will use fallback until a source appears.");
			std::this_thread::sleep_for(kDiscoveryRetry);
			continue;
		}

		NDIlib_video_frame_v2_t videoFrame;
		const NDIlib_frame_type_e frameType = NDIlib_recv_capture_v3(
			static_cast<NDIlib_recv_instance_t>(mReceiver),
			&videoFrame,
			nullptr,
			nullptr,
			kCaptureTimeoutMilliseconds);

		if (frameType == NDIlib_frame_type_video)
		{
			HandleVideoFrame(mReceiver, &videoFrame);
			NDIlib_recv_free_video_v2(static_cast<NDIlib_recv_instance_t>(mReceiver), &videoFrame);
		}
		else if (frameType == NDIlib_frame_type_error)
		{
			TryLog(LogLevel::Warning, "ndi-input", "NDI input receiver reported an error; reconnecting.");
			DisconnectReceiver();
		}
		else if (frameType == NDIlib_frame_type_source_change)
		{
			TryLog(LogLevel::Log, "ndi-input", "NDI input source changed.");
		}
	}
}

bool NdiInput::ConnectToConfiguredSource()
{
	return IsDefaultDeviceName(mConfig.device) ? ConnectToDefaultSource() : [&]() {
		NDIlib_source_t source;
		std::memset(&source, 0, sizeof(source));
		source.p_ndi_name = mConfig.device.c_str();

		NDIlib_recv_create_v3_t createSettings;
		std::memset(&createSettings, 0, sizeof(createSettings));
		createSettings.source_to_connect_to = source;
		createSettings.color_format = NDIlib_recv_color_format_BGRX_BGRA;
		createSettings.bandwidth = NDIlib_recv_bandwidth_highest;
		createSettings.allow_video_fields = false;
		createSettings.p_ndi_recv_name = "Render Cadence NDI Input";

		std::lock_guard<std::mutex> lock(mReceiverMutex);
		mReceiver = NDIlib_recv_create_v3(&createSettings);
		if (mReceiver != nullptr)
			Log("ndi-input", "NDI input connected to configured source '" + mConfig.device + "'.");
		return mReceiver != nullptr;
	}();
}

bool NdiInput::ConnectToDefaultSource()
{
	NDIlib_find_instance_t finder = NDIlib_find_create_v2();
	if (finder == nullptr)
		return false;

	uint32_t sourceCount = 0;
	const NDIlib_source_t* sources = NDIlib_find_get_current_sources(finder, &sourceCount);
	if (sourceCount == 0)
	{
		NDIlib_find_wait_for_sources(finder, 100);
		sources = NDIlib_find_get_current_sources(finder, &sourceCount);
	}

	if (sourceCount == 0 || sources == nullptr)
	{
		NDIlib_find_destroy(finder);
		return false;
	}

	NDIlib_recv_create_v3_t createSettings;
	std::memset(&createSettings, 0, sizeof(createSettings));
	createSettings.source_to_connect_to = sources[0];
	createSettings.color_format = NDIlib_recv_color_format_BGRX_BGRA;
	createSettings.bandwidth = NDIlib_recv_bandwidth_highest;
	createSettings.allow_video_fields = false;
	createSettings.p_ndi_recv_name = "Render Cadence NDI Input";

	const std::string sourceName = sources[0].p_ndi_name != nullptr ? sources[0].p_ndi_name : "unnamed";
	std::lock_guard<std::mutex> lock(mReceiverMutex);
	mReceiver = NDIlib_recv_create_v3(&createSettings);
	NDIlib_find_destroy(finder);
	if (mReceiver != nullptr)
		Log("ndi-input", "NDI input connected to default source '" + sourceName + "'.");
	return mReceiver != nullptr;
}

void NdiInput::DisconnectReceiver()
{
	std::lock_guard<std::mutex> lock(mReceiverMutex);
	if (mReceiver != nullptr)
	{
		NDIlib_recv_destroy(static_cast<NDIlib_recv_instance_t>(mReceiver));
		mReceiver = nullptr;
	}
}

void NdiInput::HandleVideoFrame(void*, const void* frame)
{
	const NDIlib_video_frame_v2_t& videoFrame = *static_cast<const NDIlib_video_frame_v2_t*>(frame);
	if (videoFrame.p_data == nullptr || videoFrame.xres <= 0 || videoFrame.yres <= 0 || videoFrame.line_stride_in_bytes <= 0)
	{
		mUnsupportedFrames.fetch_add(1, std::memory_order_relaxed);
		return;
	}

	VideoIOPixelFormat pixelFormat = VideoIOPixelFormat::Bgra8;
	if (!MapNdiFourCcToVideoIOPixelFormat(videoFrame.FourCC, pixelFormat))
	{
		mUnsupportedFrames.fetch_add(1, std::memory_order_relaxed);
		bool expected = false;
		if (mLoggedUnsupportedFrame.compare_exchange_strong(expected, true, std::memory_order_relaxed))
			TryLog(LogLevel::Warning, "ndi-input", std::string("Unsupported NDI input frame format ") + NdiFourCcName(videoFrame.FourCC) + ".");
		return;
	}

	const uint64_t frameIndex = mCapturedFrames.load(std::memory_order_relaxed);
	if (SubmitFrame(
		videoFrame.p_data,
		static_cast<unsigned>(videoFrame.xres),
		static_cast<unsigned>(videoFrame.yres),
		static_cast<unsigned>(videoFrame.line_stride_in_bytes),
		pixelFormat,
		frameIndex))
	{
		mCapturedFrames.fetch_add(1, std::memory_order_relaxed);
		LogConfiguredShapeMismatchOnce(static_cast<unsigned>(videoFrame.xres), static_cast<unsigned>(videoFrame.yres));
		LogFirstFrameOnce(static_cast<unsigned>(videoFrame.xres), static_cast<unsigned>(videoFrame.yres), pixelFormat);
	}
}

bool NdiInput::SubmitFrame(const void* bytes, unsigned width, unsigned height, unsigned rowBytes, VideoIOPixelFormat pixelFormat, uint64_t frameIndex)
{
	ConfigureMailboxIfNeeded(width, height, rowBytes, pixelFormat);
	const auto submitStart = std::chrono::steady_clock::now();
	const bool submitted = mMailbox.SubmitFrame(bytes, rowBytes, frameIndex);
	mSubmitMilliseconds.store(
		std::chrono::duration_cast<std::chrono::duration<double, std::milli>>(std::chrono::steady_clock::now() - submitStart).count(),
		std::memory_order_relaxed);
	if (!submitted)
	{
		mSubmitMisses.fetch_add(1, std::memory_order_relaxed);
		bool expected = false;
		if (mLoggedSubmitMiss.compare_exchange_strong(expected, true, std::memory_order_relaxed))
			TryLog(LogLevel::Warning, "ndi-input", "NDI input frame could not be submitted to InputFrameMailbox.");
	}
	return submitted;
}

void NdiInput::ConfigureMailboxIfNeeded(unsigned width, unsigned height, unsigned rowBytes, VideoIOPixelFormat pixelFormat)
{
	InputFrameMailboxConfig current = mMailbox.Config();
	if (current.width == width && current.height == height && current.rowBytes == rowBytes && current.pixelFormat == pixelFormat)
		return;

	current.width = width;
	current.height = height;
	current.rowBytes = rowBytes;
	current.pixelFormat = pixelFormat;
	current.capacity = mConfig.mailboxCapacity;
	current.maxReadyFrames = mConfig.maxReadyFrames;
	mMailbox.Configure(current);
	mCapturePixelFormat.store(pixelFormat, std::memory_order_release);
	TryLog(
		LogLevel::Log,
		"ndi-input",
		"NDI input mailbox adapted to " + std::to_string(width) + "x" + std::to_string(height) +
			" " + VideoIOPixelFormatName(pixelFormat) + " rowBytes=" + std::to_string(rowBytes) + ".");
}

void NdiInput::LogConfiguredShapeMismatchOnce(unsigned width, unsigned height)
{
	if (mConfig.expectedWidth == 0 || mConfig.expectedHeight == 0 ||
		(width == mConfig.expectedWidth && height == mConfig.expectedHeight))
		return;

	bool expected = false;
	if (mLoggedShapeMismatch.compare_exchange_strong(expected, true, std::memory_order_relaxed))
	{
		TryLog(
			LogLevel::Warning,
			"ndi-input",
			"NDI input source is " + std::to_string(width) + "x" + std::to_string(height) +
				", which differs from configured input " + std::to_string(mConfig.expectedWidth) + "x" +
				std::to_string(mConfig.expectedHeight) + ". Adapting to the source shape.");
	}
}

void NdiInput::LogFirstFrameOnce(unsigned width, unsigned height, VideoIOPixelFormat pixelFormat)
{
	bool expected = false;
	if (mLoggedFirstFrame.compare_exchange_strong(expected, true, std::memory_order_relaxed))
	{
		TryLog(
			LogLevel::Log,
			"ndi-input",
			"First NDI input frame submitted to InputFrameMailbox: " + std::to_string(width) + "x" +
				std::to_string(height) + " " + VideoIOPixelFormatName(pixelFormat) + ".");
	}
}

const char* NdiInput::CaptureFormatName() const
{
	return VideoIOPixelFormatName(mCapturePixelFormat.load(std::memory_order_acquire));
}
}