/*
 * Copyright 2009-2010, Stephan Aßmus <superstippi@gmx.de>
 * Copyright 2018, Dario Casalinuovo
 * All rights reserved. Distributed under the terms of the GNU L-GPL license.
 */

#include "AVFormatWriter.h"

#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include <new>

#include <Application.h>
#include <AutoDeleter.h>
#include <Autolock.h>
#include <ByteOrder.h>
#include <MediaIO.h>
#include <MediaDefs.h>
#include <MediaFormats.h>
#include <Roster.h>

extern "C" {
        #include "avformat.h"
}

#include "DemuxerTable.h"
#include "EncoderTable.h"
#include "gfx_util.h"


//#define TRACE_AVFORMAT_WRITER
#ifdef TRACE_AVFORMAT_WRITER
#       define TRACE printf
#       define TRACE_IO(a...)
#       define TRACE_PACKET printf
#else
#       define TRACE(a...)
#       define TRACE_IO(a...)
#       define TRACE_PACKET(a...)
#endif

#define ERROR(a...) fprintf(stderr, a)


static const size_t kIOBufferSize = 64 * 1024;
        // TODO: This could depend on the BMediaFile creation flags, IIRC,
        // they allow to specify a buffering mode.

typedef AVCodecID CodecID;

// #pragma mark - AVFormatWriter::StreamCookie


class AVFormatWriter::StreamCookie {
public:
                                                                StreamCookie(AVFormatContext* context,
                                                                        BLocker* streamLock);
        virtual                                         ~StreamCookie();

                        status_t                        Init(media_format* format,
                                                                        const media_codec_info* codecInfo);

                        status_t                        WriteChunk(const void* chunkBuffer,
                                                                        size_t chunkSize,
                                                                        media_encode_info* encodeInfo);

                        status_t                        AddTrackInfo(uint32 code, const void* data,
                                                                        size_t size, uint32 flags);

private:
                        AVFormatContext*        fFormatContext;
                        AVStream*                       fStream;
                        AVPacket*                       fPacket;
                        // Since different threads may write to the target,
                        // we need to protect the file position and I/O by a lock.
                        BLocker*                        fStreamLock;
};



AVFormatWriter::StreamCookie::StreamCookie(AVFormatContext* context,
                BLocker* streamLock)
        :
        fFormatContext(context),
        fStream(NULL),
        fStreamLock(streamLock)
{
        fPacket = av_packet_alloc();
}


AVFormatWriter::StreamCookie::~StreamCookie()
{
        // fStream is freed automatically when the codec context is closed
        av_packet_free(&fPacket);
}


static void
set_channel_count(AVCodecParameters* context, int count)
{
#if LIBAVCODEC_VERSION_MAJOR >= 60
        context->ch_layout.nb_channels = count;
#else
        context->channels = count;
#endif
}


