#include "Util.h"

cv::Mat Util::AVFrameToCVMat(AVFrame* yuv420Frame)
{
	//得到AVFrame信息
	int srcW = yuv420Frame->width;
	int srcH = yuv420Frame->height;
	SwsContext* swsCtx = sws_getContext(
		srcW, srcH,
		(AVPixelFormat)yuv420Frame->format,
		srcW, srcH,
		(AVPixelFormat)AV_PIX_FMT_BGR24,
		SWS_BICUBIC, NULL, NULL, NULL);

	//创建Mat对象
	cv::Mat mat(srcH, srcW, CV_8UC3);

	//创建一个临时BGR24格式的AVFrame
	AVFrame* bgr24Frame = av_frame_alloc();
	//用于分配 AVFrame 的数据内存，以确保内存对齐。
	av_image_alloc(bgr24Frame->data, bgr24Frame->linesize,
		srcW, srcH, AV_PIX_FMT_BGR24, 1);

	// 执行yuv420p转bgr24
	sws_scale(swsCtx,
		(const uint8_t* const*)yuv420Frame->data,
		yuv420Frame->linesize,
		0,
		srcH,
		bgr24Frame->data,
		bgr24Frame->linesize);

	// 将数据拷贝到cv::Mat
	for (int y = 0; y < srcH; ++y) {
		memcpy(mat.ptr(y), bgr24Frame->data[0] + y * bgr24Frame->linesize[0], srcW * 3);
	}

	//释放
	av_freep(&bgr24Frame->data[0]);
	av_frame_free(&bgr24Frame);
	sws_freeContext(swsCtx);

	return mat;
}

AVFrame* Util::CVMatToAVFrame(cv::Mat& inMat)
{
	//得到Mat信息
	AVPixelFormat dstFormat = AV_PIX_FMT_YUV420P;
	int width = inMat.cols;
	int height = inMat.rows;

	//创建AVFrame填充参数 注：调用者释放该frame
	AVFrame* frame = av_frame_alloc();
	frame->width = width;
	frame->height = height;
	frame->format = dstFormat;

	//初始化AVFrame内部空间
	int ret = av_frame_get_buffer(frame, 0);
	if (ret < 0)
	{
		std::cerr << "Could not allocate the video frame data";
		av_frame_free(&frame);
		return nullptr;
	}

	ret = av_frame_make_writable(frame);
	if (ret < 0)
	{
		std::cerr << "Av frame make writable failed.";
		av_frame_free(&frame);
		return nullptr;
	}

	//转换颜色空间为YUV420
	cv::cvtColor(inMat, inMat, cv::COLOR_BGR2YUV_I420);

	// 按YUV420格式，设置数据地址，注意内存对齐
	int frame_width = frame->linesize[0];
	int frame_size = width * height;
	unsigned char* data = inMat.data;
	// 无填充情况
	if (frame_width == width) {
		memcpy(frame->data[0], data, frame_size);
		memcpy(frame->data[1], data + frame_size, frame_size / 4);
		memcpy(frame->data[2], data + frame_size * 5 / 4, frame_size / 4);
	}
	else {
		memset(frame->data[0], 0, frame->linesize[0] * height);
		memset(frame->data[1], 0, frame->linesize[0] * height / 4);
		memset(frame->data[2], 0, frame->linesize[0] * height / 4);
		// Y Panel,Y分量的大小为width*height
		for (int i = 0; i < height; i++) {
			memcpy(frame->data[0] + i * frame->linesize[0], data + i * width, width);
		}
		unsigned char* udata = data + frame_size;
		unsigned char* vdata = udata + frame_size / 4;
		// UV Panel，U和V分量的大小为(width/2)*(height/2)
		for (int i = 0; i < height / 2; i++) {
			memcpy(frame->data[1] + i * frame->linesize[1], udata + i * width / 2, width / 2);
			memcpy(frame->data[2] + i * frame->linesize[2], vdata + i * width / 2, width / 2);
		}
	}

	return frame;
}

AVFrame* Util::CVMatToAVFrame2(cv::Mat& inMat)
{
	int width = inMat.cols;
	int height = inMat.rows;

	// Allocate AVFrame for YUV420P
	AVFrame* frame = av_frame_alloc();

	frame->format = AV_PIX_FMT_YUV420P;
	frame->width = width;
	frame->height = height;

	// Allocate memory for YUV420P data
	int y_plane_size = width * height;
	int uv_plane_size = (width / 2) * (height / 2);
	// 处理字节对齐问题，并分配内存设置行大小，以确保数据对齐满足ffmpeg的要求
	av_image_alloc(frame->data, frame->linesize, width, height, AV_PIX_FMT_YUV420P, 16);

	// Initialize SwsContext for conversion
	SwsContext* sws_ctx = sws_getContext(
		width, height, AV_PIX_FMT_BGR24,
		width, height, AV_PIX_FMT_YUV420P,
		SWS_BILINEAR, NULL, NULL, NULL
	);

	// Perform the conversion
	sws_scale(
		sws_ctx,
		(const uint8_t* const*)&inMat.data,
		(const int*)&inMat.step[0],
		0,
		height,
		frame->data,
		frame->linesize
	);

	// Clean up
	sws_freeContext(sws_ctx);

	return frame;
}