/*
 *  FFMPEGEncoder.cpp
 *  live555 Test
 *
 *  Created by Severin Schoepke on 05.06.07.
 *  Copyright 2007 Neotrivium AG. All rights reserved.
 *
 */

#include "FFMPEGEncoder.hh"

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <sys/time.h>
#include <sys/resource.h>

//#define WRITE_TO_FILE
//#define BENCHMARK

class Benchmark
{
public:
    Benchmark() {};
    virtual ~Benchmark() {};
    
    void start();
    void stop();
    void print(const char *title);

private:
    struct rusage _start, _stop;
};

void Benchmark::start()
{
  getrusage(RUSAGE_SELF, &_start);
}  

void Benchmark::stop()
{
  getrusage(RUSAGE_SELF, &_stop);
}

void Benchmark::print(const char *title)
{
  int userSecs  = _stop.ru_utime.tv_sec - _start.ru_utime.tv_sec; 
  int userUsecs = _stop.ru_utime.tv_usec - _start.ru_utime.tv_usec; 
  int sysSecs   = _stop.ru_stime.tv_sec - _start.ru_stime.tv_sec; 
  int sysUsecs  = _stop.ru_stime.tv_usec - _start.ru_stime.tv_usec;
  
  printf("%s: user time = %i:%i, system time = %i:%i\n", title, userSecs, userUsecs, sysSecs, sysUsecs);
}

FFMPEGEncoder::FFMPEGEncoder(const char *format, 
                             const char *vcodec, 
                             const char *acodec, 
                             int width, 
                             int height, 
                             int videoBitRate, 
                             int audioBitRate, 
                             int sampleRate, 
                             int numChannels)
{
    av_register_all();
    
    // search output format
    _fmt = NULL;
    _fmt = guess_format(format, NULL, NULL);
    if(!_fmt) {
        printf("output format not found.\n");
        exit(1);
    }
    
    // allocate the output media context
    _oc = av_alloc_format_context();
    if (!_oc) {
        printf("could not alloc memory.\n");
        exit(1); 
    }
    _oc->oformat        = _fmt;
    
#ifdef WRITE_TO_FILE
    snprintf(_oc->filename, sizeof(_oc->filename), "dump.%s", format);
#else
    _oc->oformat->flags |= AVFMT_NOFILE;
#endif
    
    // add video stream
    _video_st = av_new_stream(_oc, 0);
    if (!_video_st) {
        printf("could not alloc video stream.\n");
        exit(1); 
    }
    
    AVCodecContext *video_codec;
    
    video_codec               = _video_st->codec;
    video_codec->codec_id     = _fmt->video_codec;
    video_codec->codec_type   = CODEC_TYPE_VIDEO;
    
    video_codec->bit_rate     = videoBitRate;
    video_codec->width        = width;
    video_codec->height       = height;
    video_codec->time_base    = (AVRational){1,25}; //frame rate
    video_codec->gop_size     = 12;                 //num frames between I-frames   
    video_codec->pix_fmt      = PIX_FMT_YUV420P;
    
    if(!strcmp(_oc->oformat->name, "mp4") || !strcmp(_oc->oformat->name, "mov") || !strcmp(_oc->oformat->name, "3gp"))
        video_codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
    
    // add audio stream
    _audio_st = av_new_stream(_oc, 0);
    if (!_audio_st) {
        printf("could not alloc audio stream.\n");
        exit(1); 
    }
    
    AVCodecContext *audio_codec;
    
    audio_codec                 = _audio_st->codec;
    audio_codec->codec_id       = _fmt->audio_codec;
    audio_codec->codec_type     = CODEC_TYPE_AUDIO;
    
    audio_codec->bit_rate       = audioBitRate;
    audio_codec->sample_rate    = sampleRate;
    audio_codec->channels       = numChannels;
    
    // set the output parameters (must be done even if no parameters).
    if (av_set_parameters(_oc, NULL) < 0) {
        printf("invalid output format parameters.\n");
        exit(1); 
    }
    
    // prepare video codec
    AVCodec *video;
    video = avcodec_find_encoder_by_name(vcodec);
    if (!video) {
        printf("could not find video codec (%s).\n", vcodec);
        exit(1); 
    }
    
    if (avcodec_open(video_codec, video) < 0) {
        printf("could not open video codec.\n");
        exit(1); 
    }
    
    _tmp_frame = avcodec_alloc_frame();    
    if (! _tmp_frame) {
        printf("Could not alloc frame.\n");
        exit(1);
    }
    
    _video_outbuf_size   = 500000;
    _video_outbuf        = (uint8_t*)av_malloc(_video_outbuf_size);
    
    // prepare audio codec
    AVCodec *audio;
    audio = avcodec_find_encoder_by_name(acodec);
    if (!audio) {
        printf("could not find audio codec (%s).\n", acodec);
        exit(1); 
    }
    
    if (avcodec_open(audio_codec, audio) < 0) {
        printf("could not open audio codec.\n");
        exit(1); 
    }
    
    _audio_outbuf_size  = 500000;
    _audio_outbuf       = (uint8_t*)av_malloc(_audio_outbuf_size);

#ifdef WRITE_TO_FILE
    if (!(_fmt->flags & AVFMT_NOFILE)) {
        if (url_fopen(&_oc->pb, _oc->filename, URL_WRONLY) < 0) {
            printf("%s could not open file %s.\n", _oc->filename);
        }
    }

    dump_format(_oc, 0, _oc->filename, 1);
    av_write_header(_oc);
#else
    dump_format(_oc, 0, "NO FILE", 1);
#endif

    _benchmark = new Benchmark();
}