status_t
AVFormatWriter::StreamCookie::Init(media_format* format,
        const media_codec_info* codecInfo)
{
        TRACE("AVFormatWriter::StreamCookie::Init()\n");

        BAutolock _(fStreamLock);

        fPacket->stream_index = fFormatContext->nb_streams;
        fStream = avformat_new_stream(fFormatContext, NULL);

        if (fStream == NULL) {
                TRACE("  failed to add new stream\n");
                return B_ERROR;
        }

        fStream->id = fPacket->stream_index;

//      TRACE("  fStream->codecpar: %p\n", fStream->codecpar);
        // TODO: This is a hack for now! Use avcodec_find_encoder_by_name()
        // or something similar...
        fStream->codecpar->codec_id = (CodecID)codecInfo->sub_id;
        if (fStream->codecpar->codec_id == AV_CODEC_ID_NONE)
                fStream->codecpar->codec_id = raw_audio_codec_id_for(*format);

        // Setup the stream according to the media format...
        if (format->type == B_MEDIA_RAW_VIDEO) {
                fStream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
                fStream->time_base.den = (int)format->u.raw_video.field_rate;
                fStream->time_base.num = 1;

                // video size
                fStream->codecpar->width = format->u.raw_video.display.line_width;
                fStream->codecpar->height = format->u.raw_video.display.line_count;
                // pixel aspect ratio
                fStream->sample_aspect_ratio.num
                        = format->u.raw_video.pixel_width_aspect;
                fStream->sample_aspect_ratio.den
                        = format->u.raw_video.pixel_height_aspect;
                if (fStream->sample_aspect_ratio.num == 0
                        || fStream->sample_aspect_ratio.den == 0) {
                        av_reduce(&fStream->sample_aspect_ratio.num,
                                &fStream->sample_aspect_ratio.den, fStream->codecpar->width,
                                fStream->codecpar->height, 255);
                }

                fStream->codecpar->sample_aspect_ratio = fStream->sample_aspect_ratio;

                // Use the last supported pixel format of the AVCodec, which we hope
                // is the one with the best quality (true for all currently supported
                // encoders).
//              AVCodec* codec = fStream->codecpar->codec;
//              for (int i = 0; codec->pix_fmts[i] != PIX_FMT_NONE; i++)
//                      fStream->codecpar->pix_fmt = codec->pix_fmts[i];
                fStream->codecpar->format = AV_PIX_FMT_YUV420P;

        } else if (format->type == B_MEDIA_RAW_AUDIO) {
                fStream->codecpar->codec_type = AVMEDIA_TYPE_AUDIO;

                // frame rate
                fStream->codecpar->sample_rate = (int)format->u.raw_audio.frame_rate;

                // channels
                set_channel_count(fStream->codecpar, format->u.raw_audio.channel_count);

                // set fStream to the audio format we want to use. This is only a hint
                // (each encoder has a different set of accepted formats)
                switch (format->u.raw_audio.format) {
                        case media_raw_audio_format::B_AUDIO_FLOAT:
                                fStream->codecpar->format = AV_SAMPLE_FMT_FLT;
                                break;
                        case media_raw_audio_format::B_AUDIO_DOUBLE:
                                fStream->codecpar->format = AV_SAMPLE_FMT_DBL;
                                break;
                        case media_raw_audio_format::B_AUDIO_INT:
                                fStream->codecpar->format = AV_SAMPLE_FMT_S32;
                                break;
                        case media_raw_audio_format::B_AUDIO_SHORT:
                                fStream->codecpar->format = AV_SAMPLE_FMT_S16;
                                break;
                        case media_raw_audio_format::B_AUDIO_UCHAR:
                                fStream->codecpar->format = AV_SAMPLE_FMT_U8;
                                break;

                        case media_raw_audio_format::B_AUDIO_CHAR:
                        default:
                                return B_MEDIA_BAD_FORMAT;
                                break;
                }

                // Now negociate the actual format with the encoder
                // First check if the requested format is acceptable
                const AVCodec* codec = avcodec_find_encoder(fStream->codecpar->codec_id);

                if (codec == NULL)
                        return B_MEDIA_BAD_FORMAT;

                const enum AVSampleFormat *p = codec->sample_fmts;
                for (; *p != -1; p++) {
                        if (*p == fStream->codecpar->format)
                                break;
                }
                // If not, force one of the acceptable ones
                if (*p == -1) {
                        fStream->codecpar->format = codec->sample_fmts[0];

                        // And finally set the format struct to the accepted format. It is
                        // then up to the caller to make sure we get data matching that
                        // format.
                        switch (fStream->codecpar->format) {
                                case AV_SAMPLE_FMT_FLT:
                                        format->u.raw_audio.format
                                                = media_raw_audio_format::B_AUDIO_FLOAT;
                                        break;
                                case AV_SAMPLE_FMT_DBL:
                                        format->u.raw_audio.format
                                                = media_raw_audio_format::B_AUDIO_DOUBLE;
                                        break;
                                case AV_SAMPLE_FMT_S32:
                                        format->u.raw_audio.format
                                                = media_raw_audio_format::B_AUDIO_INT;
                                        break;
                                case AV_SAMPLE_FMT_S16:
                                        format->u.raw_audio.format
                                                = media_raw_audio_format::B_AUDIO_SHORT;
                                        break;
                                case AV_SAMPLE_FMT_U8:
                                        format->u.raw_audio.format
                                                = media_raw_audio_format::B_AUDIO_UCHAR;
                                        break;
                                default:
                                        return B_MEDIA_BAD_FORMAT;
                                        break;
                        }
                }

#if LIBAVCODEC_VERSION_MAJOR >= 60
                if (format->u.raw_audio.channel_mask == 0) {
                        // guess the channel mask...
                        av_channel_layout_default(&fStream->codecpar->ch_layout,
                                format->u.raw_audio.channel_count);
                } else {
                        // The bits match 1:1 for media_multi_channels and FFmpeg defines.
                        av_channel_layout_from_mask(&fStream->codecpar->ch_layout,
                                format->u.raw_audio.channel_mask);
                }
#else
                if (format->u.raw_audio.channel_mask == 0) {
                        // guess the channel mask...
                        switch (format->u.raw_audio.channel_count) {
                                default:
                                case 2:
                                        fStream->codecpar->channel_layout = AV_CH_LAYOUT_STEREO;
                                        break;
                                case 1:
                                        fStream->codecpar->channel_layout = AV_CH_LAYOUT_MONO;
                                        break;
                                case 3:
                                        fStream->codecpar->channel_layout = AV_CH_LAYOUT_SURROUND;
                                        break;
                                case 4:
                                        fStream->codecpar->channel_layout = AV_CH_LAYOUT_QUAD;
                                        break;
                                case 5:
                                        fStream->codecpar->channel_layout = AV_CH_LAYOUT_5POINT0;
                                        break;
                                case 6:
                                        fStream->codecpar->channel_layout = AV_CH_LAYOUT_5POINT1;
                                        break;
                                case 8:
                                        fStream->codecpar->channel_layout = AV_CH_LAYOUT_7POINT1;
                                        break;
                                case 10:
                                        fStream->codecpar->channel_layout = AV_CH_LAYOUT_7POINT1_WIDE;
                                        break;
                        }
                } else {
                        // The bits match 1:1 for media_multi_channels and FFmpeg defines.
                        fStream->codecpar->channel_layout = format->u.raw_audio.channel_mask;
                }
#endif
        }

        TRACE("  stream->time_base: (%d/%d)\n",
                fStream->time_base.num, fStream->time_base.den);

#if 0
        // Write the AVCodecContext pointer to the user data section of the
        // media_format. For some encoders, it seems to be necessary to use
        // the AVCodecContext of the AVStream in order to successfully encode
        // anything and write valid media files. For example some codecs need
        // to store meta data or global data in the container.
        app_info appInfo;
        if (be_app->GetAppInfo(&appInfo) == B_OK) {
                uchar* userData = format->user_data;
                *(uint32*)userData = 'ffmp';
                userData += sizeof(uint32);
                *(team_id*)userData = appInfo.team;
                userData += sizeof(team_id);
                *(AVCodecContext**)userData = fStream->codec;
        }
#endif

        return B_OK;
}


