#include "VideoCapture.hpp"
#include <QImage>
#include <QtDebug>
#include <QTimer>

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

struct VideoCapturePrivate {
	AVFormatContext* fcontext = nullptr;
	AVCodecContext* ccontext = nullptr;
	SwsContext* scontext = nullptr;
	int video_stream_index = -1;
	AVFrame* pic = nullptr;
	AVFrame* picrgb = nullptr;
	QTimer timer;
};

///////////////////////////////////////////////////////////////////////////////

VideoCapture::VideoCapture()
	: d(new VideoCapturePrivate)
{
	av_register_all();
	avformat_network_init();
	d->timer.setSingleShot(true);
	connect(&d->timer, SIGNAL(timeout()), this, SLOT(takeSnapshot()));
}

VideoCapture::~VideoCapture()
{
	delete d;
}

bool VideoCapture::start(const QString &address)
{
	// Open RTSP
	d->fcontext = avformat_alloc_context();
	if(avformat_open_input(&d->fcontext, address.toAscii(), NULL, NULL) != 0){
		qCritical() << "Failed to open RTSP input" << endl;
		return false;
	}

	// Search video stream
	if(avformat_find_stream_info(d->fcontext, NULL) < 0){
		qCritical() << "Failed to find stream info" << endl;
		return false;
	}
	d->video_stream_index = -1;
	for(unsigned i = 0; i < d->fcontext->nb_streams; ++i){
		if(d->fcontext->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO){
			d->video_stream_index = i;
			break;
		}
	}
	if(d->video_stream_index == -1){
		qCritical() << "Failed to find video stream" << endl;
		return false;
	}

	// Load codec
	AVCodec* codec = avcodec_find_decoder(CODEC_ID_H264);
	d->ccontext = avcodec_alloc_context3(codec);
	avcodec_copy_context(d->ccontext, d->fcontext->streams[d->video_stream_index]->codec);
	if(avcodec_open2(d->ccontext, codec, NULL) < 0){
		qCritical() << "Failed to open codec" << endl;
		return false;
	}

	// SwSContext and frame for conversion to RGB24
	d->scontext = sws_getContext(d->ccontext->width, d->ccontext->height, d->ccontext->pix_fmt, d->ccontext->width, d->ccontext->height, PIX_FMT_RGB24, SWS_BICUBIC, NULL, NULL, NULL);
	d->picrgb = avcodec_alloc_frame();
	avpicture_alloc((AVPicture*) d->picrgb, PIX_FMT_RGB24, d->ccontext->width, d->ccontext->height);

	// Alloc frame for decoder
	d->pic = avcodec_alloc_frame();

	d->timer.start(1000);
	return true;
}

void VideoCapture::stop()
{
	d->timer.stop();
	avcodec_free_frame(&d->pic);
	avpicture_free((AVPicture*) d->picrgb);
	avcodec_free_frame(&d->picrgb);

	avcodec_close(d->ccontext);
	av_free(d->ccontext);
	avformat_free_context(d->fcontext);
	d->ccontext = nullptr;
	d->fcontext = nullptr;
	d->scontext = nullptr;
}

void VideoCapture::takeSnapshot()
{
	AVPacket packet;
	av_read_play(d->fcontext);
	while(av_read_frame(d->fcontext, &packet) == 0){
		if(packet.stream_index == d->video_stream_index){
			int check = 0;
			if(avcodec_decode_video2(d->ccontext, d->pic, &check, &packet) > 0 && check != 0){
				sws_scale(d->scontext, d->pic->data, d->pic->linesize, 0, d->ccontext->height, d->picrgb->data, d->picrgb->linesize);
				QImage image(d->picrgb->data[0], d->ccontext->width, d->ccontext->height, d->picrgb->linesize[0], QImage::Format_RGB888);
				image.save("capture.png");
				av_free_packet(&packet);
				break;
			}
		}
		av_free_packet(&packet);
	}

	av_read_pause(d->fcontext);
	d->timer.start(1000);
}