#include "DeckLinkInput.h"

#include "DeckLinkVideoIOFormat.h"
#include "../logging/Logger.h"

#include <algorithm>
#include <chrono>
#include <new>

namespace RenderCadenceCompositor
{
namespace
{
bool FindInputDisplayMode(IDeckLinkInput* input, BMDDisplayMode targetMode, IDeckLinkDisplayMode** foundMode)
{
	if (input == nullptr || foundMode == nullptr)
		return false;

	*foundMode = nullptr;
	CComPtr<IDeckLinkDisplayModeIterator> iterator;
	if (input->GetDisplayModeIterator(&iterator) != S_OK)
		return false;

	return FindDeckLinkDisplayMode(iterator, targetMode, foundMode);
}

unsigned char ClampToByte(double value)
{
	if (value <= 0.0)
		return 0;
	if (value >= 255.0)
		return 255;
	return static_cast<unsigned char>(value + 0.5);
}

void StoreRec709UyvyAsBgra(unsigned char yByte, unsigned char uByte, unsigned char vByte, unsigned char* destination)
{
	const double y = (static_cast<double>(yByte) - 16.0) / 219.0;
	const double cb = (static_cast<double>(uByte) - 16.0) / 224.0 - 0.5;
	const double cr = (static_cast<double>(vByte) - 16.0) / 224.0 - 0.5;
	const double red = y + 1.5748 * cr;
	const double green = y - 0.1873 * cb - 0.4681 * cr;
	const double blue = y + 1.8556 * cb;
	destination[0] = ClampToByte(blue * 255.0);
	destination[1] = ClampToByte(green * 255.0);
	destination[2] = ClampToByte(red * 255.0);
	destination[3] = 255;
}
}

DeckLinkInputCallback::DeckLinkInputCallback(DeckLinkInput& owner) :
	mOwner(owner)
{
}

HRESULT STDMETHODCALLTYPE DeckLinkInputCallback::QueryInterface(REFIID iid, LPVOID* ppv)
{
	if (ppv == nullptr)
		return E_POINTER;
	if (iid == IID_IUnknown || iid == IID_IDeckLinkInputCallback)
	{
		*ppv = static_cast<IDeckLinkInputCallback*>(this);
		AddRef();
		return S_OK;
	}
	*ppv = nullptr;
	return E_NOINTERFACE;
}

ULONG STDMETHODCALLTYPE DeckLinkInputCallback::AddRef()
{
	return ++mRefCount;
}

ULONG STDMETHODCALLTYPE DeckLinkInputCallback::Release()
{
	const ULONG refCount = --mRefCount;
	if (refCount == 0)
		delete this;
	return refCount;
}

HRESULT STDMETHODCALLTYPE DeckLinkInputCallback::VideoInputFrameArrived(IDeckLinkVideoInputFrame* videoFrame, IDeckLinkAudioInputPacket*)
{
	if (videoFrame != nullptr)
		mOwner.HandleFrameArrived(videoFrame);
	return S_OK;
}

HRESULT STDMETHODCALLTYPE DeckLinkInputCallback::VideoInputFormatChanged(BMDVideoInputFormatChangedEvents, IDeckLinkDisplayMode*, BMDDetectedVideoInputFormatFlags)
{
	mOwner.HandleFormatChanged();
	return S_OK;
}

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

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

bool DeckLinkInput::Initialize(const DeckLinkInputConfig& config, std::string& error)
{
	ReleaseResources();
	mConfig = config;
	Log("decklink-input", "Initializing DeckLink input for " + config.videoFormat.displayName + ".");

	if (!DiscoverInput(config, error))
		return false;

	if (mInput->EnableVideoInput(config.videoFormat.displayMode, mCapturePixelFormat, bmdVideoInputFlagDefault) != S_OK)
	{
		error = "DeckLink input setup failed while enabling " +
			std::string(mCapturePixelFormat == bmdFormat8BitBGRA ? "BGRA8" : "UYVY8") +
			" input for " + config.videoFormat.displayName + ".";
		ReleaseResources();
		return false;
	}
	Log(
		"decklink-input",
		std::string("DeckLink input enabled in ") + (mCapturePixelFormat == bmdFormat8BitBGRA ? "BGRA8" : "UYVY8-to-BGRA8 conversion") + " mode.");

	mCallback.Attach(new (std::nothrow) DeckLinkInputCallback(*this));
	if (mCallback == nullptr)
	{
		error = "DeckLink input setup failed while creating the capture callback.";
		ReleaseResources();
		return false;
	}

	if (mInput->SetCallback(mCallback) != S_OK)
	{
		error = "DeckLink input setup failed while installing the capture callback.";
		ReleaseResources();
		return false;
	}
	Log("decklink-input", "DeckLink input callback installed.");

	return true;
}

bool DeckLinkInput::Start(std::string& error)
{
	if (mInput == nullptr)
	{
		error = "DeckLink input has not been initialized.";
		return false;
	}
	if (mRunning.load(std::memory_order_acquire))
		return true;

	if (mInput->StartStreams() != S_OK)
	{
		error = "DeckLink input stream failed to start.";
		return false;
	}
	mRunning.store(true, std::memory_order_release);
	Log("decklink-input", "DeckLink input stream started.");
	return true;
}

void DeckLinkInput::Stop()
{
	if (mInput != nullptr && mRunning.exchange(false, std::memory_order_acq_rel))
	{
		mInput->StopStreams();
		Log("decklink-input", "DeckLink input stream stopped.");
	}
}

void DeckLinkInput::ReleaseResources()
{
	Stop();
	if (mInput != nullptr)
	{
		mInput->SetCallback(nullptr);
		mInput->DisableVideoInput();
	}
	mCallback.Release();
	mInput.Release();
}

DeckLinkInputMetrics DeckLinkInput::Metrics() const
{
	DeckLinkInputMetrics 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 = mConvertMilliseconds.load(std::memory_order_relaxed);
	metrics.submitMilliseconds = mSubmitMilliseconds.load(std::memory_order_relaxed);
	metrics.captureFormat = CaptureFormatName();
	return metrics;
}

void DeckLinkInput::HandleFrameArrived(IDeckLinkVideoInputFrame* inputFrame)
{
	if (inputFrame == nullptr)
		return;

	if ((inputFrame->GetFlags() & bmdFrameHasNoInputSource) == bmdFrameHasNoInputSource)
	{
		mNoInputSourceFrames.fetch_add(1, std::memory_order_relaxed);
		bool expected = false;
		if (mLoggedNoInputSource.compare_exchange_strong(expected, true, std::memory_order_relaxed))
			LogWarning("decklink-input", "DeckLink input callback reports no input source.");
		return;
	}

	if (inputFrame->GetWidth() != static_cast<long>(mMailbox.Config().width) ||
		inputFrame->GetHeight() != static_cast<long>(mMailbox.Config().height))
	{
		mUnsupportedFrames.fetch_add(1, std::memory_order_relaxed);
		bool expected = false;
		if (mLoggedUnsupportedFrame.compare_exchange_strong(expected, true, std::memory_order_relaxed))
			LogWarning("decklink-input", "DeckLink input frame dimensions do not match the configured mailbox.");
		return;
	}

	CComPtr<IDeckLinkVideoBuffer> inputFrameBuffer;
	if (inputFrame->QueryInterface(IID_IDeckLinkVideoBuffer, reinterpret_cast<void**>(&inputFrameBuffer)) != S_OK)
	{
		mUnsupportedFrames.fetch_add(1, std::memory_order_relaxed);
		bool expected = false;
		if (mLoggedUnsupportedFrame.compare_exchange_strong(expected, true, std::memory_order_relaxed))
			LogWarning("decklink-input", "DeckLink input frame does not expose IDeckLinkVideoBuffer.");
		return;
	}

	if (inputFrameBuffer->StartAccess(bmdBufferAccessRead) != S_OK)
	{
		mUnsupportedFrames.fetch_add(1, std::memory_order_relaxed);
		bool expected = false;
		if (mLoggedUnsupportedFrame.compare_exchange_strong(expected, true, std::memory_order_relaxed))
			LogWarning("decklink-input", "DeckLink input frame buffer could not be opened for read access.");
		return;
	}

	void* bytes = nullptr;
	inputFrameBuffer->GetBytes(&bytes);
	bool submitted = false;
	if (mCapturePixelFormat == bmdFormat8BitBGRA)
		submitted = SubmitBgra8Frame(inputFrame, bytes);
	else if (mCapturePixelFormat == bmdFormat8BitYUV)
		submitted = SubmitUyvy8Frame(inputFrame, bytes);
	else
		mUnsupportedFrames.fetch_add(1, 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))
			LogWarning("decklink-input", "DeckLink input frame could not be submitted to InputFrameMailbox.");
	}
	mCapturedFrames.fetch_add(1, std::memory_order_relaxed);
	bool expectedFirstFrame = false;
	if (mLoggedFirstFrame.compare_exchange_strong(expectedFirstFrame, true, std::memory_order_relaxed))
	{
		Log(
			"decklink-input",
			std::string("First DeckLink ") + (mCapturePixelFormat == bmdFormat8BitBGRA ? "BGRA8" : "UYVY8 converted BGRA8") + " input frame submitted to InputFrameMailbox.");
	}

	inputFrameBuffer->EndAccess(bmdBufferAccessRead);
}