status_t
AVFormatWriter::StreamCookie::WriteChunk(const void* chunkBuffer,
        size_t chunkSize, media_encode_info* encodeInfo)
{
        TRACE_PACKET("AVFormatWriter::StreamCookie[%d]::WriteChunk(%p, %ld, "
                "start_time: %" B_PRIdBIGTIME ")\n", fStream->index, chunkBuffer, chunkSize,
                encodeInfo->start_time);

        BAutolock _(fStreamLock);

        fPacket->data = const_cast<uint8_t*>((const uint8_t*)chunkBuffer);
        fPacket->size = chunkSize;
        fPacket->stream_index = fStream->index;

        fPacket->pts = int64_t((double)encodeInfo->start_time
                * fStream->time_base.den / (1000000.0 * fStream->time_base.num)
                + 0.5);

        fPacket->dts = fPacket->pts;

        fPacket->flags = 0;
        if ((encodeInfo->flags & B_MEDIA_KEY_FRAME) != 0)
                fPacket->flags |= AV_PKT_FLAG_KEY;

        TRACE_PACKET("  PTS: %" PRId64 " (stream->time_base: (%d/%d)\n", fPacket->pts,
                fStream->time_base.num, fStream->time_base.den);

#if 0
        // TODO: Eventually, we need to write interleaved packets, but
        // maybe we are only supposed to use this if we have actually
        // more than one stream. For the moment, this crashes in AVPacket
        // shuffling inside libavformat. Maybe if we want to use this, we
        // need to allocate a separate AVPacket and copy the chunk buffer.
        int result = av_interleaved_write_frame(fFormatContext, fPacket);
        if (result < 0)
                TRACE("  av_interleaved_write_frame(): %d\n", result);
#else
        int result = av_write_frame(fFormatContext, fPacket);
        if (result < 0)
                TRACE("  av_write_frame(): %d\n", result);
#endif

        return result == 0 ? B_OK : B_ERROR;
}