FFMPEGEncoder::~FFMPEGEncoder()
{
    // close video
    avcodec_close(_video_st->codec);
    av_free(_video_outbuf);
    
    // close audio
    avcodec_close(_audio_st->codec);
    av_free(_audio_outbuf);
    
    // free streams
    unsigned int i;
    for(i = 0; i < _oc->nb_streams; i++) {
        av_freep(&_oc->streams[i]->codec);
        av_freep(&_oc->streams[i]);
    }
    
    av_free(_oc);
}

int FFMPEGEncoder::encodeVideo(void *output, void *input, unsigned int input_size) 
{
#ifdef BENCHMARK
    _benchmark->start();
#endif

    if(input_size > _video_outbuf_size) {
        printf("Too much data, I can't write that.\n");
        return 0;
    }
    
    AVCodecContext *c = _video_st->codec;
    
    avpicture_fill((AVPicture *)_tmp_frame, (uint8_t*)input, c->pix_fmt, c->width, c->height);
    int out_size = avcodec_encode_video(c, _video_outbuf, _video_outbuf_size, _tmp_frame);    
    
    if(out_size == 0) return 0;

#ifdef WRITE_TO_FILE
    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.size            = out_size;
    pkt.pts             = _video_pts;
    pkt.stream_index    = _video_st->index;
    pkt.data            = _video_outbuf;
    if(c->coded_frame->key_frame) 
        pkt.flags       |= PKT_FLAG_KEY;

    if (av_write_frame(_oc, &pkt) != 0) {
        printf("Error writing video frame to file.\n");
    }
#endif
    
    memcpy(output, _video_outbuf, out_size);
    
#ifdef BENCHMARK
    _benchmark->stop();
    _benchmark->print("Video Encoding");
#endif

    return out_size;
}

int FFMPEGEncoder::encodeAudio(void *output, void *input, unsigned int input_size) 
{
#ifdef BENCHMARK
    _benchmark->start();
#endif

    if(input_size > _audio_outbuf_size) {
        printf("Too much data, I can't write that.\n");
        return 0;
    }
    
    AVCodecContext *c = _audio_st->codec;
    
    int out_size = avcodec_encode_audio(c, _audio_outbuf, input_size, (int16_t*)input); 
    
    if(out_size == 0) return 0;

#ifdef WRITE_TO_FILE
    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.size            = out_size;
    pkt.pts             = _audio_pts;
    pkt.stream_index    = _audio_st->index;
    pkt.data            = _audio_outbuf;
    pkt.flags          |= PKT_FLAG_KEY;

    if (av_write_frame(_oc, &pkt) != 0) {
        printf("Error writing audio frame to file.\n");
    }
#endif
    
    memcpy(output, _audio_outbuf, out_size);    
    
#ifdef BENCHMARK
    _benchmark->stop();
    _benchmark->print("Audio Encoding");
#endif

    return out_size;
}

PixelFormat FFMPEGEncoder::pixelFormat()
{
    return _video_st->codec->pix_fmt;
}

int FFMPEGEncoder::width()
{
    return _video_st->codec->width;
}

int FFMPEGEncoder::height()
{
    return _video_st->codec->height;
}

int FFMPEGEncoder::frameRate()
{
    return _video_st->codec->time_base.den;
}

int FFMPEGEncoder::videoBitRate()
{
    return _video_st->codec->bit_rate;
}

int FFMPEGEncoder::audioBitRate()
{
    return _audio_st->codec->bit_rate;
}

int FFMPEGEncoder::audioFrameSize()
{
    return _audio_st->codec->frame_size;
}

int FFMPEGEncoder::sampleRate()
{
    return _audio_st->codec->sample_rate;
}

int FFMPEGEncoder::numChannels()
{
    return _audio_st->codec->channels;
}
