// RtspServer.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include "pch.h"
#include "framework.h"
#include "RtspServer.h"

#include <iostream>
#include <ctime>
#include <vector>
#include <functional>

#ifdef _DEBUG
#define new DEBUG_NEW
#endif

#pragma warning(disable:4819)
#pragma warning(disable:4996)

extern "C"
{
#include "libavcodec/avcodec.h"
#include "libavformat/avformat.h"
#include "libavfilter/avfilter.h"
#include "libavdevice/avdevice.h"
#include "libswscale/swscale.h"
#include "libavutil/imgutils.h"
}

#pragma comment(lib, "avcodec.lib")
#pragma comment(lib, "avformat.lib")
#pragma comment(lib, "avutil.lib")
#pragma comment(lib, "avfilter.lib")
#pragma comment(lib, "avdevice.lib")
#pragma comment(lib, "swscale.lib")
#pragma comment(lib, "swresample.lib")

#ifndef MEM_ALIGN
#define MEM_ALIGN (64)
#endif

#define CAPTURE_WIDTH (1920)
#define CAPTURE_HEIGHT (1080)

#define ENCODER_WIDTH (1920)
#define ENCODER_HEIGHT (1080)
#define ENCODER_FPS (30)
#define ENCODER_Q_SIZE (128)

class CDataPacket
{
protected:
	uint8_t* m_pData;
	int m_nCapacity;
	int m_nSize;

public:
	CDataPacket(const int nSize=0) : m_pData(NULL), m_nCapacity(0), m_nSize(0)
	{
		Resize(nSize);
	}

	CDataPacket(const CDataPacket& d) : m_pData(NULL), m_nCapacity(0), m_nSize(0)
	{
		Resize(d.m_nSize);

		if (d.m_nSize != 0)
			memcpy(m_pData, d.m_pData, d.m_nSize);
	}

	CDataPacket(CDataPacket&& d) : m_pData(NULL), m_nCapacity(0), m_nSize(0)
	{
		Swap(d);
	}

	~CDataPacket()
	{
		if (m_pData != NULL)
			av_free(m_pData);
	}

	void Resize(const int& nSize)
	{
		if (nSize <0)
		{
			TRACE("Invalid data packet size!\n");
			return;
		}

		if (nSize > m_nCapacity)
		{
			if (m_pData != NULL)
			{
				av_free(m_pData);
				m_nCapacity = 0;
				m_pData = NULL;
			}

			if (nSize > 0)
			{
				m_pData = (uint8_t*)av_mallocz(nSize);
				m_nCapacity = nSize;
			}
		}

		m_nSize = nSize;
	}

	void Clear()
	{
		Resize(0);
	}

	void Copy(const uint8_t* pPtr, const int& nSize)
	{
		if (nSize <= 0)
			return;

		if (nSize > m_nSize)
			Resize(nSize);

		memcpy(m_pData, pPtr, nSize);
	}

	bool IsEmpty() const
	{
		return (m_nSize == 0);
	}

	int GetSize() const
	{
		return m_nSize;
	}

	uint8_t* GetData()
	{
		return m_pData;
	}

	void Swap(CDataPacket& d)
	{
		std::swap(m_pData, d.m_pData);
		std::swap(m_nSize, d.m_nSize);
		std::swap(m_nCapacity, d.m_nCapacity);
	}
};

class CCapture
{
public:
	CCapture()
	{
		m_cbFunc = nullptr;

		m_bReady = false;

		m_pInputFormat = NULL;
		m_pFormatCtx = NULL;
		m_nVideoStreamIdx = -1;
		m_pCodecCtx = NULL;
		m_pSwsCtx = NULL;
		m_pFrame = NULL;
	}

	~CCapture()
	{
		Close();
	}

	bool Open(LPCSTR lpszVideo=NULL, LPCSTR lpszAudio=NULL);
	bool Close();
	bool Run();

	bool IsReady()
	{
		return m_bReady && (m_cbFunc != nullptr);
	}

	std::function<void(CDataPacket*)> m_cbFunc;

protected:
	bool m_bReady;

	AVInputFormat* m_pInputFormat;
	AVFormatContext* m_pFormatCtx;
	int m_nVideoStreamIdx;
	AVCodecContext* m_pCodecCtx;
	struct SwsContext* m_pSwsCtx;
	AVFrame* m_pFrame;
	int m_nFrameBufSize;
	uint64_t m_nFrameCnt;
	CDataPacket m_dataOut;
};

