#include "FaceCache.h"
#include "PipeLine.h"
#include "MaterialBuffer.h"
#include "PL_BlockGrouping.h"
#include <logger.h>
#include <Logger/src/logger.hpp>

#include <opencv/cv.h>
#include <opencv/cxcore.h>
#include <opencv/highgui.h>

#include <vector>

#include <PbFaceList.pb.h>

#include <libyuv.h>
#include <MediaHelper.h>

struct OutRGBBufferWrapper
{
	uint8_t* buffer;
	size_t buffSize;

	OutRGBBufferWrapper() : buffer(nullptr), buffSize(0)
	{ }

	OutRGBBufferWrapper(uint8_t* _buffer, size_t _buffSize) : buffer(_buffer), buffSize(_buffSize)
	{ }
};

struct FcPmBreackerContextPLBG
{
	plbg_output_vec_t lastOutputs;
	std::vector<OutRGBBufferWrapper> lastOutputsRGBBufs;

	uint8_t frameYUV[1920*1080*4];
	size_t frameYUVSize;
	uint8_t frameRGB[1920*1080*4];
	size_t frameRGBSize;
	bool frameTruncated;

	typedef std::map<int, float> block_score_filter_map_t;
	block_score_filter_map_t blockScoreFilter;

	FcPmBreackerContextPLBG() :
			lastOutputs(), lastOutputsRGBBufs(), frameYUV(), frameYUVSize(0), frameRGB(),frameRGBSize(0),
			frameTruncated(false),
			blockScoreFilter()
	{}

	void reset()
	{
		lastOutputs.clear();
		lastOutputsRGBBufs.clear();
		frameYUVSize = 0;
		frameRGBSize = 0;
		frameTruncated = false;
	}

	void copy_data()
	{
		uint8_t* bufPtr = frameYUV;
		for (plbg_output_vec_t::iterator iter = lastOutputs.begin(); iter != lastOutputs.end(); ++iter)
		{
			const uint8_t* originBuf = iter->croppedData;
			if (originBuf == nullptr || iter->croppedDataSize == 0)
				continue;

			if (sizeof(frameYUV) - (bufPtr - frameYUV) < iter->croppedDataSize)
			{
				frameTruncated = true;
				LOG_WARN << "yuv buffer truncated" << LOG_ENDL;
				break;
			}

			iter->croppedData = bufPtr;
			memcpy(bufPtr, originBuf, iter->croppedDataSize);
			bufPtr += iter->croppedDataSize;
			frameYUVSize += iter->croppedDataSize;
		}
	}

	void convert_data()
	{
		uint8_t* bufPtr = frameRGB;
		for (plbg_output_vec_t::iterator iter = lastOutputs.begin(); iter != lastOutputs.end(); ++iter)
		{
			if (iter->croppedData == nullptr || iter->croppedDataSize == 0)
			{
				OutRGBBufferWrapper bw;
				lastOutputsRGBBufs.push_back(bw);
				continue;
			}

			OutRGBBufferWrapper bw;
			bw.buffer = bufPtr;
			bw.buffSize = iter->rectInOriginFrame.rect.width() * iter->rectInOriginFrame.rect.height() * 2; //RGB565
			if (sizeof(frameRGB) - (bufPtr - frameRGB) < bw.buffSize)
			{
				frameTruncated = true;
				OutRGBBufferWrapper bw;
				lastOutputsRGBBufs.push_back(bw);
				LOG_WARN << "rgb buffer truncated" << LOG_ENDL;
				continue;
			}

			convertYUV420ToRGB565(*iter, bw.buffer);

			lastOutputsRGBBufs.push_back(bw);
			bufPtr += bw.buffSize;
			frameRGBSize += bw.buffSize;
		}
	}

private:
	bool convertYUV420ToRGB565(const PLBG_Output& out, uint8_t* rgbBuff)
	{
		int src_height = out.rectInOriginFrame.rect.height();
		int src_width = out.rectInOriginFrame.rect.width();
		const uint8* src_y = (const uint8*)(out.croppedData);
		const uint8* src_u = (const uint8*)(src_y + (src_height * src_width));
		const uint8* src_v = (const uint8*)(src_u + (src_height * src_width / 4));

		if (out.originframe->type == MB_Frame::MBFT_YUV420)
		{
			libyuv::I420ToRGB565(src_y, src_width,
								 src_v, MH_SUBSAMPLE1(src_width, 2),
								 src_u, MH_SUBSAMPLE1(src_width, 2),
								 rgbBuff, 2 * src_width,
								 src_width, src_height);
		}
		else if (out.originframe->type == MB_Frame::MBFT_NV12)
		{
			libyuv::NV12ToRGB565(src_y, src_width,
								 src_u, src_width,
								 rgbBuff, 2 * src_width,
								 src_width, src_height);
		}
		else
			return false;

		//{
		//	static size_t f = 0;
		//	char fname[50];
		//	sprintf(fname, "/sdcard/face-%u-w%u-h%u.rgb", ++f, src_width, src_height);
		//	FILE *pFile = fopen(fname, "wb");
		//	fwrite(rgbBuff, 1, src_height * src_width * 2, pFile);
		//	fclose(pFile);
		//	if (f > 20)exit(0);
		//}
		return true;
	}
};