status_t
AVFormatWriter::StreamCookie::AddTrackInfo(uint32 code,
        const void* data, size_t size, uint32 flags)
{
        TRACE("AVFormatWriter::StreamCookie::AddTrackInfo(%" B_PRIu32 ", %p, %ld, %" B_PRIu32 ")\n",
                code, data, size, flags);

        BAutolock _(fStreamLock);

        return B_NOT_SUPPORTED;
}


// #pragma mark - AVFormatWriter


AVFormatWriter::AVFormatWriter()
        :
        fFormatContext(avformat_alloc_context()),
        fCodecOpened(false),
        fHeaderError(-1),
        fIOContext(NULL),
        fStreamLock("stream lock")
{
        TRACE("AVFormatWriter::AVFormatWriter\n");
}


AVFormatWriter::~AVFormatWriter()
{
        TRACE("AVFormatWriter::~AVFormatWriter\n");

        // Free the streams and close the AVCodecContexts
        for (unsigned i = 0; i < fFormatContext->nb_streams; i++) {
                av_freep(&fFormatContext->streams[i]->codecpar);
                av_freep(&fFormatContext->streams[i]);
        }

        avformat_free_context(fFormatContext);
        av_free(fIOContext->buffer);
        av_free(fIOContext);
}


// #pragma mark -


status_t
AVFormatWriter::Init(const media_file_format* fileFormat)
{
        TRACE("AVFormatWriter::Init()\n");

        if (fIOContext == NULL) {
                uint8* buffer = static_cast<uint8*>(av_malloc(kIOBufferSize));
                if (buffer == NULL)
                        return B_NO_MEMORY;

                // Allocate I/O context and initialize it with buffer
                // and hook functions, pass ourself as cookie.
                fIOContext = avio_alloc_context(buffer, kIOBufferSize, 1, this,
                                0, _Write, _Seek);
                if (fIOContext == NULL) {
                        av_free(buffer);
                        TRACE("av_alloc_put_byte() failed!\n");
                        return B_ERROR;
                }

                // Setup I/O hooks. This seems to be enough.
                fFormatContext->pb = fIOContext;
        }

        // Set the AVOutputFormat according to fileFormat...
        fFormatContext->oformat = av_guess_format(fileFormat->short_name,
                fileFormat->file_extension, fileFormat->mime_type);
        if (fFormatContext->oformat == NULL) {
                TRACE("  failed to find AVOuputFormat for %s\n",
                        fileFormat->short_name);
                return B_NOT_SUPPORTED;
        }

        TRACE("  found AVOuputFormat for %s: %s\n", fileFormat->short_name,
                fFormatContext->oformat->name);

        return B_OK;
}


status_t
AVFormatWriter::SetCopyright(const char* copyright)
{
        TRACE("AVFormatWriter::SetCopyright(%s)\n", copyright);

        return B_NOT_SUPPORTED;
}


status_t
AVFormatWriter::CommitHeader()
{
        TRACE("AVFormatWriter::CommitHeader\n");

        if (fFormatContext == NULL)
                return B_NO_INIT;

        if (fCodecOpened)
                return B_NOT_ALLOWED;

        // We need to close the codecs we opened, even in case of failure.
        fCodecOpened = true;

        fHeaderError = avformat_write_header(fFormatContext, NULL);

        #ifdef TRACE_AVFORMAT_WRITER
        if (fHeaderError < 0) {
                char errorBuffer[AV_ERROR_MAX_STRING_SIZE];
                av_strerror(fHeaderError, errorBuffer, sizeof(errorBuffer));
                TRACE("  avformat_write_header(): %s\n", errorBuffer);
        } else {
                TRACE("  wrote header\n");
        }

        for (unsigned i = 0; i < fFormatContext->nb_streams; i++) {
                AVStream* stream = fFormatContext->streams[i];
                TRACE("  stream[%u] time_base: (%d/%d)\n",
                        i, stream->time_base.num, stream->time_base.den);
        }
        #endif // TRACE_AVFORMAT_WRITER

        return fHeaderError == 0 ? B_OK : B_ERROR;
}


