#include "base/sys_string_conversions.h"
#include "database/PA_SqliteHelper.h"
#include "utils/PA_ProcessLock.h"
#include "utils/PA_Helpers.h"
#include "utils/PA_StringUtils.h"
#include "MPEG2TSRecorderSink.h"


MPEG2TSRecorderSink::MPEG2TSRecorderSink( void )
{
    Init();
}

MPEG2TSRecorderSink::~MPEG2TSRecorderSink( void )
{
    if(IsOpened())
        Close();
}


int MPEG2TSRecorderSink::Open( const EncoderConfig& encoder_config, const std::string& file_path )
{
    encoder_config_ = encoder_config;
    file_path_ = file_path;
    if(CreateContext() != PA_OK)
    {
        PA_WarnLog("CreateMp4Context failed!");
        return PA_ERROR;
    }

    return PA_OK;
}



int MPEG2TSRecorderSink::CreateContext()
{
    avformat_alloc_output_context2(&output_ctx_, NULL, NULL, file_path_.c_str());
    if (!output_ctx_)
    {
        PA_WarnLog("Could not deduce output format from file extension: using MPEG.\n");
        avformat_alloc_output_context2(&output_ctx_, NULL, "mpegts", file_path_.c_str());
    }

    if (!output_ctx_)
    {
        return PA_ERROR;
    }

    output_ctx_->oformat->video_codec = AV_CODEC_ID_MPEG4;
    output_ctx_->oformat->audio_codec = CODEC_ID_AAC;

    output_format_ = output_ctx_->oformat;

    if (avio_open(&output_ctx_->pb, file_path_.c_str(), AVIO_FLAG_READ_WRITE) < 0)
    {
        PA_WarnLog("avio_open failed \n");
        goto label_error;
    }

    if (output_format_->video_codec != CODEC_ID_NONE)
    {
        video_steam_ = AddVideoStream(output_ctx_,
                                          output_format_->video_codec,
                                          encoder_config_.video_bitrate_,
                                          encoder_config_.width_,
                                          encoder_config_.height_,
                                          encoder_config_.frame_rate_);
    }

    if (output_format_->audio_codec != CODEC_ID_NONE)
    {
        audio_stream_ = AddAudioStream(output_ctx_,
                                       output_format_->audio_codec,
                                       encoder_config_.audio_bitrate_,
                                       encoder_config_.sample_rate_,
                                       encoder_config_.channels_);
    }

    if (NULL == video_steam_ || NULL == audio_stream_)
    {
        PA_WarnLog("NULL == video_steam_ || NULL == audio_stream_");
        goto label_error;
    }

    av_dump_format(output_ctx_, 0, file_path_.c_str(), 1);

    if (OpenVideo(output_ctx_, video_steam_) == -1)
    {
        PA_WarnLog("NULL == video_steam_ || NULL == audio_stream_");
        goto label_error;
    }

    if (OpenAudio(output_ctx_, audio_stream_))
    {
        PA_WarnLog("NULL == video_steam_ || NULL == audio_stream_");
        goto label_error;
    }

    if (avformat_write_header(output_ctx_, NULL))
    {
        PA_WarnLog("Call avformat_write_header function failed! \n");
        goto label_error;
    }

    return PA_OK;

label_error:

    if (NULL != video_steam_)
    {
        if(video_steam_->codec != NULL)
        {
            avcodec_close(video_steam_->codec);
        }

        video_steam_ = NULL;
    }

    if (NULL != audio_stream_)
    {
        if(audio_stream_->codec != NULL)
        {
            avcodec_close(audio_stream_->codec);
        }
        audio_stream_ = NULL;
    }

    for(unsigned i = 0; i < output_ctx_->nb_streams; i++)
    {
        av_freep(&output_ctx_->streams[i]->codec);
        av_freep(&output_ctx_->streams[i]);
    }

    if(output_ctx_ != NULL)
    {
        if (output_ctx_->pb != NULL)
        {
            avio_close(output_ctx_->pb);
        }
        av_free(output_ctx_);
        output_ctx_ = NULL;
    }

    output_format_ = NULL;

    return PA_ERROR;
}

