﻿#pragma warning (disable:26812)
#include <memory>
#include <iostream>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
}

#include "../libavcv/libavcv.h"

using std::cout;
using std::cerr;
using std::endl;

AVRational ffmpeg_load_time_base(AVCodecContext* avcodecctx, AVFormatContext* avfmtctx, const AVStream* avstream) {
	//avfmtctx->pb->pos == 先頭の時のみに呼び出し可能（但し先頭は0ではない）．
	//if (avfmtctx->pb->pos != 0)
	//	throw "avfmtctx->pb->pos != 0";
	if (avstream->first_dts != avstream->cur_dts)
		throw "Must call this function when the current DTS of AVStream is the first.";

	//復帰のために保存．
	//int64_t cur_dts = avstream->cur_dts;
	int64_t cur_pos = avfmtctx->pb->pos;

	//事前に先頭フレームへseek．
	//if (av_seek_frame(avfmtctx, avstream->index, 0, AVSEEK_FLAG_FRAME) < 0)
	//	throw "Failed: av_seek_frame";

	//avcodec_flush_buffers(avcodecctx); //Seek後にフラッシュ

	AVPacket avpacket;
	bool is_one_frame_decoded = false;
	while (is_one_frame_decoded == false && av_read_frame(avfmtctx, &avpacket) == 0) {
		//if (avpacket.stream_index != avstream->index) {
		//	av_packet_unref(&avpacket);
		//	throw "Invalid AVStream"; //この関数はビデオストリーム以外では呼び出してはいけない
		//}

		if (avcodec_send_packet(avcodecctx, &avpacket) != 0)
		{
			av_packet_unref(&avpacket);
			throw "Failed: avcodec_send_packet";
		}

		// avcodec_send_packetの呼出し後に初めて正しいtime_baseが取得できるようになる．
		// このtime_baseがStream内に記録されている為と考えられる．その為，実際のパケットをSendする必要がある（nullptrではNG）．
		// ここでは，安全を取って１フレーム分が復号できるまで，パケットをRead/Sendする．
		AVFrame* frame = av_frame_alloc();
		//取り出せるフレームが存在する間，戻り値0．
		if (avcodec_receive_frame(avcodecctx, frame) == 0)
			is_one_frame_decoded = true; //１フレーム取り出せれば終了．

		av_frame_free(&frame);
		av_packet_unref(&avpacket);
	}

	if (is_one_frame_decoded == false)
		throw "Can NOT get time_base because none of the frames are present.";

	//先頭に戻しておく
	//if (av_seek_frame(avfmtctx, avstream->index, 0, AVSEEK_FLAG_FRAME) < 0)
	//	throw "Failed: av_seek_frame";
	// AVSEEK_FLAG_FRAME|AVSEEK_FLAG_BACKWARD
	if (av_seek_frame(avfmtctx, avstream->index, 0, AVSEEK_FLAG_FRAME) < 0)
		throw "Failed: av_seek_frame";
	avcodec_flush_buffers(avcodecctx); //Seek後にフラッシュ

	//int64_t result_pos0 = avio_seek(avfmtctx->pb, cur_pos, 0);
	//if (result_pos0 != cur_pos)
	//	throw "Failed: avio_seek";

	return avcodecctx->time_base;
}