void DeckLinkInput::HandleFormatChanged()
{
	mUnsupportedFrames.fetch_add(1, std::memory_order_relaxed);
	LogWarning("decklink-input", "DeckLink input format changed; input edge does not auto-switch formats yet.");
}

bool DeckLinkInput::DiscoverInput(const DeckLinkInputConfig& config, std::string& error)
{
	CComPtr<IDeckLinkIterator> iterator;
	HRESULT result = CoCreateInstance(CLSID_CDeckLinkIterator, nullptr, CLSCTX_ALL, IID_IDeckLinkIterator, reinterpret_cast<void**>(&iterator));
	if (FAILED(result))
	{
		error = "DeckLink input discovery failed. Blackmagic DeckLink drivers may not be installed.";
		return false;
	}

	CComPtr<IDeckLink> deckLink;
	while (iterator->Next(&deckLink) == S_OK)
	{
		CComPtr<IDeckLinkInput> candidateInput;
		if (deckLink->QueryInterface(IID_IDeckLinkInput, reinterpret_cast<void**>(&candidateInput)) == S_OK && candidateInput != nullptr)
		{
			CComPtr<IDeckLinkDisplayMode> displayMode;
			if (FindInputDisplayMode(candidateInput, config.videoFormat.displayMode, &displayMode) &&
				SupportsInputFormat(candidateInput, config.videoFormat.displayMode, bmdFormat8BitBGRA))
			{
				mInput = candidateInput;
				mCapturePixelFormat = bmdFormat8BitBGRA;
				Log("decklink-input", "DeckLink input device selected for BGRA8 capture.");
				return true;
			}
			if (displayMode != nullptr &&
				SupportsInputFormat(candidateInput, config.videoFormat.displayMode, bmdFormat8BitYUV))
			{
				mInput = candidateInput;
				mCapturePixelFormat = bmdFormat8BitYUV;
				Log("decklink-input", "DeckLink input device selected for UYVY8 capture with CPU BGRA8 conversion.");
				return true;
			}
		}
		deckLink.Release();
	}

	error = "No DeckLink input device supports BGRA8 or UYVY8 capture for " + config.videoFormat.displayName + ".";
	return false;
}