bool CCapture::Open(LPCSTR lpszVideo, LPCSTR lpszAudio)
{
	std::string deviceName;
	int rval;

	if (lpszVideo != NULL)
	{
		deviceName = "video=";
		deviceName.append(lpszVideo);
		if (lpszAudio != NULL)
		{
			deviceName.append(":audio=");
			deviceName.append(lpszAudio);
		}

		m_pInputFormat = av_find_input_format("dshow");
	}
	else
	{
		m_pInputFormat = av_find_input_format("vfwcap");
	}

	AVDictionary* opts = NULL;
	rval = av_dict_parse_string(&opts, "video_size=640x480;pixel_format=yuyv422;framerate=30;", "=", ";", 0);
	rval = av_dict_set(&opts, "rtbufsize", "1000000000", NULL);
	rval = avformat_open_input(&m_pFormatCtx, deviceName.c_str(), m_pInputFormat, &opts);
	av_dict_free(&opts);
	if (rval != 0)
	{
		char buf[AV_ERROR_MAX_STRING_SIZE];
		av_strerror(rval, buf, sizeof(buf));
		TRACE("Couldn't open input source '%s'! with error '%s'\n", deviceName.c_str(), buf);
		return false;
	}

	rval = avformat_find_stream_info(m_pFormatCtx, NULL);
	if (rval < 0)
	{
		TRACE("Couldn't find stream information!\n");
		return false;
	}

	av_dump_format(m_pFormatCtx, 0, m_pFormatCtx->url, 0);

	m_nVideoStreamIdx = -1;
	for (unsigned int i=0; i<m_pFormatCtx->nb_streams; ++i)
	{
		if (m_pFormatCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO)
		{
			m_nVideoStreamIdx = i;
			break;
		}
	}
	if (m_nVideoStreamIdx == -1)
	{
		TRACE("Didn't find a video stream!\n");
		return false;
	}

	AVCodec* pCodec = avcodec_find_decoder(m_pFormatCtx->streams[m_nVideoStreamIdx]->codec->codec_id);
	if (pCodec == NULL)
	{
		TRACE("Codec not found!\n");
		return false;
	}

	m_pCodecCtx = avcodec_alloc_context3(pCodec);

	avcodec_parameters_to_context(m_pCodecCtx, m_pFormatCtx->streams[m_nVideoStreamIdx]->codecpar);

	if (avcodec_open2(m_pCodecCtx, pCodec, NULL) < 0)
	{
		TRACE("Couldn't open codec!\n");
		return false;
	}

	m_pSwsCtx = sws_getCachedContext(
		NULL,
		m_pCodecCtx->width,
		m_pCodecCtx->height,
		m_pCodecCtx->pix_fmt,
		CAPTURE_WIDTH,
		CAPTURE_HEIGHT,
		AV_PIX_FMT_NV12,
		SWS_BICUBIC,
		NULL,
		NULL,
		NULL);
	if (m_pSwsCtx == NULL)
	{
		TRACE("Can't initialize the conversion context!\n");
		return false;
	}

	m_pFrame = av_frame_alloc();
	if (m_pFrame == NULL)
	{
		TRACE("Allocating output frame failed!\n");
		return false;
	}

	m_pFrame->width = CAPTURE_WIDTH;
	m_pFrame->height = CAPTURE_HEIGHT;
	m_pFrame->format = AV_PIX_FMT_NV12;
	rval = av_frame_get_buffer(m_pFrame, MEM_ALIGN);
	if (rval < 0)
	{
		TRACE("Allocating frame data buffer failed!\n");
		return false;
	}

	m_nFrameBufSize = av_image_get_buffer_size(AV_PIX_FMT_NV12, m_pFrame->width, m_pFrame->height, MEM_ALIGN);
	m_nFrameCnt = 0;

	m_dataOut.Resize(m_nFrameBufSize);

	m_bReady = true;

	return true;
}

bool CCapture::Close()
{
	if (m_pFrame != NULL)
	{
		av_free(m_pFrame);
		m_pFrame = NULL;
	}

	sws_freeContext(m_pSwsCtx);
	m_pSwsCtx = NULL;

	avcodec_free_context(&m_pCodecCtx);
	m_pCodecCtx = NULL;

	avformat_close_input(&m_pFormatCtx);

	m_bReady = false;

	return true;
}