void test_read_allframes(AVCodecContext* avcodecctx, const AVStream* avstream, AVFormatContext* avfmtctx) {
	cv::String logstr;

	AVPacket avpacket = AVPacket();
	int frame_counter = 0;
	int pkt_counter = 0;

	while (av_read_frame(avfmtctx, &avpacket) == 0) {
		if (avpacket.stream_index != avstream->index) {
			av_packet_unref(&avpacket);
			continue; //対象ビデオストリーム以外は無視
		}
		pkt_counter++;

		logstr = cv::format("AVPacket: count=%3d dts=%6d pts=%6d %s",
			pkt_counter, avpacket.dts, avpacket.pts, avpacket.dts == avpacket.pts ? "" : "(!=)");
		cout << logstr << endl;

		auto avframeVec = ffmpeg_send_packet_receive_frames(avcodecctx, &avpacket);

		for (auto& frame : avframeVec) {
			logstr = cv::format("Info: Frame=%6d pts=%6d dts=%6d pict_type=%c (pts %s dts) CodedPicNo=%4d ",
				frame_counter, frame->pts, frame->pkt_dts, av_get_picture_type_char(frame->pict_type),
				frame->pts == frame->pkt_dts ? "==" : "!=",
				frame->coded_picture_number);
			cout << logstr << endl;
			frame_counter++;

			auto mat = convert_avframe_to_mat(frame.get(), avcodecctx->pix_fmt);
			cv::imshow("Validating", mat);
			cv::waitKey(1);
		}

		av_packet_unref(&avpacket);
	}

	// nullptrを渡して、デコーダをフラッシュする。
	cout << "Flush encoder." << endl;
	auto avframeVec = ffmpeg_send_packet_receive_frames(avcodecctx, nullptr);
	for (auto& frame : avframeVec) {
		logstr = cv::format("Info: Frame=%6d pts=%6d dts=%6d pict_type=%c (pts %s dts) CodedPicNo=%4d ",
			frame_counter, frame->pts, frame->pkt_dts, av_get_picture_type_char(frame->pict_type),
			frame->pts == frame->pkt_dts ? "==" : "!=",
			frame->coded_picture_number);
		cout << logstr << endl;
		frame_counter++;

		auto mat = convert_avframe_to_mat(frame.get(), avcodecctx->pix_fmt);
		cv::imshow("Validating", mat);
		cv::waitKey(1);
	}
}

void test_read_specifidframe_by_frameindex(AVCodecContext* avcodecctx, AVFormatContext* avfmtctx, const AVStream* avstream,
	int64_t frame_count)
{
	cv::String logstr;

	for (int64_t fidx = 0; fidx < frame_count; fidx++) {
	//for (int64_t fidx = frame_count - 1; fidx >= 0; fidx--) {
		int64_t target_idx;
		auto avframeVec = ffmpeg_seek_read_send_receive_frames_by_frameindex(avcodecctx, avfmtctx, avstream, fidx, &target_idx);
		if (target_idx == -1)
			throw cv::format("Failed: There is NOT this frame");

		auto& frame = avframeVec[target_idx];
		logstr = cv::format("Info: SpecifiedFrameIndex=%4d pts=%6d dts=%6d pict_type=%c index=%d(size=%d)",
			fidx, frame->pts, frame->pkt_dts, av_get_picture_type_char(frame->pict_type), target_idx, avframeVec.size());
		cout << logstr << endl;

		auto mat = convert_avframe_to_mat(frame.get(), avcodecctx->pix_fmt);
		cv::imshow("Validating", mat);
		cv::waitKey(1);
	}
}

void test_read_specifidframe_by_pts(AVCodecContext* avcodecctx, AVFormatContext* avfmtctx, const AVStream* avstream,
	const std::vector<int64_t>& pts_list)
{
	cv::String logstr;

	for (int64_t pts : pts_list) {

		int64_t target_idx;
		auto avframeVec = ffmpeg_seek_read_send_receive_frames_by_pts(avcodecctx, avfmtctx, avstream, pts, &target_idx);
		if (target_idx == -1)
			throw cv::format("Failed: Thre is NOT this PTS");

		auto& frame = avframeVec[target_idx];
		logstr = cv::format("Info: pts=%6d dts=%6d pict_type=%c index=%d(size=%d)",
			frame->pts, frame->pkt_dts, av_get_picture_type_char(frame->pict_type), target_idx, avframeVec.size());
		cout << logstr << endl;

		auto mat = convert_avframe_to_mat(frame.get(), avcodecctx->pix_fmt);
		cv::imshow("Validating", mat);
		cv::waitKey(1);
	}
}