bool DeckLinkInput::SupportsInputFormat(IDeckLinkInput* input, BMDDisplayMode displayMode, BMDPixelFormat pixelFormat) const
{
	if (input == nullptr)
		return false;

	BOOL supported = FALSE;
	BMDDisplayMode actualMode = bmdModeUnknown;
	const HRESULT result = input->DoesSupportVideoMode(
		bmdVideoConnectionUnspecified,
		displayMode,
		pixelFormat,
		bmdNoVideoInputConversion,
		bmdSupportedVideoModeDefault,
		&actualMode,
		&supported);
	return result == S_OK && supported != FALSE;
}

bool DeckLinkInput::SubmitBgra8Frame(IDeckLinkVideoInputFrame* inputFrame, const void* bytes)
{
	if (inputFrame == nullptr || bytes == nullptr)
		return false;
	const uint64_t frameIndex = mCapturedFrames.load(std::memory_order_relaxed);
	mConvertMilliseconds.store(0.0, std::memory_order_relaxed);
	const auto submitStart = std::chrono::steady_clock::now();
	const bool submitted = mMailbox.SubmitFrame(bytes, static_cast<unsigned>(inputFrame->GetRowBytes()), frameIndex);
	const auto submitEnd = std::chrono::steady_clock::now();
	mSubmitMilliseconds.store(
		std::chrono::duration_cast<std::chrono::duration<double, std::milli>>(submitEnd - submitStart).count(),
		std::memory_order_relaxed);
	return submitted;
}