bool CCapture::Run()
{
	AVPacket* pPacket;
	AVFrame* pFrame;
	int rval;

	pPacket = av_packet_alloc();
	if (pPacket == NULL)
	{
		TRACE("Allocating packet failed!\n");
		return false;
	}

	pFrame = av_frame_alloc();
	if (pFrame == NULL)
	{
		TRACE("Allocating frame failed!\n");
		return false;
	}

	while (1)
	{
		rval = av_read_frame(m_pFormatCtx, pPacket);
		if (rval < 0)
		{
			TRACE("Read Error!\n");
			break;
		}

		if ((pPacket->buf == NULL) || (pPacket->stream_index != m_nVideoStreamIdx))
			continue;

		rval = avcodec_send_packet(m_pCodecCtx, pPacket);
		if ((rval != 0) && (rval != AVERROR(EAGAIN)))
		{
			TRACE("Sending a packet for decoding failed!\n");
			return false;
		}

		while (1)
		{
			rval = avcodec_receive_frame(m_pCodecCtx, pFrame);
			if (rval == AVERROR(EAGAIN))
			{
				break;
			}
			else if (rval != 0)
			{
				TRACE("Reading a frame from decoder failed!\n");
				return false;
			}

			rval = sws_scale(m_pSwsCtx, pFrame->data, pFrame->linesize, 0, pFrame->height, m_pFrame->data, m_pFrame->linesize);
			if (rval <= 0)
			{
				TRACE("Scaling picture failed!\n");
				return false;
			}

			m_dataOut.Resize(m_nFrameBufSize);

			av_image_copy_to_buffer(m_dataOut.GetData(), m_dataOut.GetSize(), m_pFrame->data, m_pFrame->linesize, (enum AVPixelFormat)m_pFrame->format, m_pFrame->width, m_pFrame->height, MEM_ALIGN);

			++m_nFrameCnt;

			if (m_cbFunc != nullptr)
				m_cbFunc(&m_dataOut);
		}

		Sleep(1);

		av_packet_unref(pPacket);
		av_frame_unref(pFrame);
	}

	av_packet_free(&pPacket);
	av_frame_free(&pFrame);

	return true;
}

template<std::size_t nSize> class TCircularQueue
{
protected:
	std::size_t m_nFront;
	std::size_t m_nRear;
	std::size_t m_nAvailableSize;
	const std::size_t m_nSize;

public:
	TCircularQueue() : m_nFront(0), m_nRear(nSize-1), m_nAvailableSize(0), m_nSize(nSize)
	{
	}

	size_t Push()
	{
		if (m_nAvailableSize == m_nSize)
			Pop();

		if (++m_nRear == m_nSize)
			m_nRear = 0;

		++m_nAvailableSize;

		return m_nRear;
	}

	size_t Pop()
	{
		size_t nPos = (size_t)-1;

		if (m_nAvailableSize > 0)
		{
			nPos = m_nFront;

			if (++m_nFront == m_nSize)
				m_nFront = 0;

			--m_nAvailableSize;
		}

		return nPos;
	}

	size_t Front() const
	{
		return m_nFront;
	}

	size_t Rear() const
	{
		return m_nRear;
	}

	size_t GetAvailableSize() const
	{
		return m_nAvailableSize;
	}

	size_t GetSize() const
	{
		return m_nSize;
	}

	bool IsFull() const
	{
		return (m_nAvailableSize == m_nSize);
	}

	bool IsEmpty() const
	{
		return (m_nAvailableSize == 0);
	}
};

class CEncoder
{
public:
	CEncoder()
	{
		m_cbFunc = NULL;

		m_bReady = false;

		m_nWidth = ENCODER_WIDTH;
		m_nHeight = ENCODER_HEIGHT;
		m_frameRate = av_d2q(ENCODER_FPS, 100);
		m_nBitRate = 5 * 1024 * 1024;

		m_pCodec = NULL;
		m_pCodecCtx = NULL;
		m_pSwsCtx = NULL;
		m_pFrame = NULL;

		InitializeCriticalSectionAndSpinCount(&m_mutexOut, 100);
	}

	~CEncoder()
	{
		Close();

		DeleteCriticalSection(&m_mutexIn);
		DeleteCriticalSection(&m_mutexOut);
	}

	int GetWidth() const
	{
		return m_nWidth;
	}

	int GetHeight() const
	{
		return m_nHeight;
	}

	void SetSize(int nWidth, int nHeight)
	{
		m_nWidth = nWidth;
		m_nHeight = nHeight;
	}

	double GetFps() const
	{
		return av_q2d(m_frameRate);
	}

	void SetFps(double fps)
	{
		m_frameRate = av_d2q(fps, 100);
	}

	uint64_t GetBitRate() const
	{
		return m_nBitRate;
	}

	void SetBitRate(uint64_t nBitRate)
	{
		m_nBitRate = nBitRate;
	}

	const uint8_t* GetSPSNAL()
	{
		if (!m_dataSPSNAL.IsEmpty())
			return m_dataSPSNAL.GetData();

		return NULL;
	}

	int GetSPSNALSize()
	{
		return m_dataSPSNAL.GetSize();
	}

	const uint8_t* GetPPSNAL()
	{
		if (!m_dataPPSNAL.IsEmpty())
			return m_dataPPSNAL.GetData();

		return NULL;
	}

	int GetPPSNALSize()
	{
		return m_dataPPSNAL.GetSize();
	}

	bool Setup();
	bool Close();
	bool Run();

