/* -LICENSE-START-
** Copyright (c) 2018 Blackmagic Design
**  
** Permission is hereby granted, free of charge, to any person or organization 
** obtaining a copy of the software and accompanying documentation (the 
** "Software") to use, reproduce, display, distribute, sub-license, execute, 
** and transmit the Software, and to prepare derivative works of the Software, 
** and to permit third-parties to whom the Software is furnished to do so, in 
** accordance with:
** 
** (1) if the Software is obtained from Blackmagic Design, the End User License 
** Agreement for the Software Development Kit (“EULA”) available at 
** https://www.blackmagicdesign.com/EULA/DeckLinkSDK; or
** 
** (2) if the Software is obtained from any third party, such licensing terms 
** as notified by that third party,
** 
** and all subject to the following:
** 
** (3) the copyright notices in the Software and this entire statement, 
** including the above license grant, this restriction and the following 
** disclaimer, must be included in all copies of the Software, in whole or in 
** part, and all derivative works of the Software, unless such copies or 
** derivative works are solely in the form of machine-executable object code 
** generated by a source language processor.
** 
** (4) THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS 
** OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
** FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT 
** SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE 
** FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, 
** ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 
** DEALINGS IN THE SOFTWARE.
** 
** A copy of the Software is available free of charge at 
** https://www.blackmagicdesign.com/desktopvideo_sdk under the EULA.
** 
** -LICENSE-END-
*/
// CaptureStills.cpp 
//

#include <stdio.h>
#include <condition_variable>
#include <mutex>
#include <queue>
#include <thread>
#include <tuple>
#include <vector>

#include "platform.h"
#include "Bgra32VideoFrame.h"
#include "DeckLinkInputDevice.h"
#include "DeckLinkAPI.h"
#include "ImageWriter.h"

// Pixel format tuple encoding {BMDPixelFormat enum, Pixel format display name}
const std::vector<std::tuple<BMDPixelFormat, std::string>> kSupportedPixelFormats
{
	std::make_tuple(bmdFormat8BitYUV, "8 bit YUV (4:2:2)"),
	std::make_tuple(bmdFormat10BitYUV, "10 bit YUV (4:2:2)"),
	std::make_tuple(bmdFormat8BitARGB, "8 bit ARGB (4:4:4)"),
	std::make_tuple(bmdFormat8BitBGRA, "8 bit BGRA (4:4:4)"),
	std::make_tuple(bmdFormat10BitRGB, "10 bit RGB (4:4:4)"),
	std::make_tuple(bmdFormat12BitRGB, "12 bit RGB (4:4:4)"),
	std::make_tuple(bmdFormat12BitRGBLE, "12 bit RGB (4:4:4) Little-Endian"),
	std::make_tuple(bmdFormat10BitRGBX, "10 bit RGBX (4:4:4)"),
	std::make_tuple(bmdFormat10BitRGBXLE, "10 bit RGBX (4:4:4) Little-Endian"),
};
enum {
	kPixelFormatValue = 0,
	kPixelFormatString
};