AVStream* MPEG2TSRecorderSink::AddVideoStream(AVFormatContext *oc, enum AVCodecID codec_id , int bit_rate , int width , int height , int frame_rate)
{
    AVCodecContext *c;
    AVStream *new_stream;
    AVCodec *codec;

    new_stream = avformat_new_stream(oc, NULL);
    if (!new_stream)
    {
        PA_WarnLog("Could not alloc stream\n");
        return NULL;
    }
    new_stream->id = 0;

    c = new_stream->codec;

    /* find the video encoder */
    codec = avcodec_find_encoder(codec_id);
    if (!codec)
    {
        PA_WarnLog("add_video_stream_v codec not found\n");
        av_freep(new_stream);
        return NULL;
    }
    avcodec_get_context_defaults3(c, codec);

    c->codec_id = codec_id;
    c->codec_type = AVMEDIA_TYPE_VIDEO;

    /* put sample parameters */
    c->bit_rate = /*400000*/bit_rate;
    /* resolution must be a multiple of two */
    c->width = width;
    c->height = height;

    /* time base: this is the fundamental unit of time (in seconds) in terms
     of which frame timestamps are represented. for fixed-fps content,
     timebase should be 1/framerate and timestamp increments should be
     identically 1. */
    c->time_base.den = frame_rate;
    c->time_base.num = 1;
    c->gop_size = 12; /* emit one intra frame every twelve frames at most */
    c->pix_fmt = PIX_FMT_YUV420P;
    if (c->codec_id == CODEC_ID_MPEG2VIDEO)
    {
        /* just for testing, we also add B frames */
        c->max_b_frames = 2;
    }
    if (c->codec_id == CODEC_ID_MPEG1VIDEO){
        /* Needed to avoid using macroblocks in which some coeffs overflow.
         This does not happen with normal video, it just happens here as
         the motion of the chroma plane does not match the luma plane. */
        c->mb_decision=2;
    }
    // some formats want stream headers to be separate
    if (oc->oformat->flags & AVFMT_GLOBALHEADER)
        c->flags |= CODEC_FLAG_GLOBAL_HEADER;

    return new_stream;
}

AVStream* MPEG2TSRecorderSink::AddAudioStream(AVFormatContext *oc, enum AVCodecID codec_id,
                         int bit_rate, int sample_rate, int channels)
{
    AVCodecContext *c;
    AVStream *new_stream;
    //AVCodec *codec;

    new_stream = avformat_new_stream(oc, NULL);
    if (!new_stream)
    {
        PA_WarnLog("Could not alloc stream\n");
        return NULL;
    }
    new_stream->id = 1;

    c = new_stream->codec;

    /* find the video encoder */
    /*  codec = avcodec_find_encoder(codec_id);
     if (!codec)
     {
         fprintf(stderr, "add_video_stream_v codec not found\n");
         exit(1);
     }
     avcodec_get_context_defaults3(c, codec);*/

    c->codec_id = codec_id;
    c->codec_type = AVMEDIA_TYPE_AUDIO;

    /* put sample parameters */
    c->sample_fmt = AV_SAMPLE_FMT_S16;
    c->bit_rate = bit_rate;
    c->sample_rate = sample_rate;
    c->channels = channels;
    c->profile = FF_PROFILE_AAC_LOW;
    c->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;

    // some formats want stream headers to be separate
    if (oc->oformat->flags & AVFMT_GLOBALHEADER)
        c->flags |= CODEC_FLAG_GLOBAL_HEADER;

    return new_stream;
}

int MPEG2TSRecorderSink::OpenAudio(AVFormatContext *oc, AVStream *st)
{
    AVCodecContext *audio_codec_context;
    AVCodec *codec;

    audio_codec_context = st->codec;

    /* find the audio encoder */
    codec = avcodec_find_encoder(audio_codec_context->codec_id);
    if (!codec)
    {
        PA_WarnLog("codec not found\n");
        return -1;
    }

    /* open it */
    if (avcodec_open2(audio_codec_context, codec, NULL) < 0)
    {
        PA_WarnLog("could not open codec\n");
        return -1;
    }

    return 0;
}