bool DeckLinkInput::SubmitUyvy8Frame(IDeckLinkVideoInputFrame* inputFrame, const void* bytes)
{
	if (inputFrame == nullptr || bytes == nullptr)
		return false;

	const unsigned width = static_cast<unsigned>(inputFrame->GetWidth());
	const unsigned height = static_cast<unsigned>(inputFrame->GetHeight());
	const long sourceRowBytes = inputFrame->GetRowBytes();
	if (width == 0 || height == 0 || sourceRowBytes < static_cast<long>(width * 2u))
		return false;

	const unsigned destinationRowBytes = VideoIORowBytes(VideoIOPixelFormat::Bgra8, width);
	const auto convertStart = std::chrono::steady_clock::now();
	mConversionBuffer.resize(static_cast<std::size_t>(destinationRowBytes) * static_cast<std::size_t>(height));
	const unsigned char* sourceBytes = static_cast<const unsigned char*>(bytes);
	for (unsigned y = 0; y < height; ++y)
	{
		const unsigned char* sourceRow = sourceBytes + static_cast<std::size_t>(y) * static_cast<std::size_t>(sourceRowBytes);
		unsigned char* destinationRow = mConversionBuffer.data() + static_cast<std::size_t>(y) * static_cast<std::size_t>(destinationRowBytes);
		for (unsigned x = 0; x < width; x += 2)
		{
			const unsigned pairOffset = x * 2u;
			const unsigned char u = sourceRow[pairOffset + 0];
			const unsigned char y0 = sourceRow[pairOffset + 1];
			const unsigned char v = sourceRow[pairOffset + 2];
			const unsigned char y1 = sourceRow[pairOffset + 3];
			StoreRec709UyvyAsBgra(y0, u, v, destinationRow + static_cast<std::size_t>(x) * 4u);
			if (x + 1u < width)
				StoreRec709UyvyAsBgra(y1, u, v, destinationRow + static_cast<std::size_t>(x + 1u) * 4u);
		}
	}
	const auto convertEnd = std::chrono::steady_clock::now();
	mConvertMilliseconds.store(
		std::chrono::duration_cast<std::chrono::duration<double, std::milli>>(convertEnd - convertStart).count(),
		std::memory_order_relaxed);

	const uint64_t frameIndex = mCapturedFrames.load(std::memory_order_relaxed);
	const auto submitStart = std::chrono::steady_clock::now();
	const bool submitted = mMailbox.SubmitFrame(mConversionBuffer.data(), destinationRowBytes, frameIndex);
	const auto submitEnd = std::chrono::steady_clock::now();
	mSubmitMilliseconds.store(
		std::chrono::duration_cast<std::chrono::duration<double, std::milli>>(submitEnd - submitStart).count(),
		std::memory_order_relaxed);
	return submitted;
}

const char* DeckLinkInput::CaptureFormatName() const
{
	if (mInput == nullptr)
		return "none";
	if (mCapturePixelFormat == bmdFormat8BitBGRA)
		return "BGRA8";
	if (mCapturePixelFormat == bmdFormat8BitYUV)
		return "UYVY8";
	return "unsupported";
}
}