/***************************************************
This is the main Grabber code.  It uses AVbin and ffmpeg
to capture video and audio from video and audio files.
Because of this, mmread and supporting code is now
distributed under the LGPL.  See

The code supports any number of audio or video streams and
is a cross platform solution to replace DDGrab.cpp.

This code was intended to be used inside of a matlab interface,
but can be used as a generic grabber class for anyone who needs
one.

Copyright 2008 Micah Richert

This file is part of mmread.

mmread is free software; you can redistribute it and/or modify it
under the terms of the GNU Lesser General Public License as
published by the Free Software Foundation; either version 3 of
the License, or (at your option) any later version.

mmread is distributed WITHOUT ANY WARRANTY.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public
License along with mmread.  If not, see <http://www.gnu.org/licenses/>.
**************************************************/

#ifdef MATLAB_MEX_FILE
#include "mex.h"
#define FFprintf(...) mexPrintf(__VA_ARGS__)
#else
#define FFprintf(...) printf(__VA_ARGS__)
#endif

//#ifndef mwSize
//#define mwSize int
//#endif

#define DEBUG 0

extern "C" {
	#include <avbin.h>
	#include <libavformat/avformat.h>

	struct _AVbinFile {
	    AVFormatContext *context;
	    AVPacket *packet;
	};

	struct _AVbinStream {
		int type;
		AVFormatContext *format_context;
		AVCodecContext *codec_context;
		AVFrame *frame;
	};
}

#include <stdlib.h>
#include <string.h>
#include <vector>
#include <map>
using namespace std;

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>

map<unsigned int,double> keyframes;
unsigned int startDecodingAt;

class Grabber
{
public:
	Grabber(bool isAudio, AVbinStream* stream, bool trySeeking, double rate, int bytesPerWORD, AVbinStreamInfo info, AVbinTimestamp start_time)
	{
		this->stream = stream;
		frameNr = 0;
		packetNr = 0;
		done = false;
		this->bytesPerWORD = bytesPerWORD;
		this->rate = rate;
		startTime = 0;
		stopTime = 0;
		this->isAudio = isAudio;
		this->info = info;
		this->trySeeking = trySeeking;
		this->start_time = start_time>0?start_time:0;
	};

	~Grabber()
	{
		// clean up any remaining memory...
		if (DEBUG) FFprintf("freeing frame data...\n");
		for (vector<uint8_t*>::iterator i=frames.begin();i != frames.end(); i++) free(*i);
	}

	AVbinStream* stream;
	AVbinStreamInfo info;
	AVbinTimestamp start_time;

	vector<uint8_t*> frames;
	vector<unsigned int> frameBytes;
	vector<double> frameTimes;

	vector<unsigned int> frameNrs;

	unsigned int frameNr;
	unsigned int packetNr;
	bool done;
	bool isAudio;
	bool trySeeking;

	int bytesPerWORD;
	double rate;
	double startTime, stopTime;

