1/*
2 * Copyright 2009-2010, Stephan A��mus <superstippi@gmx.de>
3 * Copyright 2018, Dario Casalinuovo
4 * All rights reserved. Distributed under the terms of the GNU L-GPL license.
5 */
6
7#include "AVFormatWriter.h"
8
9#include <stdio.h>
10#include <string.h>
11#include <stdlib.h>
12
13#include <new>
14
15#include <Application.h>
16#include <AutoDeleter.h>
17#include <Autolock.h>
18#include <ByteOrder.h>
19#include <MediaIO.h>
20#include <MediaDefs.h>
21#include <MediaFormats.h>
22#include <Roster.h>
23
24extern "C" {
25	#include "avformat.h"
26}
27
28#include "DemuxerTable.h"
29#include "EncoderTable.h"
30#include "gfx_util.h"
31
32
33//#define TRACE_AVFORMAT_WRITER
34#ifdef TRACE_AVFORMAT_WRITER
35#	define TRACE printf
36#	define TRACE_IO(a...)
37#	define TRACE_PACKET printf
38#else
39#	define TRACE(a...)
40#	define TRACE_IO(a...)
41#	define TRACE_PACKET(a...)
42#endif
43
44#define ERROR(a...) fprintf(stderr, a)
45
46
47static const size_t kIOBufferSize = 64 * 1024;
48	// TODO: This could depend on the BMediaFile creation flags, IIRC,
49	// they allow to specify a buffering mode.
50
51typedef AVCodecID CodecID;
52
53// #pragma mark - AVFormatWriter::StreamCookie
54
55
56class AVFormatWriter::StreamCookie {
57public:
58								StreamCookie(AVFormatContext* context,
59									BLocker* streamLock);
60	virtual						~StreamCookie();
61
62			status_t			Init(media_format* format,
63									const media_codec_info* codecInfo);
64
65			status_t			WriteChunk(const void* chunkBuffer,
66									size_t chunkSize,
67									media_encode_info* encodeInfo);
68
69			status_t			AddTrackInfo(uint32 code, const void* data,
70									size_t size, uint32 flags);
71
72private:
73			AVFormatContext*	fFormatContext;
74			AVStream*			fStream;
75			AVPacket*			fPacket;
76			// Since different threads may write to the target,
77			// we need to protect the file position and I/O by a lock.
78			BLocker*			fStreamLock;
79};
80
81
82
83AVFormatWriter::StreamCookie::StreamCookie(AVFormatContext* context,
84		BLocker* streamLock)
85	:
86	fFormatContext(context),
87	fStream(NULL),
88	fStreamLock(streamLock)
89{
90	fPacket = av_packet_alloc();
91}
92
93
94AVFormatWriter::StreamCookie::~StreamCookie()
95{
96	// fStream is freed automatically when the codec context is closed
97	av_packet_free(&fPacket);
98}
99
100
101static void
102set_channel_count(AVCodecParameters* context, int count)
103{
104#if LIBAVCODEC_VERSION_MAJOR >= 60
105	context->ch_layout.nb_channels = count;
106#else
107	context->channels = count;
108#endif
109}
110
111
112status_t
113AVFormatWriter::StreamCookie::Init(media_format* format,
114	const media_codec_info* codecInfo)
115{
116	TRACE("AVFormatWriter::StreamCookie::Init()\n");
117
118	BAutolock _(fStreamLock);
119
120	fPacket->stream_index = fFormatContext->nb_streams;
121	fStream = avformat_new_stream(fFormatContext, NULL);
122
123	if (fStream == NULL) {
124		TRACE("  failed to add new stream\n");
125		return B_ERROR;
126	}
127
128	fStream->id = fPacket->stream_index;
129
130//	TRACE("  fStream->codecpar: %p\n", fStream->codecpar);
131	// TODO: This is a hack for now! Use avcodec_find_encoder_by_name()
132	// or something similar...
133	fStream->codecpar->codec_id = (CodecID)codecInfo->sub_id;
134	if (fStream->codecpar->codec_id == AV_CODEC_ID_NONE)
135		fStream->codecpar->codec_id = raw_audio_codec_id_for(*format);
136
137	// Setup the stream according to the media format...
138	if (format->type == B_MEDIA_RAW_VIDEO) {
139		fStream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
140		fStream->time_base.den = (int)format->u.raw_video.field_rate;
141		fStream->time_base.num = 1;
142
143		// video size
144		fStream->codecpar->width = format->u.raw_video.display.line_width;
145		fStream->codecpar->height = format->u.raw_video.display.line_count;
146		// pixel aspect ratio
147		fStream->sample_aspect_ratio.num
148			= format->u.raw_video.pixel_width_aspect;
149		fStream->sample_aspect_ratio.den
150			= format->u.raw_video.pixel_height_aspect;
151		if (fStream->sample_aspect_ratio.num == 0
152			|| fStream->sample_aspect_ratio.den == 0) {
153			av_reduce(&fStream->sample_aspect_ratio.num,
154				&fStream->sample_aspect_ratio.den, fStream->codecpar->width,
155				fStream->codecpar->height, 255);
156		}
157
158		fStream->codecpar->sample_aspect_ratio = fStream->sample_aspect_ratio;
159
160		// Use the last supported pixel format of the AVCodec, which we hope
161		// is the one with the best quality (true for all currently supported
162		// encoders).
163//		AVCodec* codec = fStream->codecpar->codec;
164//		for (int i = 0; codec->pix_fmts[i] != PIX_FMT_NONE; i++)
165//			fStream->codecpar->pix_fmt = codec->pix_fmts[i];
166		fStream->codecpar->format = AV_PIX_FMT_YUV420P;
167
168	} else if (format->type == B_MEDIA_RAW_AUDIO) {
169		fStream->codecpar->codec_type = AVMEDIA_TYPE_AUDIO;
170
171		// frame rate
172		fStream->codecpar->sample_rate = (int)format->u.raw_audio.frame_rate;
173
174		// channels
175		set_channel_count(fStream->codecpar, format->u.raw_audio.channel_count);
176
177		// set fStream to the audio format we want to use. This is only a hint
178		// (each encoder has a different set of accepted formats)
179		switch (format->u.raw_audio.format) {
180			case media_raw_audio_format::B_AUDIO_FLOAT:
181				fStream->codecpar->format = AV_SAMPLE_FMT_FLT;
182				break;
183			case media_raw_audio_format::B_AUDIO_DOUBLE:
184				fStream->codecpar->format = AV_SAMPLE_FMT_DBL;
185				break;
186			case media_raw_audio_format::B_AUDIO_INT:
187				fStream->codecpar->format = AV_SAMPLE_FMT_S32;
188				break;
189			case media_raw_audio_format::B_AUDIO_SHORT:
190				fStream->codecpar->format = AV_SAMPLE_FMT_S16;
191				break;
192			case media_raw_audio_format::B_AUDIO_UCHAR:
193				fStream->codecpar->format = AV_SAMPLE_FMT_U8;
194				break;
195
196			case media_raw_audio_format::B_AUDIO_CHAR:
197			default:
198				return B_MEDIA_BAD_FORMAT;
199				break;
200		}
201
202		// Now negociate the actual format with the encoder
203		// First check if the requested format is acceptable
204		const AVCodec* codec = avcodec_find_encoder(fStream->codecpar->codec_id);
205
206		if (codec == NULL)
207			return B_MEDIA_BAD_FORMAT;
208
209		const enum AVSampleFormat *p = codec->sample_fmts;
210		for (; *p != -1; p++) {
211			if (*p == fStream->codecpar->format)
212				break;
213		}
214		// If not, force one of the acceptable ones
215		if (*p == -1) {
216			fStream->codecpar->format = codec->sample_fmts[0];
217
218			// And finally set the format struct to the accepted format. It is
219			// then up to the caller to make sure we get data matching that
220			// format.
221			switch (fStream->codecpar->format) {
222				case AV_SAMPLE_FMT_FLT:
223					format->u.raw_audio.format
224						= media_raw_audio_format::B_AUDIO_FLOAT;
225					break;
226				case AV_SAMPLE_FMT_DBL:
227					format->u.raw_audio.format
228						= media_raw_audio_format::B_AUDIO_DOUBLE;
229					break;
230				case AV_SAMPLE_FMT_S32:
231					format->u.raw_audio.format
232						= media_raw_audio_format::B_AUDIO_INT;
233					break;
234				case AV_SAMPLE_FMT_S16:
235					format->u.raw_audio.format
236						= media_raw_audio_format::B_AUDIO_SHORT;
237					break;
238				case AV_SAMPLE_FMT_U8:
239					format->u.raw_audio.format
240						= media_raw_audio_format::B_AUDIO_UCHAR;
241					break;
242				default:
243					return B_MEDIA_BAD_FORMAT;
244					break;
245			}
246		}
247
248#if LIBAVCODEC_VERSION_MAJOR >= 60
249		if (format->u.raw_audio.channel_mask == 0) {
250			// guess the channel mask...
251			av_channel_layout_default(&fStream->codecpar->ch_layout,
252				format->u.raw_audio.channel_count);
253		} else {
254			// The bits match 1:1 for media_multi_channels and FFmpeg defines.
255			av_channel_layout_from_mask(&fStream->codecpar->ch_layout,
256				format->u.raw_audio.channel_mask);
257		}
258#else
259		if (format->u.raw_audio.channel_mask == 0) {
260			// guess the channel mask...
261			switch (format->u.raw_audio.channel_count) {
262				default:
263				case 2:
264					fStream->codecpar->channel_layout = AV_CH_LAYOUT_STEREO;
265					break;
266				case 1:
267					fStream->codecpar->channel_layout = AV_CH_LAYOUT_MONO;
268					break;
269				case 3:
270					fStream->codecpar->channel_layout = AV_CH_LAYOUT_SURROUND;
271					break;
272				case 4:
273					fStream->codecpar->channel_layout = AV_CH_LAYOUT_QUAD;
274					break;
275				case 5:
276					fStream->codecpar->channel_layout = AV_CH_LAYOUT_5POINT0;
277					break;
278				case 6:
279					fStream->codecpar->channel_layout = AV_CH_LAYOUT_5POINT1;
280					break;
281				case 8:
282					fStream->codecpar->channel_layout = AV_CH_LAYOUT_7POINT1;
283					break;
284				case 10:
285					fStream->codecpar->channel_layout = AV_CH_LAYOUT_7POINT1_WIDE;
286					break;
287			}
288		} else {
289			// The bits match 1:1 for media_multi_channels and FFmpeg defines.
290			fStream->codecpar->channel_layout = format->u.raw_audio.channel_mask;
291		}
292#endif
293	}
294
295	TRACE("  stream->time_base: (%d/%d)\n",
296		fStream->time_base.num, fStream->time_base.den);
297
298#if 0
299	// Write the AVCodecContext pointer to the user data section of the
300	// media_format. For some encoders, it seems to be necessary to use
301	// the AVCodecContext of the AVStream in order to successfully encode
302	// anything and write valid media files. For example some codecs need
303	// to store meta data or global data in the container.
304	app_info appInfo;
305	if (be_app->GetAppInfo(&appInfo) == B_OK) {
306		uchar* userData = format->user_data;
307		*(uint32*)userData = 'ffmp';
308		userData += sizeof(uint32);
309		*(team_id*)userData = appInfo.team;
310		userData += sizeof(team_id);
311		*(AVCodecContext**)userData = fStream->codec;
312	}
313#endif
314
315	return B_OK;
316}
317
318
319status_t
320AVFormatWriter::StreamCookie::WriteChunk(const void* chunkBuffer,
321	size_t chunkSize, media_encode_info* encodeInfo)
322{
323	TRACE_PACKET("AVFormatWriter::StreamCookie[%d]::WriteChunk(%p, %ld, "
324		"start_time: %" B_PRIdBIGTIME ")\n", fStream->index, chunkBuffer, chunkSize,
325		encodeInfo->start_time);
326
327	BAutolock _(fStreamLock);
328
329	fPacket->data = const_cast<uint8_t*>((const uint8_t*)chunkBuffer);
330	fPacket->size = chunkSize;
331	fPacket->stream_index = fStream->index;
332
333	fPacket->pts = int64_t((double)encodeInfo->start_time
334		* fStream->time_base.den / (1000000.0 * fStream->time_base.num)
335		+ 0.5);
336
337	fPacket->dts = fPacket->pts;
338
339	fPacket->flags = 0;
340	if ((encodeInfo->flags & B_MEDIA_KEY_FRAME) != 0)
341		fPacket->flags |= AV_PKT_FLAG_KEY;
342
343	TRACE_PACKET("  PTS: %" PRId64 " (stream->time_base: (%d/%d)\n", fPacket->pts,
344		fStream->time_base.num, fStream->time_base.den);
345
346#if 0
347	// TODO: Eventually, we need to write interleaved packets, but
348	// maybe we are only supposed to use this if we have actually
349	// more than one stream. For the moment, this crashes in AVPacket
350	// shuffling inside libavformat. Maybe if we want to use this, we
351	// need to allocate a separate AVPacket and copy the chunk buffer.
352	int result = av_interleaved_write_frame(fFormatContext, fPacket);
353	if (result < 0)
354		TRACE("  av_interleaved_write_frame(): %d\n", result);
355#else
356	int result = av_write_frame(fFormatContext, fPacket);
357	if (result < 0)
358		TRACE("  av_write_frame(): %d\n", result);
359#endif
360
361	return result == 0 ? B_OK : B_ERROR;
362}
363
364
365status_t
366AVFormatWriter::StreamCookie::AddTrackInfo(uint32 code,
367	const void* data, size_t size, uint32 flags)
368{
369	TRACE("AVFormatWriter::StreamCookie::AddTrackInfo(%" B_PRIu32 ", %p, %ld, %" B_PRIu32 ")\n",
370		code, data, size, flags);
371
372	BAutolock _(fStreamLock);
373
374	return B_NOT_SUPPORTED;
375}
376
377
378// #pragma mark - AVFormatWriter
379
380
381AVFormatWriter::AVFormatWriter()
382	:
383	fFormatContext(avformat_alloc_context()),
384	fCodecOpened(false),
385	fHeaderError(-1),
386	fIOContext(NULL),
387	fStreamLock("stream lock")
388{
389	TRACE("AVFormatWriter::AVFormatWriter\n");
390}
391
392
393AVFormatWriter::~AVFormatWriter()
394{
395	TRACE("AVFormatWriter::~AVFormatWriter\n");
396
397	// Free the streams and close the AVCodecContexts
398	for (unsigned i = 0; i < fFormatContext->nb_streams; i++) {
399		av_freep(&fFormatContext->streams[i]->codecpar);
400		av_freep(&fFormatContext->streams[i]);
401	}
402
403	avformat_free_context(fFormatContext);
404	av_free(fIOContext->buffer);
405	av_free(fIOContext);
406}
407
408
409// #pragma mark -
410
411
412status_t
413AVFormatWriter::Init(const media_file_format* fileFormat)
414{
415	TRACE("AVFormatWriter::Init()\n");
416
417	if (fIOContext == NULL) {
418		uint8* buffer = static_cast<uint8*>(av_malloc(kIOBufferSize));
419		if (buffer == NULL)
420			return B_NO_MEMORY;
421
422		// Allocate I/O context and initialize it with buffer
423		// and hook functions, pass ourself as cookie.
424		fIOContext = avio_alloc_context(buffer, kIOBufferSize, 1, this,
425				0, _Write, _Seek);
426		if (fIOContext == NULL) {
427			av_free(buffer);
428			TRACE("av_alloc_put_byte() failed!\n");
429			return B_ERROR;
430		}
431
432		// Setup I/O hooks. This seems to be enough.
433		fFormatContext->pb = fIOContext;
434	}
435
436	// Set the AVOutputFormat according to fileFormat...
437	fFormatContext->oformat = av_guess_format(fileFormat->short_name,
438		fileFormat->file_extension, fileFormat->mime_type);
439	if (fFormatContext->oformat == NULL) {
440		TRACE("  failed to find AVOuputFormat for %s\n",
441			fileFormat->short_name);
442		return B_NOT_SUPPORTED;
443	}
444
445	TRACE("  found AVOuputFormat for %s: %s\n", fileFormat->short_name,
446		fFormatContext->oformat->name);
447
448	return B_OK;
449}
450
451
452status_t
453AVFormatWriter::SetCopyright(const char* copyright)
454{
455	TRACE("AVFormatWriter::SetCopyright(%s)\n", copyright);
456
457	return B_NOT_SUPPORTED;
458}
459
460
461status_t
462AVFormatWriter::CommitHeader()
463{
464	TRACE("AVFormatWriter::CommitHeader\n");
465
466	if (fFormatContext == NULL)
467		return B_NO_INIT;
468
469	if (fCodecOpened)
470		return B_NOT_ALLOWED;
471
472	// We need to close the codecs we opened, even in case of failure.
473	fCodecOpened = true;
474
475	fHeaderError = avformat_write_header(fFormatContext, NULL);
476
477	#ifdef TRACE_AVFORMAT_WRITER
478	if (fHeaderError < 0) {
479		char errorBuffer[AV_ERROR_MAX_STRING_SIZE];
480		av_strerror(fHeaderError, errorBuffer, sizeof(errorBuffer));
481		TRACE("  avformat_write_header(): %s\n", errorBuffer);
482	} else {
483		TRACE("  wrote header\n");
484	}
485
486	for (unsigned i = 0; i < fFormatContext->nb_streams; i++) {
487		AVStream* stream = fFormatContext->streams[i];
488		TRACE("  stream[%u] time_base: (%d/%d)\n",
489			i, stream->time_base.num, stream->time_base.den);
490	}
491	#endif // TRACE_AVFORMAT_WRITER
492
493	return fHeaderError == 0 ? B_OK : B_ERROR;
494}
495
496
497status_t
498AVFormatWriter::Flush()
499{
500	TRACE("AVFormatWriter::Flush\n");
501
502	return B_NOT_SUPPORTED;
503}
504
505
506status_t
507AVFormatWriter::Close()
508{
509	TRACE("AVFormatWriter::Close\n");
510
511	if (fFormatContext == NULL)
512		return B_NO_INIT;
513
514	if (!fCodecOpened)
515		return B_NOT_ALLOWED;
516
517	// From ffmpeg documentation: [av_write_trailer] may only be called
518	// after a successful call to avformat_write_header.
519	if (fHeaderError != 0)
520		return B_ERROR;
521
522	int result = av_write_trailer(fFormatContext);
523	if (result < 0)
524		TRACE("  av_write_trailer(): %d\n", result);
525	return result == 0 ? B_OK : B_ERROR;
526}
527
528
529status_t
530AVFormatWriter::AllocateCookie(void** _cookie, media_format* format,
531	const media_codec_info* codecInfo)
532{
533	TRACE("AVFormatWriter::AllocateCookie()\n");
534
535	if (fCodecOpened)
536		return B_NOT_ALLOWED;
537
538	BAutolock _(fStreamLock);
539
540	if (_cookie == NULL)
541		return B_BAD_VALUE;
542
543	StreamCookie* cookie = new(std::nothrow) StreamCookie(fFormatContext,
544		&fStreamLock);
545
546	status_t ret = cookie->Init(format, codecInfo);
547	if (ret != B_OK) {
548		delete cookie;
549		return ret;
550	}
551
552	*_cookie = cookie;
553	return B_OK;
554}
555
556
557status_t
558AVFormatWriter::FreeCookie(void* _cookie)
559{
560	BAutolock _(fStreamLock);
561
562	StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
563	delete cookie;
564
565	return B_OK;
566}
567
568
569// #pragma mark -
570
571
572status_t
573AVFormatWriter::SetCopyright(void* cookie, const char* copyright)
574{
575	TRACE("AVFormatWriter::SetCopyright(%p, %s)\n", cookie, copyright);
576
577	return B_NOT_SUPPORTED;
578}
579
580
581status_t
582AVFormatWriter::AddTrackInfo(void* _cookie, uint32 code,
583	const void* data, size_t size, uint32 flags)
584{
585	TRACE("AVFormatWriter::AddTrackInfo(%" B_PRIu32 ", %p, %ld, %" B_PRIu32 ")\n",
586		code, data, size, flags);
587
588	if (fHeaderError != 0)
589		return B_ERROR;
590
591	StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
592	return cookie->AddTrackInfo(code, data, size, flags);
593}
594
595
596status_t
597AVFormatWriter::WriteChunk(void* _cookie, const void* chunkBuffer,
598	size_t chunkSize, media_encode_info* encodeInfo)
599{
600	TRACE_PACKET("AVFormatWriter::WriteChunk(%p, %ld, %p)\n", chunkBuffer,
601		chunkSize, encodeInfo);
602
603	if (fHeaderError != 0)
604		return B_ERROR;
605
606	StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
607	return cookie->WriteChunk(chunkBuffer, chunkSize, encodeInfo);
608}
609
610
611// #pragma mark - I/O hooks
612
613
614/*static*/ int
615AVFormatWriter::_Write(void* cookie, uint8* buffer, int bufferSize)
616{
617	TRACE_IO("AVFormatWriter::_Write(%p, %p, %d)\n",
618		cookie, buffer, bufferSize);
619
620	AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);
621
622	ssize_t written = writer->fTarget->Write(buffer, bufferSize);
623
624	TRACE_IO("  written: %ld\n", written);
625	return (int)written;
626
627}
628
629
630/*static*/ off_t
631AVFormatWriter::_Seek(void* cookie, off_t offset, int whence)
632{
633	TRACE_IO("AVFormatWriter::_Seek(%p, %lld, %d)\n",
634		cookie, offset, whence);
635
636	AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);
637
638	BMediaIO* mediaIO = dynamic_cast<BMediaIO*>(writer->fTarget);
639	if (mediaIO == NULL)
640		return -1;
641
642	// Support for special file size retrieval API without seeking anywhere:
643	if (whence == AVSEEK_SIZE) {
644		off_t size;
645		if (mediaIO->GetSize(&size) == B_OK)
646			return size;
647
648		return -1;
649	}
650
651	off_t position = mediaIO->Seek(offset, whence);
652	TRACE_IO("  position: %lld\n", position);
653	if (position < 0)
654		return -1;
655
656	return position;
657}
658
659
660