	bool IsReady()
	{
		return m_bReady;
	}

	void ReceiveFrame(CDataPacket* pData);

	void SendPacket(AVPacket* pPacket);
	bool FetchPacket(uint8_t** ppFrameBuffer, uint32_t* pFrameSize);

	std::function<void()> m_cbFunc;

protected:
	bool m_bReady;

	int m_nWidth;
	int m_nHeight;
	AVRational m_frameRate;
	uint64_t m_nBitRate;

	AVCodec* m_pCodec;
	AVCodecContext* m_pCodecCtx;
	struct SwsContext* m_pSwsCtx;
	AVFrame* m_pFrame;
	uint64_t m_nFrameCnt;
	std::clock_t m_lFrameIntvl;

	CDataPacket m_dataSPSNAL;
	CDataPacket m_dataPPSNAL;

	CDataPacket m_dataIn;
	bool m_bInState;
	CRITICAL_SECTION m_mutexIn;

	CDataPacket m_qOut[ENCODER_Q_SIZE];
	TCircularQueue<ENCODER_Q_SIZE> m_qOutPos;
	CRITICAL_SECTION m_mutexOut;

	CDataPacket m_dataOut;
	uint64_t m_nOutCnt;

	static const uint8_t* FindNextNAL(const uint8_t* pStart, const uint8_t* pEnd);
};

bool CEncoder::Setup()
{
	int rval;

	m_pCodec = avcodec_find_encoder_by_name("libx264");
	if (m_pCodec == NULL)
	{
		TRACE("Failed to find 'libx264' codec!\n");
		return false;
	}

	m_pCodecCtx = avcodec_alloc_context3(m_pCodec);
	if (m_pCodecCtx == NULL)
	{
		TRACE("Could not allocate video codec context!\n");
		return false;
	}

	m_pCodecCtx->bit_rate = m_nBitRate;
	m_pCodecCtx->flags &= ~AV_CODEC_FLAG_GLOBAL_HEADER; // Set AV_CODEC_FLAG_GLOBAL_HEADER will disable b_repeat_header in libx264
	m_pCodecCtx->time_base = av_inv_q(m_frameRate);
	m_pCodecCtx->width = m_nWidth;
	m_pCodecCtx->height = m_nHeight;
	m_pCodecCtx->gop_size = 10;
	m_pCodecCtx->pix_fmt = AV_PIX_FMT_NV12;
	m_pCodecCtx->max_b_frames = 1;
	m_pCodecCtx->framerate = m_frameRate;
	m_pCodecCtx->profile = FF_PROFILE_H264_BASELINE;

	AVDictionary* opts = NULL;
	av_dict_set(&opts, "preset", "veryfast", 0);
	av_dict_set(&opts, "profile", "high", 0);
	av_dict_set(&opts, "vcm", "1", 0);
	rval = avcodec_open2(m_pCodecCtx, NULL, &opts);
	av_dict_free(&opts);
	if (rval != 0)
	{
		TRACE("Couldn't open codec!\n");
		return false;
	}

	m_pSwsCtx = sws_getContext(
		CAPTURE_WIDTH,
		CAPTURE_HEIGHT,
		AV_PIX_FMT_NV12,
		m_pCodecCtx->width,
		m_pCodecCtx->height,
		AV_PIX_FMT_NV12,
		SWS_BILINEAR,
		NULL,
		NULL,
		NULL);
	if (m_pSwsCtx == NULL)
	{
		TRACE("Can't initialize the conversion context!\n");
		return false;
	}

	m_pFrame = av_frame_alloc();
	if (m_pFrame == NULL)
	{
		TRACE("Allocating input frame failed!\n");
		return false;
	}

	m_pFrame->width = m_pCodecCtx->width;
	m_pFrame->height = m_pCodecCtx->height;
	m_pFrame->format = m_pCodecCtx->pix_fmt;

	rval = av_frame_get_buffer(m_pFrame, MEM_ALIGN);
	if (rval < 0)
	{
		TRACE("Allocating frame data buffer failed!\n");
		return false;
	}

	av_frame_make_writable(m_pFrame);

	AVPacket* pPacket = av_packet_alloc();
	if (pPacket == NULL)
	{
		TRACE("Allocating packet failed!\n");
		return false;
	}

	m_nFrameCnt = 0;

	while (m_dataSPSNAL.IsEmpty() || m_dataPPSNAL.IsEmpty())
	{
		m_pFrame->pts = m_nFrameCnt;

		rval = avcodec_send_frame(m_pCodecCtx, m_pFrame);
		if ((rval != 0) && (rval != AVERROR(EAGAIN)))
		{
			TRACE("Sending frame for encoding error!\n");
			av_packet_free(&pPacket);
			return false;
		}

		++m_nFrameCnt;

		while (1)
		{
			rval = avcodec_receive_packet(m_pCodecCtx, pPacket);
			if (rval != 0)
			{
				if (rval != AVERROR(EAGAIN))
					TRACE("Receiving encoded packet failed!\n");

				break;
			}

			if (pPacket->size < 4)
				continue;

			const uint8_t* pStart = pPacket->data;
			const uint8_t* pEnd = &pStart[pPacket->size];
			const uint8_t* pNALStart = pStart;
			const uint8_t* pNALEnd = pEnd;

			while (pNALStart < pEnd)
			{
				pNALStart = FindNextNAL(pNALStart, pEnd);

				if (pNALStart == pEnd)
				{
					continue;
				}
				else
				{
					while (*pNALStart == 0)
						++pNALStart; // skip all 0x00

					++pNALStart; // skip 0x01
				}

				pNALEnd = FindNextNAL(pNALStart, pEnd);

				if (((*pNALStart) & 0x1F) == 0x07)
					m_dataSPSNAL.Copy(pNALStart, static_cast<int>(pNALEnd-pNALStart));

				if (((*pNALStart) & 0x1F) == 0x08)
					m_dataPPSNAL.Copy(pNALStart, static_cast<int>(pNALEnd-pNALStart));

				pNALStart = pNALEnd;
			}

			SendPacket(pPacket);
		}
	}

	av_packet_free(&pPacket);

	m_lFrameIntvl = (std::clock_t)(CLOCKS_PER_SEC / av_q2d(m_frameRate));

	m_bInState = false;
	InitializeCriticalSectionAndSpinCount(&m_mutexIn, 1000);

	m_nOutCnt = 0;

	m_bReady = true;

	return true;
}

