path: root/src/recorder.cpp

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#define __STDC_CONSTANT_MACROS
#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>

#include <SDL.h>

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

#include "osc_graphics.h"
#include "osc_server.h"

#include "recorder.h"

/*
 * Macros
 */
#define FFMPEG_ERROR(AVERROR, FMT, ...) do {		\
	char buf[255];					\
							\
	av_strerror(AVERROR, buf, sizeof(buf));		\
	ERROR_MSG(FMT ": %s", ##__VA_ARGS__, buf);	\
} while (0)

extern "C" {

static int start_handler(const char *path, const char *types,
			 lo_arg **argv, int argc,
			 void *data, void *user_data);
static int stop_handler(const char *path, const char *types,
			lo_arg **argv, int argc,
			void *data, void *user_data);

}

Recorder::Recorder() : Mutex(), ffmpeg(NULL), stream(NULL), sws_context(NULL),
				encodeFrame(NULL), encodeFrameBuffer(NULL),
				pkt(NULL)
{
	static bool initialized = false;

	if (!initialized) {
		avcodec_register_all();
		av_register_all();

		initialized = true;
	}
}

void
Recorder::register_methods()
{
	osc_server.add_method("ss", start_handler, this, "/recorder/start");
	osc_server.add_method("", stop_handler, this, "/recorder/stop");
}

static int
start_handler(const char *path, const char *types,
	      lo_arg **argv, int argc, void *data, void *user_data)
{
	Recorder *obj = (Recorder *)user_data;

	obj->start(&argv[0]->s, &argv[1]->s);

	return 0;
}

static inline enum PixelFormat
get_pix_fmt(SDL_PixelFormat *format)
{
	switch (format->BitsPerPixel) {
	case 8:
		return PIX_FMT_PAL8;
	case 24:
		return PIX_FMT_RGB24;
	case 32:
		return PIX_FMT_RGB32;
	}

	return PIX_FMT_NONE;
}

void
Recorder::start(const char *filename, const char *codecname)
{
	AVCodec *videoCodec;
	int err;

	stop();

	if (!filename || !*filename)
		return;

	lock();

	ffmpeg = avformat_alloc_context();

	ffmpeg->oformat = av_guess_format(NULL, filename, NULL);
	if (!ffmpeg->oformat)
		ffmpeg->oformat = av_guess_format("dvd", NULL, NULL);

	//ffmpeg->preload = (int)(0.5 * AV_TIME_BASE);
	ffmpeg->max_delay = (int)(0.7 * AV_TIME_BASE);

#if LIBAVFORMAT_VERSION_MAJOR < 54
	err = url_fopen(&ffmpeg->pb, filename, URL_WRONLY);
	if (err < 0) {
		FFMPEG_ERROR(-err, "url_fopen");
		exit(EXIT_FAILURE);
	}
#else
	err = avio_open(&ffmpeg->pb, filename, AVIO_FLAG_WRITE);
	if (err < 0) {
		FFMPEG_ERROR(-err, "avio_open");
		exit(EXIT_FAILURE);
	}
#endif

	stream = av_new_stream(ffmpeg, 0);
	if (!stream) {
		ERROR_MSG("Could not open stream!");
		exit(EXIT_FAILURE);
	}

	stream->codec = avcodec_alloc_context();

#if LIBAVCODEC_VERSION_MAJOR < 54
	avcodec_get_context_defaults2(stream->codec, CODEC_TYPE_VIDEO);
	stream->codec->codec_type = CODEC_TYPE_VIDEO;
#else
	avcodec_get_context_defaults2(stream->codec, AVMEDIA_TYPE_VIDEO);
	stream->codec->codec_type = AVMEDIA_TYPE_VIDEO;
#endif

	stream->codec->bit_rate = 6000000;

	/* resolution must be a multiple of two */
	stream->codec->width = screen->w;
	stream->codec->height = screen->h;

	/* set time_base */
	stream->codec->time_base.num = 1;
	stream->codec->time_base.den = config_framerate;

	/* emit one intra frame every twelve frames at most */
	stream->codec->gop_size = 12;

	switch (stream->codec->codec_id) {
	case CODEC_ID_MPEG2VIDEO:
		/* set mpeg2 codec parameters */
		stream->codec->max_b_frames = 2;
		break;
	case CODEC_ID_MPEG1VIDEO:
		/* set mpeg1 codec parameters */
		/* needed to avoid using macroblocks in which some coeffs overflow
		   this doesnt happen with normal video, it just happens here as the
		   motion of the chroma plane doesnt match the luma plane */
		stream->codec->mb_decision = 2;
		break;
	default:
		break;
	}

	/* some formats want stream headers to be separate */
	if (ffmpeg->oformat->flags & AVFMT_GLOBALHEADER)
		stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;

	/* find the video encoder */
	if (codecname && *codecname)
		videoCodec = avcodec_find_encoder_by_name(codecname);
	else
		videoCodec = avcodec_find_encoder(ffmpeg->oformat->video_codec);
	if (!videoCodec) {
		ERROR_MSG("Could not find encoder %s (%d)!",
			  codecname ? : "NULL",
			  ffmpeg->oformat->video_codec);
		exit(EXIT_FAILURE);
	}
	stream->codec->codec_id = videoCodec->id;

	enum PixelFormat screen_fmt = get_pix_fmt(screen->format);

	/* set pixel format */
	stream->codec->pix_fmt = PIX_FMT_YUV420P;
	for (const PixelFormat *fmt = videoCodec->pix_fmts; *fmt != -1; fmt++) {
		if (*fmt == screen_fmt) {
			stream->codec->pix_fmt = screen_fmt;
			break;
		}
	}

	if (stream->codec->pix_fmt != screen_fmt) {
		WARNING_MSG("Pixel format conversion necessary!");

		/* NOTE: sizes shouldn't differ */
		sws_context = sws_getContext(screen->w, screen->h, screen_fmt,
					     stream->codec->width,
					     stream->codec->height,
					     stream->codec->pix_fmt,
					     SWS_BILINEAR, NULL, NULL, NULL);
	}

	/* open the codec */
	err = avcodec_open(stream->codec, videoCodec);
	if (err < 0) {
		FFMPEG_ERROR(-err, "avcodec_open");
		exit(EXIT_FAILURE);
	}

	encodeFrame = avcodec_alloc_frame();

	if (sws_context) {
		int size = avpicture_get_size(stream->codec->pix_fmt,
					      stream->codec->width,
					      stream->codec->height) +
			   FF_INPUT_BUFFER_PADDING_SIZE;
		avpicture_fill((AVPicture *)encodeFrame, new uint8_t[size],
			       stream->codec->pix_fmt,
			       stream->codec->width,
			       stream->codec->height);
	}

	encodeFrameBufferSize = stream->codec->width * stream->codec->height * 4 +
				FF_INPUT_BUFFER_PADDING_SIZE;
	encodeFrameBuffer = new uint8_t[encodeFrameBufferSize];

	pkt = new AVPacket;

	/* try to write a header */
#if LIBAVFORMAT_VERSION_MAJOR < 54
	err = av_set_parameters(ffmpeg, NULL);
	if (err < 0) {
		FFMPEG_ERROR(-err, "av_set_parameters");
		exit(EXIT_FAILURE);
	}

	av_write_header(ffmpeg);
#else
	avformat_write_header(ffmpeg, NULL);
#endif

	start_time = SDL_GetTicks();

	unlock();
}