FaceCacheForPLBG::FaceCacheForPLBG() : _ctx(new FcPmBreackerContextPLBG), resultBuffer(nullptr)
{
	resultBuffer = new uint8_t[FACECACHEFORPLBG_RESULT_BUFFSIZE];
}

FaceCacheForPLBG::~FaceCacheForPLBG()
{
	delete (FcPmBreackerContextPLBG*)_ctx;
	delete resultBuffer;
	resultBuffer = nullptr;
}


// returns count of face
int FaceCacheForPLBG::cachePm(const PipeMaterial& pm)
{
	FcPmBreackerContextPLBG& ctx(*(FcPmBreackerContextPLBG*)_ctx);
	ctx.reset();

	plbg_output_vec_t* plbgOut = (plbg_output_vec_t*)pm.buffer;
	if (plbgOut->size() == 0)
		return 0;

	ctx.lastOutputs = *plbgOut;
	ctx.copy_data();
	ctx.convert_data();

	return  ctx.lastOutputs.size();
}

size_t FaceCacheForPLBG::getFaceCount(const PipeMaterial& pm) const
{
	FcPmBreackerContextPLBG& ctx(*(FcPmBreackerContextPLBG*)_ctx);
	plbg_output_vec_t* plbgOut = (plbg_output_vec_t*)pm.buffer;

	bool hasFiltedBlock = false;
	int min_st_track_id = 0;

	for (plbg_output_vec_t::iterator iterBGOut = plbgOut->begin(); iterBGOut != plbgOut->end(); ++iterBGOut)
	{
		int st_track_id = (int)(int64_t)(iterBGOut->rectInOriginFrame.userData);
		float score = iterBGOut->score;

		FcPmBreackerContextPLBG::block_score_filter_map_t::iterator iterFilter = ctx.blockScoreFilter.find(st_track_id);
		if (iterFilter == ctx.blockScoreFilter.end())
		{
			ctx.blockScoreFilter.insert(std::make_pair(st_track_id, score));
			hasFiltedBlock = true;
		}
		else
		{
			if (iterFilter->second < score)
			{
				iterFilter->second = score;
				hasFiltedBlock = true;
			}
		}

		if (min_st_track_id > st_track_id)
			min_st_track_id = st_track_id;
	}

	for(FcPmBreackerContextPLBG::block_score_filter_map_t::iterator iterFilter = ctx.blockScoreFilter.begin(); iterFilter != ctx.blockScoreFilter.end(); )
	{
		if (iterFilter->first < min_st_track_id)
			iterFilter = ctx.blockScoreFilter.erase(iterFilter);
		else
			++iterFilter;
	}

	if (hasFiltedBlock)
		return plbgOut->size();
	else
		return 0;
}