bool CEncoder::Close()
{
	DeleteCriticalSection(&m_mutexIn);

	av_frame_free(&m_pFrame);
	m_pFrame = NULL;

	sws_freeContext(m_pSwsCtx);
	m_pSwsCtx = NULL;

	avcodec_free_context(&m_pCodecCtx);
	m_pCodecCtx = NULL;

	m_bReady = false;

	return true;
}

bool CEncoder::Run()
{
	clock_t clkPre = std::clock() - m_lFrameIntvl;
	clock_t clkCur;
	clock_t clkIntvl = CLOCKS_PER_SEC / 1000;
	long int msec;
	int rval;

	AVPacket* pPacket = av_packet_alloc();
	if (pPacket == NULL)
	{
		TRACE("Allocating packet failed!\n");
		return false;
	}

	const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(AV_PIX_FMT_NV12);

	AVFrame* pFrame = av_frame_alloc();
	if (pFrame == NULL)
	{
		TRACE("Allocating frame failed!\n");
		av_packet_free(&pPacket);
		return false;
	}

	pFrame->width = CAPTURE_WIDTH;
	pFrame->height = CAPTURE_HEIGHT;
	pFrame->format = AV_PIX_FMT_NV12;

	int nFrameBufSize = av_image_get_buffer_size(AV_PIX_FMT_NV12, CAPTURE_WIDTH, CAPTURE_HEIGHT, MEM_ALIGN);

	CDataPacket dataIn(nFrameBufSize);

	while (1)
	{
		if (m_bInState)
		{
			if (m_dataIn.GetSize() != nFrameBufSize)
			{
				TRACE("Incoming data size is not right! in_data_sz = %d, frame_data_sz = %d\n", m_dataIn.GetSize(), nFrameBufSize);
			}

			if (TryEnterCriticalSection(&m_mutexIn) == 0)
			{
				Sleep(1);
				if (TryEnterCriticalSection(&m_mutexIn) == 0)
					continue;
			}
			dataIn.Swap(m_dataIn);
			m_bInState = false;
			LeaveCriticalSection(&m_mutexIn);

			av_image_fill_arrays(pFrame->data, pFrame->linesize, dataIn.GetData(), (enum AVPixelFormat)pFrame->format, pFrame->width, pFrame->height, MEM_ALIGN);

			if ((pFrame->width == m_pFrame->width) && (pFrame->height == m_pFrame->height))
			{
				uint8_t* pSrc;
				uint8_t* pDst;

				for (int y=0; y<pFrame->height; ++y)
				{
					pSrc = &pFrame->data[0][y*pFrame->linesize[0]];
					pDst = &m_pFrame->data[0][y*m_pFrame->linesize[0]];
					memcpy(pDst, pSrc, pFrame->linesize[0]);
				}

				for (int y=0; y<(pFrame->height/2); ++y)
				{
					pSrc = &pFrame->data[1][y*pFrame->linesize[1]];
					pDst = &m_pFrame->data[1][y*m_pFrame->linesize[1]];
					memcpy(pDst, pSrc, pFrame->linesize[1]);
				}
			}
			else
			{
				sws_scale(m_pSwsCtx, pFrame->data, pFrame->linesize, 0, pFrame->height, m_pFrame->data, m_pFrame->linesize);
			}
		}

		m_pFrame->pts = m_nFrameCnt;
		m_pFrame->display_picture_number = static_cast<int>(m_nFrameCnt % 0x7FFFFFFF);

		clkCur = std::clock();
		if ((clkCur - clkPre + clkIntvl) < m_lFrameIntvl)
		{
			msec = static_cast<long>((1000 / (double)CLOCKS_PER_SEC) * (clkCur - clkPre + clkIntvl));
			if ((msec > 0) && (msec < m_lFrameIntvl))
				Sleep(msec);
		}
		clkPre = clkCur;

		rval = avcodec_send_frame(m_pCodecCtx, m_pFrame);
		if ((rval != 0) && (rval != AVERROR(EAGAIN)))
		{
			TRACE("Sending frame for encoding error!\n");
			av_packet_free(&pPacket);
			av_frame_free(&pFrame);
			return false;
		}

		++m_nFrameCnt;

		rval = 0;
		while (rval >= 0)
		{
			rval = avcodec_receive_packet(m_pCodecCtx, pPacket);
			if (rval != 0)
			{
				if (rval != AVERROR(EAGAIN))
					TRACE("Receiving encoded packet failed!\n");

				break;
			}

			SendPacket(pPacket);
		}
	}

	av_packet_free(&pPacket);
	av_frame_free(&pFrame);

	return true;
}