void CaptureStills(DeckLinkInputDevice* deckLinkInput, const int captureInterval, const int framesToCapture,
				   const std::string& captureDirectory, const std::string& filenamePrefix)
{
	int							captureFrameCount		= 0;
	HRESULT						result					= S_OK;
	bool						captureRunning			= true;
	
	IDeckLinkVideoFrame*		receivedVideoFrame		= NULL;
	IDeckLinkVideoConversion*	deckLinkFrameConverter	= NULL;
	IDeckLinkVideoFrame*		bgra32Frame				= NULL;

	// Create frame conversion instance
	result = GetDeckLinkVideoConversion(&deckLinkFrameConverter);
	if (result != S_OK)
		return;

	while (captureRunning)
	{
		bool captureCancelled;
		if (!deckLinkInput->WaitForVideoFrameArrived(&receivedVideoFrame, captureCancelled))
		{
			fprintf(stderr, "Timeout waiting for valid frame\n");
			captureRunning = false;
		}

		else if (captureCancelled)
			captureRunning = false;

		else if ((++captureFrameCount % captureInterval) == 0)
		{
			std::string outputFileName;
			result = ImageWriter::GetNextFilenameWithPrefix(captureDirectory, filenamePrefix, outputFileName);
			if (result != S_OK)
			{
				fprintf(stderr, "Unable to get filename\n");
				captureRunning = false;
			}
			else
			{
				fprintf(stderr, "Capturing frame #%d to %s\n", captureFrameCount, outputFileName.c_str());

				if (receivedVideoFrame->GetPixelFormat() == bmdFormat8BitBGRA)
				{
					// Frame is already 8-bit BGRA - no conversion required
					bgra32Frame = receivedVideoFrame;
					bgra32Frame->AddRef();
				}
				else
				{
					bgra32Frame = new Bgra32VideoFrame(receivedVideoFrame->GetWidth(), receivedVideoFrame->GetHeight(), receivedVideoFrame->GetFlags());

					result = deckLinkFrameConverter->ConvertFrame(receivedVideoFrame, bgra32Frame);
					if (FAILED(result))
					{
						fprintf(stderr, "Frame conversion to BGRA was unsuccessful\n");
						captureRunning = false;
					}
				}

				result = ImageWriter::WriteBgra32VideoFrameToPNG(bgra32Frame, outputFileName);
				if (FAILED(result))
				{
					fprintf(stderr, "Image encoding to file was unsuccessful\n");
					captureRunning = false;
				}

				bgra32Frame->Release();
				
				if ((captureFrameCount / captureInterval) >= framesToCapture)
				{
					fprintf(stderr, "Completed Capture\n");
					captureRunning = false;
				}
			}
		}

		if (receivedVideoFrame != NULL)
		{
			receivedVideoFrame->Release();
			receivedVideoFrame = NULL;
		}
	}

	if (deckLinkFrameConverter != NULL)
	{
		deckLinkFrameConverter->Release();
		deckLinkFrameConverter = NULL;
	}
}

