#include "MediaFile.h"
#include "ffcpp.h"
#include <stdexcept>

namespace ffcpp {

	MediaFile::MediaFile(const std::string& src, Mode mode) {
		_mode = mode;
		if(mode == Mode::Read) {
			int res = avformat_open_input(&_formatCtx, src.c_str(), nullptr, nullptr);
			throwIfError(res, "cannot open input file");

			res = avformat_find_stream_info(_formatCtx, nullptr);
			throwIfError(res, "cannot find stream info");
		} else if(mode == Mode::Write) {
			int res = avformat_alloc_output_context2(&_formatCtx, nullptr, nullptr, src.c_str());
			throwIfError(res, "cannot allocate format context");

			if (!(_formatCtx->oformat->flags & AVFMT_NOFILE)) {
				res = avio_open(&_formatCtx->pb, src.c_str(), AVIO_FLAG_WRITE);
				throwIfError(res, "cannot init avio context");
			}

			if(_formatCtx->oformat->flags & AVFMT_GLOBALHEADER)
				_formatCtx->flags |= CODEC_FLAG_GLOBAL_HEADER;

			av_dump_format(_formatCtx, 0, src.c_str(), 1);
		}
	}

	bool MediaFile::hasStream(AVMediaType type) const {
		for(size_t i = 0; i < _formatCtx->nb_streams; ++i) {
			if(_formatCtx->streams[i]->codec->codec_type == type) {
				return true;
			}
		}

		return false;
	}

	bool MediaFile::hasVideo() const {
		return hasStream(AVMEDIA_TYPE_VIDEO);
	}

	bool MediaFile::hasAudio() const {
		return hasStream(AVMEDIA_TYPE_AUDIO);
	}

	Stream MediaFile::getStream(AVMediaType type, size_t index) const {
		for(size_t i = 0, curIndex = 0; i < _formatCtx->nb_streams; ++i) {
			if(_formatCtx->streams[i]->codec->codec_type == type) {
				if(curIndex == index) {
					return Stream(_formatCtx->streams[i]);
				} else {
					curIndex++;
				}
			}
		}

		throw std::runtime_error("cannot find stream");
	}

	Stream MediaFile::videoStream(size_t index /* = 0 */) const {
		return getStream(AVMEDIA_TYPE_VIDEO, index);
	}

	Stream MediaFile::audioStream(size_t index /* = 0 */) const {
		return getStream(AVMEDIA_TYPE_AUDIO, index);
	}

	MediaFile::operator AVFormatContext*() const {
		return _formatCtx;
	}

	MediaFile::~MediaFile() {
		if(_mode == Mode::Write) {
			avformat_free_context(_formatCtx);
		} else {
			avformat_close_input(&_formatCtx);
		}
	}

	Stream MediaFile::addStream(AVCodecID codecID, int width, int height, AVPixelFormat pixelFormat) {
		AVCodec* codec = avcodec_find_encoder(codecID);
		if(!codec) throw std::runtime_error("cannot find codec");

		AVStream* stream = avformat_new_stream(_formatCtx, codec);
		if(!stream) throw  std::runtime_error("cannot create stream");

		AVCodecContext* ctx = stream->codec;
		ctx->width = width;
		ctx->height = height;
		ctx->pix_fmt = pixelFormat;

		return Stream(stream, codec);
	}

	void MediaFile::writeHeader() {
		int res = avformat_write_header(_formatCtx, nullptr);
		throwIfError(res, "error writing header");
	}

	void MediaFile::writeTrailer() {
		int res = av_write_trailer(_formatCtx);
		throwIfError(res, "error writing trailer");
	}
}