	int Grab(AVbinPacket* packet)
	{
		if (done) return 0;
		if (!packet->data) return 1;

		frameNr++;
		packetNr++;
		if (DEBUG) FFprintf("frameNr %d %d %d\n",frameNr,packetNr,packet->size);
		int offset=0, len=0;
		double timestamp = (packet->timestamp-start_time)/1000.0/1000.0;
		if (DEBUG) FFprintf("time %lld %lld %lf\n",packet->timestamp,start_time,timestamp);

		// either no frames are specified (capture all), or we have time specified
		if (stopTime)
		{
			if (isAudio)
			{
				// time is being used...
				if (timestamp >= startTime)
				{
					// if we've reached the start...
					offset = max(0,((int)((startTime-timestamp)*rate))*bytesPerWORD);
					len = ((int)((stopTime-timestamp)*rate))*bytesPerWORD;
					// if we have gone past our stop time...

					done = len < 0;
				}
			} else {
				done = stopTime <= timestamp;
				len = (startTime <= timestamp)?0x7FFFFFFF:0;
				if (DEBUG) FFprintf("startTime: %lf, stopTime: %lf, current: %lf, done: %d, len: %d\n",startTime,stopTime,timestamp,done,len);
			}
		} else {
			// capture everything... video or audio
			len = 0x7FFFFFFF;
		}

		if (isAudio)
		{
			if (trySeeking && (len<=0 || done)) return 0;

			uint8_t audiobuf[1024*1024];
			int uint8_tsleft = sizeof(audiobuf);
 			int uint8_tsout = uint8_tsleft;
			int uint8_tsread;
			uint8_t* audiodata = audiobuf;
			if (DEBUG) FFprintf("avbin_decode_audio\n");
			while ((uint8_tsread = avbin_decode_audio(stream, packet->data, packet->size, audiodata, &uint8_tsout)) > 0)
			{
				packet->data += uint8_tsread;
				packet->size -= uint8_tsread;
				audiodata += uint8_tsout;
				uint8_tsleft -= uint8_tsout;
				uint8_tsout = uint8_tsleft;
			}

			int nrBytes = audiodata-audiobuf;
			len = min(len,nrBytes);
			offset = min(offset,nrBytes);

			uint8_t* tmp = (uint8_t*)malloc(len);
			if (!tmp) return 2;

			memcpy(tmp,audiobuf+offset,len);

			frames.push_back(tmp);
			frameBytes.push_back(len);
			frameTimes.push_back(timestamp);

		} else {
			bool skip = false;
			if (frameNrs.size() > 0)
			{
				//frames are being specified
				// check to see if the frame is in our list
				bool foundNr = false;
				unsigned int lastFrameNr = 0;
				for (int i=0;i<frameNrs.size();i++)
				{
					if (frameNrs.at(i) == frameNr) foundNr = true;
					if (frameNrs.at(i) > lastFrameNr) lastFrameNr = frameNrs.at(i);
				}

				done = frameNr > lastFrameNr;
				if (!foundNr) {
					if (DEBUG) FFprintf("Skipping frame %d\n",frameNr);
					skip = true;
				}
			}
			if ((trySeeking && skip && packetNr < startDecodingAt && packetNr != 1) || done ) return 0;

			if (DEBUG) FFprintf("allocate frame %d\n",frames.size());
			uint8_t* videobuf = (uint8_t*)malloc(bytesPerWORD);
			if (!videobuf) return 2;
			if (DEBUG) FFprintf("avbin_decode_video\n");

			if (avbin_decode_video(stream, packet->data, packet->size,videobuf)<=0)
			{
				if (DEBUG) FFprintf("avbin_decode_video FAILED!!!\n");
				// silently ignore decode errors
				frameNr--;
				free(videobuf);
				return 3;
			}

			if (stream->frame->key_frame)
			{
				keyframes[packetNr] = timestamp;
			}

			if (skip || len==0)
			{
				free(videobuf);
				return 0;
			}
			frames.push_back(videobuf);
			frameBytes.push_back(min(len,bytesPerWORD));
			frameTimes.push_back(timestamp);
		}

		return 0;
	}
};

typedef map<int,Grabber*> streammap;

class FFGrabber
{
public:
	FFGrabber();

	int build(char* filename, bool disableVideo, bool disableAudio, bool tryseeking);
	int doCapture();

	int getVideoInfo(unsigned int id, int* width, int* height, double* rate, int* nrFramesCaptured, int* nrFramesTotal, double* totalDuration);
	int getAudioInfo(unsigned int id, int* nrChannels, double* rate, int* bits, int* nrFramesCaptured, int* nrFramesTotal, int* subtype, double* totalDuration);
	void getCaptureInfo(int* nrVideo, int* nrAudio);
	// data must be freed by caller
	int getVideoFrame(unsigned int id, unsigned int frameNr, uint8_t** data, unsigned int* nrBytes, double* time);
	// data must be freed by caller
	int getAudioFrame(unsigned int id, unsigned int frameNr, uint8_t** data, unsigned int* nrBytes, double* time);
	void setFrames(unsigned int* frameNrs, int nrFrames);
	void setTime(double startTime, double stopTime);
	void disableVideo();
	void disableAudio();
	void cleanUp(); // must be called at the end, in order to render anything afterward.

#ifdef MATLAB_MEX_FILE
	void setMatlabCommand(char * matlabCommand);
	void runMatlabCommand(Grabber* G);
#endif
private:
	streammap streams;
	vector<Grabber*> videos;
	vector<Grabber*> audios;

	AVbinFile* file;
	AVbinFileInfo fileinfo;

	bool stopForced;
	bool tryseeking;
	vector<unsigned int> frameNrs;
	double startTime, stopTime;

	char* filename;
	struct stat filestat;


#ifdef MATLAB_MEX_FILE
	char* matlabCommand;
	mxArray* prhs[5];
#endif
};


FFGrabber::FFGrabber()
{
	stopForced = false;
	tryseeking = true;
	file = NULL;
	filename = NULL;

	if (DEBUG) FFprintf("avbin_init\n");
 	if (avbin_init()) FFprintf("avbin_init init failed!!!\n");

	av_log_set_level(AV_LOG_QUIET);
}