void DisplayUsage(DeckLinkInputDevice* selectedDeckLinkInput, const std::vector<std::string>& deviceNames,
	const int selectedDeviceIndex, const int selectedDisplayModeIndex, const bool supportsFormatDetection)
{
	HRESULT								result = E_FAIL;
	std::string							selectedDisplayModeName;
	std::vector<IDeckLinkDisplayMode*>	displayModes;

	fprintf(stderr,
		"Usage: ./CaptureStills -d <device id> -m <mode id> [OPTIONS]\n"
		"\n"
		"    -d <device id>:\n"
		);

	if (deviceNames.empty())
	{
		fprintf(stderr, "        No DeckLink devices found. Please check that Desktop Video is installed.\n");
	}
	else
	{
		// Loop through all available devices
		for (size_t i = 0; i < deviceNames.size(); i++)
		{
			fprintf(stderr,
				"       %c%2d:  %s\n",
				((int)i == selectedDeviceIndex) ? '*' : ' ',
				(int)i,
				deviceNames[i].c_str()
				);
		}
	}

	fprintf(stderr,
		"    -m <mode id>: (%s)\n", ((selectedDeviceIndex >= 0) && (selectedDeviceIndex < (int)deviceNames.size())) ? deviceNames[selectedDeviceIndex].c_str() : ""
		);

	// Loop through all available display modes on the delected DeckLink device
	if (selectedDeckLinkInput == NULL)
	{
		fprintf(stderr, "        No DeckLink device selected\n");
	}
	else
	{
		if (supportsFormatDetection)
		{
			fprintf(stderr, "       %c-1:  auto detect format\n",
				(selectedDisplayModeIndex == -1) ? '*' : ' '
				);
		}

		displayModes = selectedDeckLinkInput->GetDisplayModeList();

		for (size_t i = 0; i < displayModes.size(); i++)
		{
			dlstring_t displayModeName;

			result = displayModes[i]->GetName(&displayModeName);
			if (result == S_OK)
			{
				BMDTimeValue frameRateDuration;
				BMDTimeValue frameRateScale;

				displayModes[i]->GetFrameRate(&frameRateDuration, &frameRateScale);

				fprintf(stderr,
					"       %c%2d:  %-20s \t %4li x %4li \t %.2f FPS\n",
					((int)i == selectedDisplayModeIndex) ? '*' : ' ',
					(int)i,
					DlToCString(displayModeName),
					displayModes[i]->GetWidth(),
					displayModes[i]->GetHeight(),
					(double)frameRateScale / (double)frameRateDuration
					);

				if ((int)i == selectedDisplayModeIndex)
					selectedDisplayModeName = DlToStdString(displayModeName);

				DeleteString(displayModeName);
			}
		}
	}

	fprintf(stderr, "    -p <pixelformat>: ");

	if (selectedDeckLinkInput == NULL)
		fprintf(stderr, "\n        No DeckLink device selected\n");

	else if ((selectedDisplayModeIndex < -1) || (selectedDisplayModeIndex >= (int)displayModes.size()))
		fprintf(stderr, "\n        Invalid display mode selected\n");

	else if (selectedDisplayModeIndex == -1)
		fprintf(stderr, "\n        Auto-detect mode selected\n");

	else
	{
		fprintf(stderr, "(%s)\n", selectedDisplayModeName.c_str());

		for (unsigned int i = 0; i < kSupportedPixelFormats.size(); i++)
		{
			// Check whether pixel format is supported for display mode
			dlbool_t displayModeSupported;
			result = selectedDeckLinkInput->GetDeckLinkInput()->DoesSupportVideoMode(bmdVideoConnectionUnspecified,
																					 displayModes[selectedDisplayModeIndex]->GetDisplayMode(),
																					 std::get<kPixelFormatValue>(kSupportedPixelFormats[i]),
																					 bmdNoVideoInputConversion,
																					 bmdSupportedVideoModeDefault,
																					 NULL,
																					 &displayModeSupported);

			if ((result == S_OK) && (displayModeSupported))
			{
				fprintf(stderr,
					"        %2d:  %s%s\n",
					i,
					std::get<kPixelFormatString>(kSupportedPixelFormats[i]).c_str(),
					(i == 0) ? " (default)" : ""
					);
			}
		}
	}

	fprintf(stderr,
		"    -n <frames>          Number of frames to capture (default is 1)\n"
		"    -i <interval>        Capture frame interval rate (default is 1 - every frame)\n"
		"    -f <prefix>          Filename prefix (default is \"image_\")\n"
		"    <capturedirectory>\n"
		"\n"
		"Capture image stills to a specified directory. eg:\n"
		"\n"
		"    ./CaptureStills -d 0 -m 2 -n 10 -i 60 ~/Pictures/\n\n"
		);
}


