/////////////////////////////////////////////////////////////////////////////
// Name:        mediaenc_ffmpeg.cpp
// Purpose:     FFMPEG Media Encoder
// Author:      Alex Thuering
// Created:     04.08.2007
// RCS-ID:      $Id: mediaenc_ffmpeg.cpp,v 1.19 2010/05/03 06:06:48 ntalex Exp $
// Copyright:   (c) Alex Thuering
// Licence:     GPL
/////////////////////////////////////////////////////////////////////////////

#include "mediaenc_ffmpeg.h"
#include <wx/wx.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

#define INT64_C(val) val##LL
#define UINT64_C(val) val##ULL
extern "C" {
#define __STDC_CONSTANT_MACROS
#define __STDC_LIMIT_MACROS
#ifdef HAVE_FFMPEG_AVUTIL_H
#include <avformat.h>
#include <swscale.h>
#else
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
#endif
}

#define AUDIO_BUF_SIZE 50000
#define VIDEO_BUF_SIZE 1835008

wxFfmpegMediaEncoder::wxFfmpegMediaEncoder(int threadCount) {
	m_threadCount = threadCount;
	m_outputCtx = NULL;
	m_videoStm = NULL;
	m_audioStm = NULL;
    m_samples = NULL;
    m_audioOutbuf = NULL;
    m_picture = NULL;
    m_imgConvertCtx = NULL;
    m_videoOutbuf = NULL;
}

wxFfmpegMediaEncoder::~wxFfmpegMediaEncoder() {
    EndEncode();
}

wxString wxFfmpegMediaEncoder::GetBackendVersion() {
	return wxString::Format(wxT("FFmpeg: libavformat %d.%d.%d, libavcodec %d.%d.%d, libavutil %d.%d.%d"),
			LIBAVFORMAT_VERSION_INT>>16, LIBAVFORMAT_VERSION_INT>>8&0xFF, LIBAVFORMAT_VERSION_INT&0xFF,
			LIBAVCODEC_VERSION_INT>>16, LIBAVCODEC_VERSION_INT>>8&0xFF, LIBAVCODEC_VERSION_INT&0xFF,
			LIBAVUTIL_VERSION_INT>>16, LIBAVUTIL_VERSION_INT>>8&0xFF, LIBAVUTIL_VERSION_INT&0xFF);
}

bool wxFfmpegMediaEncoder::BeginEncode(const wxString& fileName, VideoFormat videoFormat,
		AudioFormat audioFormat, AspectRatio aspectRatio, int videoBitrate) {
    EndEncode();
    AVOutputFormat* outputFormat = guess_format("dvd", NULL, NULL);
    if (audioFormat == afNONE)
    	outputFormat = guess_format(NULL, fileName.mb_str(), NULL);
    if (!outputFormat) {
      wxLogError(wxT("Cannot open output format"));
      return false;
    }
    outputFormat->video_codec = videoFormat == vfNONE ? CODEC_ID_NONE : CODEC_ID_MPEG2VIDEO;
    if (audioFormat == afNONE)
    	outputFormat->audio_codec = CODEC_ID_NONE;
    else if (audioFormat == afAC3)
    	outputFormat->audio_codec = CODEC_ID_AC3;
    else
    	outputFormat->audio_codec = CODEC_ID_MP2;
    
    m_outputCtx = avformat_alloc_context();
    if (!m_outputCtx) {
        wxLogError(wxT("memory error"));
        return false;
    }
    
    m_outputCtx->oformat = outputFormat;
    snprintf(m_outputCtx->filename, sizeof(m_outputCtx->filename), "%s", (const char*) fileName.fn_str());
    
    // add video and audio streams
    if (!addVideoStream(outputFormat->video_codec, videoFormat, aspectRatio, videoBitrate))
        return false;
    if (!addAudioStream(outputFormat->audio_codec))
        return false;
    
    if (av_set_parameters(m_outputCtx, NULL) < 0) {
        wxLogError(wxT("Invalid output format parameters"));
        return false;
    }
    dump_format(m_outputCtx, 0, (const char*) fileName.fn_str(), 1);
    m_outputCtx->packet_size = 2048;
    
    // open the audio and video codecs
    if (m_videoStm && !OpenVideoEncoder())
    	return false;
    if (m_audioStm && !OpenAudioEncoder())
    	return false;
    
    // open the output file
    if (url_fopen(&m_outputCtx->pb, fileName.mb_str(), URL_WRONLY) < 0) {
        wxLogError(wxT("Could not open '%s'"), fileName.c_str());
        return false;
    }
    // write the stream header
    m_outputCtx->packet_size = 2048;
    m_outputCtx->mux_rate = 10080000;
    m_outputCtx->preload= (int)(0.5*AV_TIME_BASE);
    m_outputCtx->max_delay= (int)(0.7*AV_TIME_BASE);
    av_write_header(m_outputCtx);
    return true;
}