void FFGrabber::cleanUp()
{
	if (!file) return; // nothing to cleanup.

	for (streammap::iterator i = streams.begin(); i != streams.end(); i++)
	{
		avbin_close_stream(i->second->stream);
		delete i->second;
	}

 	streams.clear();
 	videos.clear();
 	audios.clear();

 	avbin_close_file(file);
 	file = NULL;

#ifdef MATLAB_MEX_FILE
	if (matlabCommand) free(matlabCommand);
	matlabCommand = NULL;
#endif
}

int FFGrabber::getVideoInfo(unsigned int id, int* width, int* height, double* rate, int* nrFramesCaptured, int* nrFramesTotal, double* totalDuration)
{
	if (!width || !height || !nrFramesCaptured || !nrFramesTotal) return -1;

	if (id >= videos.size()) return -2;
	Grabber* CB = videos.at(id);

	if (!CB) return -1;

	*width  = CB->info.video.width;
	*height = CB->info.video.height;
	*rate = CB->rate;
	*nrFramesCaptured = CB->frames.size();
	*nrFramesTotal = CB->frameNr;

	*totalDuration = fileinfo.duration/1000.0/1000.0;
	if (stopForced) *nrFramesTotal = (int)(-(*rate)*(*totalDuration));

	return 0;
}

int FFGrabber::getAudioInfo(unsigned int id, int* nrChannels, double* rate, int* bits, int* nrFramesCaptured, int* nrFramesTotal, int* subtype, double* totalDuration)
{
	if (!nrChannels || !rate || !bits || !nrFramesCaptured || !nrFramesTotal) return -1;

	if (id >= audios.size()) return -2;
	Grabber* CB = audios.at(id);

	if (!CB) return -1;

	*nrChannels = CB->info.audio.channels;
	*rate = CB->info.audio.sample_rate;
	*bits = CB->info.audio.sample_bits;
	*subtype = CB->info.audio.sample_format;
	*nrFramesCaptured = CB->frames.size();
	*nrFramesTotal = CB->frameNr;

	*totalDuration = fileinfo.duration/1000.0/1000.0;

	return 0;
}

void FFGrabber::getCaptureInfo(int* nrVideo, int* nrAudio)
{
	if (!nrVideo || !nrAudio) return;

	*nrVideo = videos.size();
	*nrAudio = audios.size();
}

// data must be freed by caller
int FFGrabber::getVideoFrame(unsigned int id, unsigned int frameNr, uint8_t** data, unsigned int* nrBytes, double* time)
{
	if (DEBUG) FFprintf("getting Video frame %d\n",frameNr);

	if (!data || !nrBytes) return -1;

	if (id >= videos.size()) return -2;
	Grabber* CB = videos[id];
	if (!CB) return -1;
	if (CB->frameNr == 0) return -2;
	if (frameNr < 0 || frameNr >= CB->frames.size()) return -2;

	uint8_t* tmp = CB->frames[frameNr];
	if (!tmp) return -2;

	*nrBytes = CB->frameBytes[frameNr];
	*time = CB->frameTimes[frameNr];

	*data = tmp;
	CB->frames[frameNr] = NULL;

	return 0;
}

// data must be freed by caller
int FFGrabber::getAudioFrame(unsigned int id, unsigned int frameNr, uint8_t** data, unsigned int* nrBytes, double* time)
{
	if (!data || !nrBytes) return -1;

	if (id >= audios.size()) return -2;
	Grabber* CB = audios[id];
	if (!CB) return -1;
	if (CB->frameNr == 0) return -2;
	if (frameNr < 0 || frameNr >= CB->frames.size()) return -2;

	uint8_t* tmp = CB->frames[frameNr];
	if (!tmp) return -2;

	*nrBytes = CB->frameBytes[frameNr];
	*time = CB->frameTimes[frameNr];

	*data = tmp;
	CB->frames[frameNr] = NULL;

	return 0;
}

void FFGrabber::setFrames(unsigned int* frameNrs, int nrFrames)
{
	if (!frameNrs) return;

	unsigned int minFrame=nrFrames>0?frameNrs[0]:0;

	this->frameNrs.clear();
	for (int i=0; i<nrFrames; i++) this->frameNrs.push_back(frameNrs[i]);

	for (int j=0; j < videos.size(); j++)
	{
		Grabber* CB = videos.at(j);
		if (CB)
		{
			CB->frames.clear();
			CB->frameNrs.clear();
			for (int i=0; i<nrFrames; i++)
			{
				CB->frameNrs.push_back(frameNrs[i]);
				minFrame=frameNrs[i]<minFrame?frameNrs[i]:minFrame;
			}
			CB->frameNr = 0;
			CB->packetNr = 0;
		}
	}

	if (tryseeking && nrFrames > 0)
	{