/////////////////////////////////////////////////////////////////////////////
// Name:        mediaenc_ffmpeg.cpp
// Purpose:     FFMPEG Media Encoder
// Author:      Alex Thuering
// Created:     04.08.2007
// RCS-ID:      $Id: mediaenc_ffmpeg.cpp,v 1.7 2008/02/25 19:22:51 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
#include <ffmpeg/avformat.h>
#include <ffmpeg/swscale.h>
}

#define AUDIO_BUF_SIZE 50000
#define VIDEO_BUF_SIZE 1835008

wxFfmpegMediaEncoder::wxFfmpegMediaEncoder() {
	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();
}

bool wxFfmpegMediaEncoder::BeginEncode(const wxString& fileName, VideoFormat videoFormat,
		AudioFormat audioFormat, 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;
    
    m_outputCtx = av_alloc_format_context();
    if (!m_outputCtx) {
        wxLogError(wxT("memmory 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, 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 and allocate the necessary encode buffers
    if (m_videoStm)
        OpenVideoEncoder();
    if (m_audioStm)
        OpenAudioEncoder();
    
    // open the output file
    if (url_fopen(&m_outputCtx->pb, fileName.mb_str(), URL_WRONLY) < 0) {
        wxLogError(wxString::Format(wxT("Could not open '%s'"), fileName.c_str()));
        return false;
    }
    // write the stream header
    m_outputCtx->packet_size = 2048;
    m_outputCtx->mux_rate = 10080000;
    av_write_header(m_outputCtx);
    return true;
}

bool wxFfmpegMediaEncoder::addVideoStream(int codecId, VideoFormat videoFormat, 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;
    }

    AVCodecContext* c = m_videoStm->codec;
    c->codec_id = (CodecID) codecId;
    c->codec_type = CODEC_TYPE_VIDEO;
    c->bit_rate = videoBitrate * 1024;
    c->width = 720;
    c->height = videoFormat == vfPAL ? 576 : 480;
    c->time_base.den = videoFormat == vfPAL ? 25 : 30000;
    c->time_base.num = videoFormat == vfPAL ? 1 : 1001;
    c->gop_size = videoFormat == vfPAL ? 15 : 18;
    c->pix_fmt = PIX_FMT_YUV420P;
    c->rc_buffer_size = VIDEO_BUF_SIZE;
    c->rc_max_rate = 9 * 1024 * 1024;
    c->rc_min_rate = 0;
    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;
    }

    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(int 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) {
	// 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);
	
	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;
		//wxLogError(wxString::Format(wxT("%d %f %f %f"), frameCount, duration, audioPts, videoPts));
		
		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;
}