bool wxFfmpegMediaEncoder::addVideoStream(int codecId, VideoFormat videoFormat, AspectRatio aspectRatio,
		int videoBitrate) {
    if (codecId == CODEC_ID_NONE) {
    	m_videoStm = NULL;
        return true;
    }
    m_videoStm = av_new_stream(m_outputCtx, 0);
    if (!m_videoStm)
    {
        wxLogError(wxT("Could not alloc stream"));
        return false;
    }
    
    if (m_threadCount > 1)
    	avcodec_thread_init(m_videoStm->codec, m_threadCount);
    
    AVCodecContext* c = m_videoStm->codec;
    c->codec_id = (CodecID) codecId;
    c->codec_type = CODEC_TYPE_VIDEO;
    c->bit_rate = videoBitrate * 1000;
    wxSize frameSize = GetFrameSize(videoFormat);
    c->width = frameSize.GetWidth();
    c->height = frameSize.GetHeight();
    c->time_base.den = isNTSC(videoFormat) ? 30000 : 25;
    c->time_base.num = isNTSC(videoFormat) ? 1001 : 1;
    c->gop_size = isNTSC(videoFormat) ? 15 : 12;
    c->pix_fmt = PIX_FMT_YUV420P;
    c->rc_buffer_size = VIDEO_BUF_SIZE;
    c->rc_max_rate = 9000000;
    c->rc_min_rate = 0;
    double ar = aspectRatio == ar16_9 ? (double)16/9 : (double)4/3;
    c->sample_aspect_ratio = av_d2q(ar*c->height/c->width, 255);
    m_videoStm->sample_aspect_ratio = c->sample_aspect_ratio;
    return true;
}

bool wxFfmpegMediaEncoder::addAudioStream(int codecId) {
    if (codecId == CODEC_ID_NONE) {
    	m_audioStm = NULL;
        return true;
    }
    m_audioStm = av_new_stream(m_outputCtx, 1);
    if (!m_audioStm)
    {
        wxLogError(wxT("Could not alloc stream"));
        return false;
    }

    if (m_threadCount > 1)
    	avcodec_thread_init(m_audioStm->codec, m_threadCount);
    
    AVCodecContext* c = m_audioStm->codec;
    c->codec_id = (CodecID) codecId;
    c->codec_type = CODEC_TYPE_AUDIO;
    c->bit_rate = 64000;
    c->sample_rate = 48000;
    c->channels = 2;
    return true;
}

