/* See LICENSE file for copyright and license details. */

#include <string.h>

#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/avutil.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>

#include "decode.h"
#include "util.h"

int
dec_open(Decoder *d, const char *path)
{
	const AVCodec *codec;
	AVStream *st;
	int ret;
	int i;

	memset(d, 0, sizeof(*d));
	d->stream_idx = -1;

	ret = avformat_open_input(&d->fmt_ctx, path, NULL, NULL);
	if (ret < 0) {
		warn("cannot open '%s': %s", path, av_err2str(ret));
		return -1;
	}

	ret = avformat_find_stream_info(d->fmt_ctx, NULL);
	if (ret < 0) {
		warn("no stream info in '%s'", path);
		goto fail;
	}

	/* find first video stream */
	for (i = 0; i < (int)d->fmt_ctx->nb_streams; i++) {
		if (d->fmt_ctx->streams[i]->codecpar->codec_type
		    == AVMEDIA_TYPE_VIDEO) {
			d->stream_idx = i;
			break;
		}
	}
	if (d->stream_idx < 0) {
		warn("no video stream in '%s'", path);
		goto fail;
	}

	st = d->fmt_ctx->streams[d->stream_idx];
	codec = avcodec_find_decoder(st->codecpar->codec_id);
	if (!codec) {
		warn("unsupported codec in '%s'", path);
		goto fail;
	}

	d->codec_ctx = avcodec_alloc_context3(codec);
	if (!d->codec_ctx) {
		warn("cannot alloc codec context");
		goto fail;
	}

	ret = avcodec_parameters_to_context(d->codec_ctx,
	                                    st->codecpar);
	if (ret < 0) {
		warn("cannot copy codec params");
		goto fail;
	}

	ret = avcodec_open2(d->codec_ctx, codec, NULL);
	if (ret < 0) {
		warn("cannot open codec: %s", av_err2str(ret));
		goto fail;
	}

	d->width = d->codec_ctx->width;
	d->height = d->codec_ctx->height;

	/* scaler for 8x8 grayscale (perceptual hash) */
	d->sws_gray = sws_getContext(
		d->width, d->height, d->codec_ctx->pix_fmt,
		8, 8, AV_PIX_FMT_GRAY8,
		SWS_BILINEAR, NULL, NULL, NULL);
	if (!d->sws_gray) {
		warn("cannot create grayscale scaler");
		goto fail;
	}

	/* scaler for full-res RGB (image saving) */
	d->sws_rgb = sws_getContext(
		d->width, d->height, d->codec_ctx->pix_fmt,
		d->width, d->height, AV_PIX_FMT_RGB24,
		SWS_BILINEAR, NULL, NULL, NULL);
	if (!d->sws_rgb) {
		warn("cannot create RGB scaler");
		goto fail;
	}

	return 0;

fail:
	dec_close(d);
	return -1;
}

int
dec_read_frame(Decoder *d, AVFrame *frame)
{
	AVPacket *pkt;
	int ret;

	pkt = av_packet_alloc();
	if (!pkt)
		return -1;

	for (;;) {
		ret = av_read_frame(d->fmt_ctx, pkt);
		if (ret < 0) {
			av_packet_free(&pkt);
			/* flush decoder */
			avcodec_send_packet(d->codec_ctx, NULL);
			ret = avcodec_receive_frame(d->codec_ctx,
			                            frame);
			return (ret == 0) ? 0 : -1;
		}

		if (pkt->stream_index != d->stream_idx) {
			av_packet_unref(pkt);
			continue;
		}

		ret = avcodec_send_packet(d->codec_ctx, pkt);
		av_packet_unref(pkt);
		if (ret < 0) {
			av_packet_free(&pkt);
			return -1;
		}

		ret = avcodec_receive_frame(d->codec_ctx, frame);
		if (ret == AVERROR(EAGAIN))
			continue;
		if (ret < 0) {
			av_packet_free(&pkt);
			return -1;
		}

		av_packet_free(&pkt);
		return 0;
	}
}

double
dec_frame_time(Decoder *d, AVFrame *frame)
{
	AVStream *st;
	double tb;

	st = d->fmt_ctx->streams[d->stream_idx];
	tb = av_q2d(st->time_base);

	if (frame->pts != AV_NOPTS_VALUE)
		return frame->pts * tb;
	if (frame->pkt_dts != AV_NOPTS_VALUE)
		return frame->pkt_dts * tb;
	return 0.0;
}

void
dec_close(Decoder *d)
{
	if (d->sws_rgb)
		sws_freeContext(d->sws_rgb);
	if (d->sws_gray)
		sws_freeContext(d->sws_gray);
	if (d->codec_ctx)
		avcodec_free_context(&d->codec_ctx);
	if (d->fmt_ctx)
		avformat_close_input(&d->fmt_ctx);
	memset(d, 0, sizeof(*d));
}