int MPEG2TSRecorderSink::OpenVideo(AVFormatContext *oc, AVStream *st)
{
    AVCodec *codec;
    AVCodecContext *c;

    c = st->codec;

    /* find the video encoder */
    codec = avcodec_find_encoder(c->codec_id);
    if (!codec)
    {
        PA_WarnLog("codec not found\n");
        return -1;
    }

    /* open the codec */
    if (avcodec_open2(c, codec, NULL) < 0)
    {
        PA_WarnLog("could not open codec\n");
        return -1;
    }
    return 0;
}

int MPEG2TSRecorderSink::WriteAudioFrame(char* data , int len , unsigned int timestemp)
{
    if(output_ctx_ == NULL)
        return PA_ERROR;

    AVPacket pkt;
    av_init_packet(&pkt);

    pkt.stream_index= 1;
    pkt.data= (uint8_t *)data;
    pkt.size = len;

    if (last_audio_timestamp_ == 0)
    {
        last_audio_timestamp_  = timestemp;
        pkt.pts = 0;
    }
    else
    {
        pkt.pts = (timestemp - last_audio_timestamp_) / (av_q2d(audio_stream_->time_base)*1000);
    }

    /* write the compressed frame in the media file */
    if (av_write_frame(output_ctx_, &pkt) != 0)
    {
        PA_WarnLog("Error while writing audio frame\n");
        return PA_ERROR;
    }

    return PA_OK;
}

int MPEG2TSRecorderSink::WriteVideoFrame(char* data , int len , bool keyframe , unsigned int timestemp)
{
    if(output_ctx_ == NULL)
        return PA_ERROR;

    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.stream_index = 0;

    if (keyframe)
    {
        pkt.flags |= AV_PKT_FLAG_KEY;
    }

    pkt.data= (uint8_t *)data;
    pkt.size = len;

    if (last_video_timestamp_  == 0)
    {
        last_video_timestamp_ = timestemp;
        pkt.pts = 0;
        video_timestamp_ = 0;
    }
    else
    {
        pkt.pts = (timestemp - last_video_timestamp_) / (av_q2d(video_steam_->time_base)*1000);
        if (video_timestamp_ == pkt.pts)
        {
            pkt.pts += 1;
        }
        video_timestamp_ = pkt.pts;
    }

    /* write the compressed frame in the media file */
    if (av_write_frame(output_ctx_, &pkt) != 0)
    {
        PA_WarnLog("Error while writing audio frame\n");
        return PA_ERROR;
    }

    return PA_OK;
}

void MPEG2TSRecorderSink::Init()
{
    output_format_ = NULL;
    output_ctx_ = NULL;

    last_audio_timestamp_ = 0;
    last_video_timestamp_ = 0;

     video_timestamp_ = 0;;

    audio_stream_ = NULL;
    video_steam_ = NULL;
}

void MPEG2TSRecorderSink::Finalize()
{
    if(output_ctx_ != NULL)
    {
        av_write_trailer(output_ctx_);
    }
}

void MPEG2TSRecorderSink::Close()
{
    if (NULL != video_steam_)
    {
        if(video_steam_->codec != NULL)
        {
            avcodec_close(video_steam_->codec);
        }

        video_steam_ = NULL;
    }

    if (NULL != audio_stream_)
    {
        if(audio_stream_->codec != NULL)
        {
            avcodec_close(audio_stream_->codec);
        }
        audio_stream_ = NULL;
    }


    if(output_ctx_ != NULL)
    {
        for(unsigned i = 0; i < output_ctx_->nb_streams; i++)
        {
            av_freep(&output_ctx_->streams[i]->codec);
            av_freep(&output_ctx_->streams[i]);
        }

        if (output_ctx_->pb != NULL)
        {
            avio_close(output_ctx_->pb);
        }
        av_free(output_ctx_);
        output_ctx_ = NULL;
    }

    output_format_ = NULL;
}

bool MPEG2TSRecorderSink::IsOpened()
{
    return output_format_ != NULL;
}