bool wxFfmpegMediaEncoder::OpenAudioEncoder() {
    AVCodecContext* c = m_audioStm->codec;

    // find the audio encoder and open it
    AVCodec* codec = avcodec_find_encoder(c->codec_id);
    if (!codec)
    {
        wxLogError(wxT("Audio codec not found"));
        return false;
    }
    if (avcodec_open(c, codec) < 0)
    {
        wxLogError(wxT("Could not open audio codec"));
        return false;
    }

    m_audioOutbuf = (uint8_t*) av_malloc(AUDIO_BUF_SIZE);
    
    int audioInputFrameSize = c->frame_size;
    // ugly hack for PCM codecs (will be removed ASAP with new PCM
    // support to compute the input frame size in samples
    if (c->frame_size <= 1)
    {
        audioInputFrameSize = AUDIO_BUF_SIZE / c->channels;
        switch(m_audioStm->codec->codec_id) {
        case CODEC_ID_PCM_S16LE:
        case CODEC_ID_PCM_S16BE:
        case CODEC_ID_PCM_U16LE:
        case CODEC_ID_PCM_U16BE:
            audioInputFrameSize >>= 1;
            break;
        default:
            break;
        }
    }
    m_samples = (int16_t*) av_malloc(audioInputFrameSize * 2 * c->channels);
    int16_t *q = m_samples;
    for (int j = 0; j < audioInputFrameSize * c->channels; j++)
      *q++ = 0;
    return true;
}

void wxFfmpegMediaEncoder::CloseAudioEncoder() {
    if (!m_audioStm)
        return;
    avcodec_close(m_audioStm->codec);
    if (m_samples)
        av_free(m_samples);
    m_samples = NULL;
    if (m_audioOutbuf)
        av_free(m_audioOutbuf);
    m_audioOutbuf = NULL;
    m_audioStm = NULL;
}

AVFrame* allocPicture(PixelFormat pix_fmt, int width, int height) {
    AVFrame* picture = avcodec_alloc_frame();
    if (!picture)
        return NULL;
    int size = avpicture_get_size(pix_fmt, width, height);
    uint8_t* picture_buf = (uint8_t*) av_malloc(size);
    if (!picture_buf) {
        av_free(picture);
        return NULL;
    }
    avpicture_fill((AVPicture *)picture, picture_buf, pix_fmt, width, height);
    return picture;
}

bool wxFfmpegMediaEncoder::OpenVideoEncoder() {
    AVCodecContext* c = m_videoStm->codec;

    // find the video encoder and open it
    AVCodec* codec = avcodec_find_encoder(c->codec_id);
    if (!codec) {
        wxLogError(wxT("Video codec not found"));
        return false;
    }
    if (avcodec_open(c, codec) < 0) {
        wxLogError(wxT("Could not open video codec"));
        return false;
    }

    m_videoOutbuf = (uint8_t*) av_malloc(VIDEO_BUF_SIZE);

    // allocate the encoded raw picture
    m_picture = allocPicture(c->pix_fmt, c->width, c->height);
    if (!m_picture)
    {
        wxLogError(wxT("Could not allocate picture"));
        return false;
    }
    
    m_imgConvertCtx = sws_getContext(c->width, c->height, PIX_FMT_RGB24,
        c->width, c->height, c->pix_fmt, SWS_BICUBIC, NULL, NULL, NULL);
    if (!m_imgConvertCtx)
    {
        wxLogError(wxT("Cannot initialize the conversion context"));
        return false;
    }
    return true;
}

void wxFfmpegMediaEncoder::CloseVideoEncoder() {
    if (!m_videoStm)
        return;
    avcodec_close(m_videoStm->codec);
    if (m_imgConvertCtx)
        sws_freeContext(m_imgConvertCtx);
    if (m_picture)
    {
        av_free(m_picture->data[0]);
        av_free(m_picture);
        m_picture = NULL;
    }
    if (m_videoOutbuf)
        av_free(m_videoOutbuf);
    m_videoOutbuf = NULL;
    m_videoStm = NULL;
}