void CEncoder::ReceiveFrame(CDataPacket* pData)
{
	EnterCriticalSection(&m_mutexIn);
	m_dataIn.Swap(*pData);
	m_bInState = true;
	LeaveCriticalSection(&m_mutexIn);
}

void CEncoder::SendPacket(AVPacket* pPacket)
{
	CDataPacket dataOut;
	uint8_t *p = pPacket->data;
	int size = pPacket->size;

	if ((p[0] == 0x00) && (p[1] == 0x00))
	{
		if (p[2] == 0x01)
		{
			p += 3;
			size -= 3;
		}
		else if ((p[2] == 0x00) && (p[3] == 0x01))
		{
			p += 4;
			size -= 4;
		}
	}

	if (size <= 0)
		return;

	dataOut.Copy(p, size);

	EnterCriticalSection(&m_mutexOut);
	m_qOut[m_qOutPos.Push()].Swap(dataOut);
	LeaveCriticalSection(&m_mutexOut);

	if (m_cbFunc != nullptr)
		m_cbFunc();
}

bool CEncoder::FetchPacket(uint8_t** ppFrameBuffer, uint32_t* pFrameSize)
{
	if (!m_qOutPos.IsEmpty())
	{
		EnterCriticalSection(&m_mutexOut);
		size_t nPos = m_qOutPos.Pop();
		if (nPos != (size_t)-1)
			m_dataOut.Swap(m_qOut[nPos]);
		LeaveCriticalSection(&m_mutexOut);

		++m_nOutCnt;

		*ppFrameBuffer = m_dataOut.GetData();
		*pFrameSize = m_dataOut.GetSize();

		return true;
	}
	else
	{
		*ppFrameBuffer = NULL;
		*pFrameSize = 0;

		return false;
	}
}

const uint8_t* CEncoder::FindNextNAL(const uint8_t* pStart, const uint8_t* pEnd)
{
	const uint8_t *p = pStart;

	// Simply lookup "0x000001" pattern
	while (p <= (pEnd - 3))
	{
		if (p[0] == 0x00)
		{
			if (p[1] == 0x00)
			{
				if ((p[2] == 0x01) || ((p[2] == 0x00) && (p[3] == 0x01)))
					return p;
				else
					p += 3;
			}
			else
			{
				p += 2;
			}
		}
		else
		{
			++p;
		}
	}

	return pEnd;
}

static CCapture* s_pCapture;
static CEncoder* s_pEncoder;

static void CaptureOnReceiveFrame(CDataPacket *pData)
{
	if (s_pEncoder->IsReady())
		s_pEncoder->ReceiveFrame(pData);
}

static UINT CaptureThreadProc(LPVOID lParam)
{
	s_pCapture->Run();

	return 0;
}

static UINT EncoderThreadProc(LPVOID lParam)
{
	s_pEncoder->Run();

	return 0;
}

#include "UsageEnvironment.hh"
#include "BasicUsageEnvironment.hh"
#include "GroupsockHelper.hh"
#include "liveMedia.hh"

#pragma comment(lib, "libcrypto.lib")
#pragma comment(lib, "libssl.lib")
#pragma comment(lib, "BasicUsageEnvironment.lib")
#pragma comment(lib, "groupsock.lib")
#pragma comment(lib, "liveMedia.lib")
#pragma comment(lib, "UsageEnvironment.lib")