bool FaceCacheForPLBG::getFaceListPb(uint8_t* buffer, size_t& buffMaxSize)
{
	FcPmBreackerContextPLBG& ctx(*(FcPmBreackerContextPLBG*)_ctx);
	if (ctx.frameTruncated)
	{
		LOG_WARN << "ctx buffer truncated" << LOG_ENDL;
		return false;
	}

	PbFaceList pbFaceList;
	pbFaceList.set_magic(pbFaceList.magic());
	pbFaceList.set_image_count(ctx.lastOutputs.size());
	pbFaceList.set_src_width(0);
	pbFaceList.set_src_height(0);

	for (int i = 0; i < ctx.lastOutputs.size(); i++)
	{
		PbFaceList_FaceListImage &pbFaceListImage(*(pbFaceList.add_images()));
		pbFaceListImage.set_idx(i);
		pbFaceListImage.set_st_track_id((int)(int64_t)(ctx.lastOutputs[i].rectInOriginFrame.userData));
		pbFaceListImage.set_size(ctx.lastOutputs[i].croppedDataSize);
		pbFaceListImage.set_type(PbFaceList_FaceListImage_ImageType_MBFT_NV12);
		pbFaceListImage.set_width(ctx.lastOutputs[i].rectInOriginFrame.rect.width());
		pbFaceListImage.set_height(ctx.lastOutputs[i].rectInOriginFrame.rect.height());
		pbFaceListImage.set_top_left_x(ctx.lastOutputs[i].rectInOriginFrame.rect.leftTop.X);
		pbFaceListImage.set_top_left_y(ctx.lastOutputs[i].rectInOriginFrame.rect.leftTop.Y);

		pbFaceListImage.add_img(ctx.lastOutputs[i].croppedData, ctx.lastOutputs[i].croppedDataSize);

		//LOGP(ERROR, "bg output=%f", ctx.lastOutputs[i].score);

		//{
		//	static size_t f = 0;
		//	char fname[50];
		//	sprintf(fname, "/sdcard/face-%u-w%u-h%u.yuv420", ++f, ctx.lastOutputs[i].rectInOriginFrame.width(), ctx.lastOutputs[i].rectInOriginFrame.height());
		//	FILE *pFile = fopen(fname, "wb");
		//	fwrite(ctx.lastOutputs[i].croppedData, 1, ctx.lastOutputs[i].croppedDataSize, pFile);
		//	fclose(pFile);
		//	if (f > 10)exit(0);
		//}
	}

	size_t s = pbFaceList.ByteSize();
	buffMaxSize = std::min(s, buffMaxSize);
	pbFaceList.SerializeToArray(buffer, buffMaxSize);

	return buffMaxSize > 0;
}

bool FaceCacheForPLBG::getFaceListImage(std::vector<NativeImgIdx>& imgIdxes, uint8_t* buffImg, size_t& buffImgMaxSize)
{
	FcPmBreackerContextPLBG& ctx(*(FcPmBreackerContextPLBG*)_ctx);
	if (ctx.frameTruncated)
	{
		LOG_WARN << "ctx buffer truncated" << LOG_ENDL;
		return false;
	}
	if (buffImgMaxSize < ctx.frameRGBSize)
	{
		LOG_WARN << "buffImg truncated" << LOG_ENDL;
		return false;
	}

	imgIdxes.clear();
	size_t totalSize = 0;

	for (int i = 0; i < ctx.lastOutputs.size(); i++)
	{
		NativeImgIdx imgidx;
		imgidx.st_track_id =  (int)(int64_t)(ctx.lastOutputs[i].rectInOriginFrame.userData);
		imgidx.offset = ctx.lastOutputsRGBBufs[i].buffer - ctx.frameRGB;
		imgidx.size = ctx.lastOutputsRGBBufs[i].buffSize;
		imgidx.type = MB_Frame::MBFT_RGB565;
		imgidx.width = ctx.lastOutputs[i].rectInOriginFrame.rect.width();
		imgidx.height = ctx.lastOutputs[i].rectInOriginFrame.rect.height();

		imgIdxes.push_back(imgidx);
		totalSize += ctx.lastOutputsRGBBufs[i].buffSize;
	}

	memcpy(buffImg, ctx.frameRGB, totalSize);
	buffImgMaxSize = totalSize;

	return  true;
}