bool wxFfmpegMediaEncoder::EncodeImage(wxImage image, int frames, bool* canceled) {
	// convert wxImage to the codec pixel format
	AVCodecContext *c = m_videoStm->codec;
	uint8_t *rgbSrc[3]= {image.GetData(), NULL, NULL};
	int rgbStride[3]={3*image.GetWidth(), 0, 0}; 
	sws_scale(m_imgConvertCtx, rgbSrc, rgbStride, 0, c->height, m_picture->data, m_picture->linesize);
	// encode audio and video
	double duration = ((double) m_videoStm->pts.val) * m_videoStm->time_base.num / m_videoStm->time_base.den
		+ ((double) frames) * m_videoStm->codec->time_base.num / m_videoStm->codec->time_base.den;
	while (true) {
		double audioPts = m_audioStm ? ((double) m_audioStm->pts.val) * m_audioStm->time_base.num / m_audioStm->time_base.den : 0.0;
		double videoPts = ((double) m_videoStm->pts.val) * m_videoStm->time_base.num / m_videoStm->time_base.den;
		wxYieldIfNeeded();
		if (canceled != NULL && *canceled)
			return false;
		
		if ((!m_audioStm || audioPts >= duration) && (!m_videoStm || videoPts >= duration))
			break;
	
		// write interleaved audio and video frames
		if (m_audioStm && audioPts < videoPts) {
			if (!writeAudioFrame())
				return false;
		} else {
			if (!writeVideoFrame())
				return false;
		}
	}
	return true;
}

bool wxFfmpegMediaEncoder::writeAudioFrame() {
    AVPacket pkt;
    av_init_packet(&pkt);

    AVCodecContext* c = m_audioStm->codec;
    pkt.size= avcodec_encode_audio(c, m_audioOutbuf, AUDIO_BUF_SIZE, m_samples);
    pkt.pts= av_rescale_q(c->coded_frame->pts, c->time_base, m_audioStm->time_base);
    pkt.flags |= PKT_FLAG_KEY;
    pkt.stream_index= m_audioStm->index;
    pkt.data= m_audioOutbuf;

    // write the compressed frame in the media file
    if (av_write_frame(m_outputCtx, &pkt) != 0) {
        wxLogError(wxT("Error while writing audio frame"));
        return false;
    }
    return true;
}

bool wxFfmpegMediaEncoder::writeVideoFrame() {
    AVCodecContext *c = m_videoStm->codec;

    // encode the image
    int out_size = avcodec_encode_video(c, m_videoOutbuf, VIDEO_BUF_SIZE, m_picture);
    if (out_size < 0)
      return false;
    // if zero size, it means the image was buffered
    if (out_size > 0) {
        AVPacket pkt;
        av_init_packet(&pkt);

        pkt.pts= av_rescale_q(c->coded_frame->pts, c->time_base, m_videoStm->time_base);
        if(c->coded_frame->key_frame)
            pkt.flags |= PKT_FLAG_KEY;
        pkt.stream_index = m_videoStm->index;
        pkt.data = m_videoOutbuf;
        pkt.size = out_size;

        // write the compressed frame in the media file
        int ret = av_write_frame(m_outputCtx, &pkt);
        if (ret != 0) {
            wxLogError(wxT("Error while writing video frame"));
            return false;
        }
    }
    return true;
}

void wxFfmpegMediaEncoder::EndEncode() {
    CloseVideoEncoder();
    CloseAudioEncoder();
    
    if (!m_outputCtx)
        return;
    
    // write the trailer
    av_write_trailer(m_outputCtx);

    // free the streams
    for (unsigned int i = 0; i < m_outputCtx->nb_streams; i++) {
        av_freep(&m_outputCtx->streams[i]->codec);
        av_freep(&m_outputCtx->streams[i]);
    }
    
    // close the output file
#if LIBAVFORMAT_VERSION_INT >= (52<<16)
    url_fclose(m_outputCtx->pb);
#else
    url_fclose(&m_outputCtx->pb);
#endif
    
    // free the stream
    av_free(m_outputCtx);
    m_outputCtx = NULL;
}