int main(int argc, char* argv[])
{
	cv::String logstr;
	if (argc < 2) {
		cerr << "Error: 第一引数にMP4(H.264)ファイルを指定して下さい。" << endl;
		return -1;
	}
	const char* movie_file_path = argv[1];

	//if set nullptr, search a proper one.
	const char* decoder_name = "h264"; //"h264" "h264_qsv" "h264_cuvid"

	AVFormatContext* avfmtctx = nullptr;
	if (avformat_open_input(&avfmtctx, movie_file_path, nullptr, nullptr) != 0) {
		std::cerr << "Failed: avformat_open_input" << std::endl;
		return -1;
	}

	if (avformat_find_stream_info(avfmtctx, nullptr) < 0) {
		std::cerr << "Failed: avformat_find_stream_info" << std::endl;
		return -1;
	}
	logstr = cv::format("Info: avfmtctx->nb_streams=%d", avfmtctx->nb_streams);
	std::cout << logstr << std::endl;

	// ビデオストリームを特定
	AVStream* avstream = nullptr;
	for (int i = 0; i < (int)avfmtctx->nb_streams; ++i) {
		if (avfmtctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
			avstream = avfmtctx->streams[i];
			break;
		}
	}
	if (avstream == nullptr) {
		std::cerr << "Error: Not found any video stream." << std::endl;
		return -1;
	}

	if (avstream->codecpar->codec_id == AVCodecID::AV_CODEC_ID_H264) {
		std::cout << "Info: codec_id = " << avstream->codecpar->codec_id
			<< "(" << avcodec_get_name(avstream->codecpar->codec_id) << ")" << std::endl; //->codec_id = 27(h264)
	}
	else
	{
		std::cout << "Failed: No implementation other than h264." << std::endl;
		return -1;
	}
	AVRational time_base_of_avstream = avstream->time_base;
	AVRational avg_frame_rate_of_avstream = avstream->avg_frame_rate; //avformat_find_stream_infoで取得済み．
	AVRational r_frame_rate_of_avstream = avstream->r_frame_rate;

	AVCodec* avcodec = nullptr;
	if (decoder_name == nullptr)
	{
		avcodec = avcodec_find_decoder(avstream->codecpar->codec_id);
		if (avcodec == nullptr) {
			std::cerr << "Failed: Can NOT find any decoders." << std::endl;
			return -1;
		}
	}
	else
	{
		//Intel QuickSyncVideo(QSV), etc
		avcodec = avcodec_find_decoder_by_name(decoder_name);
		if (avcodec == nullptr) {
			std::cerr << "Failed: Not found the decoder - " << decoder_name << std::endl;
			return -1;
		}
	}

	AVCodecContext* avcodecctx = avcodec_alloc_context3(avcodec);
	if (avcodecctx == nullptr) {
		std::cerr << "Failed: avcodec_alloc_context3" << std::endl;
		return -1;
	}

	//AVStreamのコーデックパラメータをAVCodecContextへセット
	if (avcodec_parameters_to_context(avcodecctx, avstream->codecpar) < 0) {
		std::cerr << "Failed: avcodec_parameters_to_context" << std::endl;
		return -1;
	}

	//AVCodecをオープン
	if (avcodec_open2(avcodecctx, avcodec, nullptr) != 0) {
		std::cerr << "Failed: avcodec_open2" << std::endl;
		return -1;
	}

	AVRational time_base_avcodecctx = ffmpeg_load_time_base(avcodecctx, avfmtctx, avstream);
	auto fps_avcodecctx = avcodecctx->framerate; //[注]アクセスタイミングによって不正

	/* テスト：最初から最後までシーケンシャルにデコードする．*/
	//test_read_allframes(avcodecctx, avstream, avfmtctx);

	/* テスト：任意のフレームへフレームインデクスでシークする．*/
	test_read_specifidframe_by_frameindex(avcodecctx, avfmtctx, avstream, avstream->nb_frames);

	/* テスト：任意のフレームへPTSでシークする．*/
	//std::vector<int64_t> pts_list;
	//for (int64_t nb_frame = 0; nb_frame < avstream->nb_frames; nb_frame++) {
	//	pts_list.push_back(ffmpeg_frameindex_to_pts(avstream, nb_frame));
	//}
	//test_read_specifidframe_by_pts(avcodecctx, avfmtctx, avstream, pts_list);

	avcodec_free_context(&avcodecctx);
	avformat_close_input(&avfmtctx);

	return 0;
}