class CaptureFrameSource : public FramedSource
{
public:
	static CaptureFrameSource* createNew(UsageEnvironment& env, CEncoder*encoder);

	CaptureFrameSource(UsageEnvironment& env, CEncoder* encoder);
	~CaptureFrameSource() = default;

private:
	static void deliverFrameStub(void* clientData);

	void doGetNextFrame() override;
	void deliverFrame();
	void doStopGettingFrames() override;
	void onFrame();

private:
	CEncoder *fEncoder;
	EventTriggerId fEventTriggerId;
};

CaptureFrameSource* CaptureFrameSource::createNew(UsageEnvironment& env, CEncoder *encoder) 
{
	return new CaptureFrameSource(env, encoder);
}

CaptureFrameSource::CaptureFrameSource(UsageEnvironment& env, CEncoder* encoder)
	: FramedSource(env), fEncoder(encoder)
{
	fEventTriggerId = envir().taskScheduler().createEventTrigger(CaptureFrameSource::deliverFrameStub);
	fEncoder->m_cbFunc = std::bind(&CaptureFrameSource::onFrame, this);
}

void CaptureFrameSource::deliverFrameStub(void* clientData)
{
	((CaptureFrameSource*)clientData)->deliverFrame();
}

void CaptureFrameSource::doStopGettingFrames()
{
	FramedSource::doStopGettingFrames();
}

void CaptureFrameSource::onFrame()
{
	envir().taskScheduler().triggerEvent(fEventTriggerId, this);
}

void CaptureFrameSource::doGetNextFrame()
{
	deliverFrame();
}

void CaptureFrameSource::deliverFrame()
{
	if (!isCurrentlyAwaitingData())
		return;

	uint8_t* newFrameDataStart;
	unsigned newFrameSize = 0;

	if (fEncoder->FetchPacket(&newFrameDataStart, &newFrameSize))
	{
		if (newFrameDataStart != NULL)
		{
			if (newFrameSize > fMaxSize)
			{
				fFrameSize = fMaxSize;
				fNumTruncatedBytes = newFrameSize - fMaxSize;
			}
			else
			{
				fFrameSize = newFrameSize;
			}

			gettimeofday(&fPresentationTime, NULL);
			memcpy(fTo, newFrameDataStart, fFrameSize);
		}
		else
		{
			fFrameSize = 0;
			fTo = NULL;
			handleClosure(this);
		}
	}
	else
	{
		fFrameSize = 0;
	}

	if (fFrameSize > 0)
		FramedSource::afterGetting(this);
}

class CaptureMediaSubsession : public OnDemandServerMediaSubsession
{
public:
	static CaptureMediaSubsession* createNew(UsageEnvironment& env, StreamReplicator* replicator);

	void setPPSNAL(const uint8_t* ptr, int size);
	void setSPSNAL(const uint8_t* ptr, int size);
	void setBitRate(uint64_t br);

protected:
	CaptureMediaSubsession(UsageEnvironment& env, StreamReplicator* replicator);
	FramedSource* createNewStreamSource(unsigned clientSessionId, unsigned& estBitrate) override;
	RTPSink* createNewRTPSink(Groupsock* rtpGroupsock, unsigned char rtpPayloadTypeIfDynamic, FramedSource* inputSource) override;

	StreamReplicator *fReplicator;
	unsigned long fBitRate;
	const uint8_t* fPPSNAL;
	int fPPSNALSize;
	const uint8_t* fSPSNAL;
	int fSPSNALSize;
};

CaptureMediaSubsession* CaptureMediaSubsession::createNew(UsageEnvironment& env, StreamReplicator* replicator)
{
	return new CaptureMediaSubsession(env, replicator);
}

CaptureMediaSubsession::CaptureMediaSubsession(UsageEnvironment& env, StreamReplicator* replicator)
	: OnDemandServerMediaSubsession(env, True), fReplicator(replicator), fBitRate(1024), fPPSNAL(NULL), fPPSNALSize(0), fSPSNAL(NULL), fSPSNALSize(0)
{
}

FramedSource* CaptureMediaSubsession::createNewStreamSource(unsigned clientSessionId, unsigned& estBitrate)
{
	estBitrate = fBitRate;
	FramedSource* source = fReplicator->createStreamReplica();
	return H264VideoStreamDiscreteFramer::createNew(envir(), source);
}

RTPSink* CaptureMediaSubsession::createNewRTPSink(Groupsock* rtpGroupsock, unsigned char rtpPayloadTypeIfDynamic, FramedSource* inputSource)
{
	return H264VideoRTPSink::createNew(envir(), rtpGroupsock, rtpPayloadTypeIfDynamic, fSPSNAL, fSPSNALSize, fPPSNAL, fPPSNALSize);
}