static int
stop_handler(const char *path, const char *types,
	     lo_arg **argv, int argc, void *data, void *user_data)
{
	Recorder *obj = (Recorder *)user_data;

	obj->stop();

	return 0;
}

void
Recorder::stop()
{
	lock();

	if (stream && stream->codec)
		avcodec_flush_buffers(stream->codec);

	if (ffmpeg)
		av_write_trailer(ffmpeg);

	if (pkt) {
		av_free_packet(pkt);
		delete pkt;
		pkt = NULL;
	}
	if (encodeFrameBuffer) {
		delete encodeFrameBuffer;
		encodeFrameBuffer = NULL;
	}
	if (sws_context) {
		sws_freeContext(sws_context);
		sws_context = NULL;

		if (encodeFrame)
			delete encodeFrame->data[0];
	}
	if (encodeFrame) {
		av_free(encodeFrame);
		encodeFrame = NULL;
	}
	if (stream) {
		if (stream->codec)
			avcodec_close(stream->codec);
		av_free(stream);
		stream = NULL;
	}
	if (ffmpeg) {
#if LIBAVFORMAT_VERSION_MAJOR < 54
		url_fclose(ffmpeg->pb);
#else
		avio_close(ffmpeg->pb);
#endif
		av_free(ffmpeg);
		ffmpeg = NULL;
	}

	unlock();
}

void
Recorder::record(SDL_Surface *surf)
{
	lock();

	if (!ffmpeg) {
		unlock();
		return;
	}

	int64_t pts = (SDL_GetTicks() - start_time)/FRAME_DELAY;
	if (pts <= encodeFrame->pts) {
		unlock();
		return;
	}
	encodeFrame->pts = pts;

	if (sws_context) {
		const uint8_t *const data[] = {(uint8_t *)surf->pixels, NULL};
		const int pitch[] = {surf->pitch, 0};

		sws_scale(sws_context, data, pitch, 0, surf->h,
			  encodeFrame->data, encodeFrame->linesize);
	} else {
		encodeFrame->data[0] = (uint8_t *)surf->pixels;
		encodeFrame->data[1] = NULL;
		encodeFrame->linesize[0] = surf->pitch;
		encodeFrame->linesize[1] = 0;
	}

	int out_size = avcodec_encode_video(stream->codec, encodeFrameBuffer,
					    encodeFrameBufferSize, encodeFrame);

	if (out_size > 0) {
		av_init_packet(pkt);

		/* set correct stream index for this packet */
		pkt->stream_index = stream->index;
		/* set keyframe flag if needed */
		if (stream->codec->coded_frame->key_frame)
			pkt->flags |= AV_PKT_FLAG_KEY;
		/* write encoded data into packet */
		pkt->data = encodeFrameBuffer;
		/* set the correct size of this packet */
		pkt->size = out_size;
		/* set the correct duration of this packet */
		pkt->duration = AV_TIME_BASE / stream->time_base.den;

		/* if needed info is available, write pts for this packet */
		if (stream->codec->coded_frame->pts != AV_NOPTS_VALUE)
			pkt->pts = av_rescale_q(stream->codec->coded_frame->pts,
					        stream->codec->time_base,
					        stream->time_base);

		av_write_frame(ffmpeg, pkt);
	}

	unlock();
}

Recorder::~Recorder()
{
	osc_server.del_method("ss", "/recorder/start");
	osc_server.del_method("",   "/recorder/stop");

	stop();
}