int main(int argc, char* argv[])
{
	// Configuration Flags
	bool						displayHelp				= false;
	int							deckLinkIndex			= -1;
	int							displayModeIndex		= -2;
	int							framesToCapture			= 1;
	int							captureInterval			= 1;
	int							pixelFormatIndex		= 0;
	bool						enableFormatDetection	= false;
	std::string					filenamePrefix;
	std::string					captureDirectory;

	HRESULT						result;
	int							exitStatus = 1;
	int							idx;
	bool						supportsFormatDetection = false;

	std::thread					captureStillsThread;
	std::thread					keyPressThread;

	IDeckLinkIterator*			deckLinkIterator		= NULL;
	IDeckLink*					deckLink				= NULL;
	DeckLinkInputDevice*		selectedDeckLinkInput	= NULL;

	BMDDisplayMode				selectedDisplayMode		= bmdModeNTSC;
	std::string					selectedDisplayModeName;
	std::vector<std::string>	deckLinkDeviceNames;


	result = GetDeckLinkIterator(&deckLinkIterator);
	if (result != S_OK)
		goto bail;

	for (int i = 1; i < argc; i++)
	{
		if (strcmp(argv[i], "-d") == 0)
			deckLinkIndex = atoi(argv[++i]);

		else if (strcmp(argv[i], "-m") == 0)
			displayModeIndex = atoi(argv[++i]);

		else if (strcmp(argv[i], "-i") == 0)
			captureInterval = atoi(argv[++i]);

		else if (strcmp(argv[i], "-n") == 0)
			framesToCapture = atoi(argv[++i]);

		else if (strcmp(argv[i], "-p") == 0)
			pixelFormatIndex = atoi(argv[++i]);

		else if (strcmp(argv[i], "-f") == 0)
			filenamePrefix = argv[++i];

		else if ((strcmp(argv[i], "?") == 0) || (strcmp(argv[i], "-h") == 0))
			displayHelp = true;

		else if (i == argc - 1)
		{
			captureDirectory = argv[i];
		}
	}

	if (captureDirectory.empty())
	{
		fprintf(stderr, "You must set a capture directory\n");
		displayHelp = true;
	}
	else if (!IsPathDirectory(captureDirectory))
	{
		fprintf(stderr, "Invalid directory specified\n");
		displayHelp = true;
	}

	if (deckLinkIndex < 0)
	{
		fprintf(stderr, "You must select a device\n");
		displayHelp = true;
	}

	// Obtain the required DeckLink device
	idx = 0;

	while ((result = deckLinkIterator->Next(&deckLink)) == S_OK)
	{
		dlstring_t deckLinkName;

		result = deckLink->GetDisplayName(&deckLinkName);
		if (result == S_OK)
		{
			deckLinkDeviceNames.push_back(DlToStdString(deckLinkName));
			DeleteString(deckLinkName);
		}

		if (idx++ == deckLinkIndex)
		{
			// Check that selected device supports capture
			IDeckLinkProfileAttributes*	deckLinkAttributes = NULL;
			int64_t						ioSupportAttribute = 0;
			dlbool_t					formatDetectionSupportAttribute;

			result = deckLink->QueryInterface(IID_IDeckLinkProfileAttributes, (void**)&deckLinkAttributes);

			if (result != S_OK)
			{
				fprintf(stderr, "Unable to get IDeckLinkAttributes interface\n");
				goto bail;
			}

			// Check whether device supports cpature
			result = deckLinkAttributes->GetInt(BMDDeckLinkVideoIOSupport, &ioSupportAttribute);

			if ((result != S_OK) || ((ioSupportAttribute & bmdDeviceSupportsCapture) == 0))
			{
				fprintf(stderr, "Selected device does not support capture\n");
				displayHelp = true;
			}
			else
			{
				// Check if input mode detection is supported.
				result = deckLinkAttributes->GetFlag(BMDDeckLinkSupportsInputFormatDetection, &formatDetectionSupportAttribute);
				supportsFormatDetection = (result == S_OK) && formatDetectionSupportAttribute;

				selectedDeckLinkInput = new DeckLinkInputDevice(deckLink);
			}
			
			deckLinkAttributes->Release();
		}

		deckLink->Release();
	}

	// Get display modes from the selected decklink output 
	if (selectedDeckLinkInput != NULL)
	{
		result = selectedDeckLinkInput->Init();
		if (result != S_OK)
		{
			fprintf(stderr, "Unable to initialize DeckLink input interface");
			goto bail;
		}

		// Get the display mode
		if ((displayModeIndex < -1) || (displayModeIndex >= (int)selectedDeckLinkInput->GetDisplayModeList().size()))
		{
			fprintf(stderr, "You must select a valid display mode\n");
			displayHelp = true;
		}
		else if (displayModeIndex == -1)
		{
			if (!supportsFormatDetection)
			{
				fprintf(stderr, "Format detection is not supported on this device\n");
				displayHelp = true;
			}
			else
			{
				enableFormatDetection = true;

				// Format detection still needs a valid mode to start with
				selectedDisplayMode = bmdModeNTSC;
				selectedDisplayModeName = "Automatic mode detection";
				pixelFormatIndex = 0;
			}
		}
		else if ((pixelFormatIndex < 0) || (pixelFormatIndex >= (int)kSupportedPixelFormats.size()))
		{
			fprintf(stderr, "You must select a valid pixel format\n");
			displayHelp = true;
		}
		else
		{
			dlbool_t				displayModeSupported;
			dlstring_t				displayModeNameStr;
			IDeckLinkDisplayMode*	displayMode = selectedDeckLinkInput->GetDisplayModeList()[displayModeIndex];

			result = displayMode->GetName(&displayModeNameStr);
			if (result == S_OK)
			{
				selectedDisplayModeName = DlToStdString(displayModeNameStr);
				DeleteString(displayModeNameStr);
			}

			selectedDisplayMode = displayMode->GetDisplayMode();

			// Check display mode is supported with given options
			result = selectedDeckLinkInput->GetDeckLinkInput()->DoesSupportVideoMode(bmdVideoConnectionUnspecified,
																					 selectedDisplayMode,
																					 std::get<kPixelFormatValue>(kSupportedPixelFormats[pixelFormatIndex]),
																					 bmdNoVideoInputConversion,
																					 bmdSupportedVideoModeDefault,
																					 NULL,
																					 &displayModeSupported);
			if ((result != S_OK) || (!displayModeSupported))
			{
				fprintf(stderr, "Display mode %s with pixel format %s is not supported by device\n", 
					selectedDisplayModeName.c_str(),
					std::get<kPixelFormatString>(kSupportedPixelFormats[pixelFormatIndex]).c_str()
					);
				displayHelp = true;
			}
		}
	}
	else
	{
		fprintf(stderr, "Invalid input device selected\n");
		displayHelp = true;
	}

	if (filenamePrefix.empty())
	{
		filenamePrefix = "image_";
	}

	if (displayHelp)
	{
		DisplayUsage(selectedDeckLinkInput, deckLinkDeviceNames, deckLinkIndex, displayModeIndex, supportsFormatDetection);
		goto bail;
	}

	// Start capturing
	result = selectedDeckLinkInput->StartCapture(selectedDisplayMode, std::get<kPixelFormatValue>(kSupportedPixelFormats[pixelFormatIndex]), enableFormatDetection);
	if (result != S_OK)
		goto bail;

	// Print the selected configuration
	fprintf(stderr, "Capturing with the following configuration:\n"
		" - Capture device: %s\n"
		" - Video mode: %s\n"
		" - Pixel format: %s\n"
		" - Frames to capture: %d\n"
		" - Capture interval: %d\n"
		" - Filename prefix: %s\n"
		" - Capture directory: %s\n",
		selectedDeckLinkInput->GetDeviceName().c_str(),
		selectedDisplayModeName.c_str(),
		std::get<kPixelFormatString>(kSupportedPixelFormats[pixelFormatIndex]).c_str(),
		framesToCapture,
		captureInterval,
		filenamePrefix.c_str(),
		captureDirectory.c_str()
		);

	fprintf(stderr, "Starting capture, press <RETURN> to stop/exit\n");

	// Start thread for capture processing
	captureStillsThread = std::thread([&]{
		CaptureStills(selectedDeckLinkInput, captureInterval, framesToCapture, captureDirectory, filenamePrefix);
	});

	keyPressThread = std::thread([&]{
		getchar();
		selectedDeckLinkInput->CancelCapture();
	});
	
	// Wait on return of main capture stills thread
	captureStillsThread.join();
	selectedDeckLinkInput->StopCapture();

	keyPressThread.join();

	// All Okay.
	exitStatus = 0;

bail:
	if (selectedDeckLinkInput != NULL)
	{
		selectedDeckLinkInput->Release();
		selectedDeckLinkInput = NULL;
	}

	if (deckLinkIterator != NULL)
	{
		deckLinkIterator->Release();
		deckLinkIterator = NULL;
	}

	return exitStatus;
}