status_t
AVFormatWriter::Flush()
{
        TRACE("AVFormatWriter::Flush\n");

        return B_NOT_SUPPORTED;
}


status_t
AVFormatWriter::Close()
{
        TRACE("AVFormatWriter::Close\n");

        if (fFormatContext == NULL)
                return B_NO_INIT;

        if (!fCodecOpened)
                return B_NOT_ALLOWED;

        // From ffmpeg documentation: [av_write_trailer] may only be called
        // after a successful call to avformat_write_header.
        if (fHeaderError != 0)
                return B_ERROR;

        int result = av_write_trailer(fFormatContext);
        if (result < 0)
                TRACE("  av_write_trailer(): %d\n", result);
        return result == 0 ? B_OK : B_ERROR;
}


status_t
AVFormatWriter::AllocateCookie(void** _cookie, media_format* format,
        const media_codec_info* codecInfo)
{
        TRACE("AVFormatWriter::AllocateCookie()\n");

        if (fCodecOpened)
                return B_NOT_ALLOWED;

        BAutolock _(fStreamLock);

        if (_cookie == NULL)
                return B_BAD_VALUE;

        StreamCookie* cookie = new(std::nothrow) StreamCookie(fFormatContext,
                &fStreamLock);

        status_t ret = cookie->Init(format, codecInfo);
        if (ret != B_OK) {
                delete cookie;
                return ret;
        }

        *_cookie = cookie;
        return B_OK;
}


status_t
AVFormatWriter::FreeCookie(void* _cookie)
{
        BAutolock _(fStreamLock);

        StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
        delete cookie;

        return B_OK;
}


// #pragma mark -


status_t
AVFormatWriter::SetCopyright(void* cookie, const char* copyright)
{
        TRACE("AVFormatWriter::SetCopyright(%p, %s)\n", cookie, copyright);

        return B_NOT_SUPPORTED;
}


status_t
AVFormatWriter::AddTrackInfo(void* _cookie, uint32 code,
        const void* data, size_t size, uint32 flags)
{
        TRACE("AVFormatWriter::AddTrackInfo(%" B_PRIu32 ", %p, %ld, %" B_PRIu32 ")\n",
                code, data, size, flags);

        if (fHeaderError != 0)
                return B_ERROR;

        StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
        return cookie->AddTrackInfo(code, data, size, flags);
}


status_t
AVFormatWriter::WriteChunk(void* _cookie, const void* chunkBuffer,
        size_t chunkSize, media_encode_info* encodeInfo)
{
        TRACE_PACKET("AVFormatWriter::WriteChunk(%p, %ld, %p)\n", chunkBuffer,
                chunkSize, encodeInfo);

        if (fHeaderError != 0)
                return B_ERROR;

        StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
        return cookie->WriteChunk(chunkBuffer, chunkSize, encodeInfo);
}


// #pragma mark - I/O hooks


/*static*/ int
AVFormatWriter::_Write(void* cookie, uint8* buffer, int bufferSize)
{
        TRACE_IO("AVFormatWriter::_Write(%p, %p, %d)\n",
                cookie, buffer, bufferSize);

        AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);

        ssize_t written = writer->fTarget->Write(buffer, bufferSize);

        TRACE_IO("  written: %ld\n", written);
        return (int)written;

}


/*static*/ off_t
AVFormatWriter::_Seek(void* cookie, off_t offset, int whence)
{
        TRACE_IO("AVFormatWriter::_Seek(%p, %lld, %d)\n",
                cookie, offset, whence);

        AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);

        BMediaIO* mediaIO = dynamic_cast<BMediaIO*>(writer->fTarget);
        if (mediaIO == NULL)
                return -1;

        // Support for special file size retrieval API without seeking anywhere:
        if (whence == AVSEEK_SIZE) {
                off_t size;
                if (mediaIO->GetSize(&size) == B_OK)
                        return size;

                return -1;
        }

        off_t position = mediaIO->Seek(offset, whence);
        TRACE_IO("  position: %lld\n", position);
        if (position < 0)
                return -1;

        return position;
}