void CaptureMediaSubsession::setPPSNAL(const uint8_t* ptr, int size)
{
	fPPSNAL = ptr;
	fPPSNALSize = size;
}

void CaptureMediaSubsession::setSPSNAL(const uint8_t* ptr, int size)
{
	fSPSNAL = ptr;
	fSPSNALSize = size;
}

void CaptureMediaSubsession::setBitRate(uint64_t br)
{
	if (br > 102400)
	{
		fBitRate = static_cast<unsigned long>(br / 1024);
	}
	else
	{
		TRACE("Estimated bit rate %lld too small, set to 100K!\n", br);
		fBitRate = 100;
	}
}

static char eventLoopWatchVariable = 0;

static UINT RtspThreadProc(LPVOID lParam)
{
	TaskScheduler* scheduler = BasicTaskScheduler::createNew();
	UsageEnvironment* env = BasicUsageEnvironment::createNew(*scheduler);

	ASSERT(scheduler);
	ASSERT(env);

	OutPacketBuffer::increaseMaxSizeTo(5242880);

	RTSPServer* rtspServer = RTSPServer::createNew(*env, 554);
	if (rtspServer == NULL)
	{
		*env << "Failed to create RTSP server: " << env->getResultMsg() << "\n";
		return -1;
	}
	else
	{
		rtspServer->setUpTunnelingOverHTTP(8080);

		CaptureFrameSource* source = CaptureFrameSource::createNew(*env, s_pEncoder);
		StreamReplicator* replicator = StreamReplicator::createNew(*env, source, false);

		ServerMediaSession* sms = ServerMediaSession::createNew(*env, "stream1");
		CaptureMediaSubsession* sub = CaptureMediaSubsession::createNew(*env, replicator);

		sub->setSPSNAL(s_pEncoder->GetSPSNAL(), s_pEncoder->GetSPSNALSize());
		sub->setPPSNAL(s_pEncoder->GetPPSNAL(), s_pEncoder->GetPPSNALSize());
		if (s_pEncoder->GetBitRate() > 102400)
			sub->setBitRate(s_pEncoder->GetBitRate());
		else
			sub->setBitRate(static_cast<uint64_t>(s_pEncoder->GetWidth()*s_pEncoder->GetHeight()*s_pEncoder->GetFps()/20));

		sms->addSubsession(sub);
		rtspServer->addServerMediaSession(sms);

		char* url = rtspServer->rtspURL(sms);
		*env << "Play this stream using the URL \"" << url << "\"\n";
		delete[] url;

		env->taskScheduler().doEventLoop(&eventLoopWatchVariable);

		Medium::close(rtspServer);
		Medium::close(replicator);
	}

	env->reclaim();
	delete scheduler;

	return 0;
}

void RtspServerInit()
{
	av_register_all();
	avcodec_register_all();
	avdevice_register_all();
	av_log_set_level(AV_LOG_ERROR);

	s_pEncoder = new CEncoder;
	ASSERT(s_pEncoder);
	s_pEncoder->Setup();

	s_pCapture = new CCapture;
	ASSERT(s_pCapture);
	s_pCapture->m_cbFunc = CaptureOnReceiveFrame;
	s_pCapture->Open();

	if (s_pCapture->IsReady())
		AfxBeginThread((AFX_THREADPROC)CaptureThreadProc, (LPVOID)NULL);

	if (s_pEncoder->IsReady())
		AfxBeginThread((AFX_THREADPROC)EncoderThreadProc, (LPVOID)NULL);

	AfxBeginThread((AFX_THREADPROC)RtspThreadProc, (LPVOID)NULL);
}

// The one and only application object

CWinApp theApp;

using namespace std;

int main()
{
	int nRetCode = 0;

	HMODULE hModule = ::GetModuleHandle(nullptr);

	if (hModule != nullptr)
	{
		// initialize MFC and print and error on failure
		if (!AfxWinInit(hModule, nullptr, ::GetCommandLine(), 0))
		{
			// TODO: code your application's behavior here.
			wprintf(L"Fatal Error: MFC initialization failed\n");
			nRetCode = 1;
		}
		else
		{
			// TODO: code your application's behavior here.
			wprintf(L"Open source license\n");
			wprintf(L"FFMPEG Windows Builds: GPL v3\n");
			wprintf(L"live555: LGPL\n\n");

			RtspServerInit();

			while (1)
			{
				Sleep(100);
			}
		}
	}
	else
	{
		// TODO: change error code to suit your needs
		wprintf(L"Fatal Error: GetModuleHandle failed\n");
		nRetCode = 1;
	}

	return nRetCode;
}
