#pragma once

#include "XMovieCode.h"

#include <QDebug>



#define STREAM_DURATION   5.0
#define STREAM_FRAME_RATE 24
#define STREAM_NB_FRAMES  ((int)(STREAM_DURATION * STREAM_FRAME_RATE))
#define STREAM_PIX_FMT    AV_PIX_FMT_YUV420P /* default pix_fmt */
#define OUTPUT_CODEC AV_CODEC_ID_H264


int XMovieCode::s_frameCount = 0;

XMovieCode::XMovieCode( const char* filename ) :
    m_filename( filename ),
    m_encoder( avcodec_find_encoder( OUTPUT_CODEC ))
{
    av_log_set_level(AV_LOG_VERBOSE);

    int ret(0);

   /** Allocate the output media context */
    ret = avformat_alloc_output_context2( &m_formatCtx, m_outputFormat, "mp4", m_filename.c_str());

    if (!m_formatCtx)
        return;

    m_outputFormat = m_formatCtx->oformat;

    /** Allocate the codec context */
    m_codecCtx = avcodec_alloc_context3( m_encoder);

    /** Add the video stream using H264 codec */
    addStream();

    if( !m_streamOut )
        return;

    configureEncoder();

    // Print detailed information about input and output
//    av_dump_format( m_formatCtx, 0, m_filename.c_str(), 1);

    /** Open the output media file, if needed */
    if (!( m_outputFormat->flags & AVFMT_NOFILE))
    {
        ret = avio_open( &m_formatCtx->pb, m_filename.c_str(), AVIO_FLAG_WRITE);

        checkError( ret, "Could not open file : " );
    }
    else
    {
        return;
    }


    /** Write media header */
    ret = avformat_write_header( m_formatCtx, NULL );

    checkError( ret,"Error occurred when opening output file: ");
}



XMovieCode::~XMovieCode()
{
}

/* Add an output stream. */
void XMovieCode::addStream()
{
    if (!( m_encoder ))
    {
        fprintf(stderr, "Could not find encoder for '%s'\n",
            avcodec_get_name( OUTPUT_CODEC ));
        return;
    }

    /** Get the stream for codec */
    m_streamOut = avformat_new_stream(m_formatCtx, m_encoder);

    if (!m_streamOut) {
        fprintf(stderr, "Could not allocate stream\n");
        return;
    }

    m_streamOut->id = m_formatCtx->nb_streams - 1;


    if( m_encoder->type == AVMEDIA_TYPE_VIDEO )
    {
        m_streamOut->codecpar->codec_id = OUTPUT_CODEC;
        m_streamOut->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
        m_streamOut->codecpar->bit_rate = 400000;
        m_streamOut->codecpar->width = 800;
        m_streamOut->codecpar->height = 640;
        m_streamOut->codecpar->format = STREAM_PIX_FMT;
        m_streamOut->time_base = { 1, STREAM_FRAME_RATE };

        avcodec_parameters_to_context( m_codecCtx, m_streamOut->codecpar);
    }

    if (m_streamOut->codecpar->codec_id == OUTPUT_CODEC)
    {
//      av_opt_set( &m_opts, "preset", "slow", 0 );
      av_opt_set( &m_opts, "preset", "ultrafast", 0 );
    }
}


void XMovieCode::configureEncoder()
{
    m_codecCtx->gop_size = 0;
    m_codecCtx->max_b_frames = 0;
    m_codecCtx->time_base = { 1, STREAM_FRAME_RATE };
    m_codecCtx->framerate = { STREAM_FRAME_RATE, 1 };
    m_codecCtx->pix_fmt = STREAM_PIX_FMT;
//    m_codecCtx->max_samples = STREAM_DURATION * STREAM_FRAME_RATE;

    int ret = avcodec_parameters_from_context( m_streamOut->codecpar, m_codecCtx );

    checkError( ret, "avcodec_parameters_from_context returned : " );

    /* open the codec */
    ret = avcodec_open2(m_codecCtx, m_encoder, &m_opts );

    checkError( ret, "Could not open video codec: ");

    ret = avcodec_parameters_from_context( m_streamOut->codecpar, m_codecCtx );

    checkError( ret,"avcodec_parameters_from_context returned : " );

    m_streamOut->time_base = m_codecCtx->time_base;
}



void XMovieCode::encodeImage(const Image &image)
{
    if( m_setStartTime )
    {
        m_startTime = image.timeStamp();
        m_setStartTime = false;
    }

    if (!m_streamOut )
       return;

    AVFrame* frame = av_frame_alloc();

    createFrame( image, frame );

    frame->pts = int64_t( 1000 * s_frameCount / STREAM_FRAME_RATE );

    /**
     * Add a video frame
     */
    writeFrame( frame );

    if( frame )
    {
        av_free( frame->data[0]);
        av_frame_free( &frame );
    }
}



void XMovieCode::writeFrame( AVFrame* frame )
{
    int ret;


    if ( s_frameCount >= STREAM_NB_FRAMES)
    {
        /* No more frames to compress.*/
        s_frameCount = STREAM_NB_FRAMES;
    }


    if (m_formatCtx->oformat->flags & 0x0020 )
    {
        /* Raw video case - directly store the picture in the packet */
        AVPacket pkt;
        av_init_packet(&pkt);

        pkt.flags |= AV_PKT_FLAG_KEY;
        pkt.stream_index = m_streamOut->index;
        pkt.data = frame->data[0];
        pkt.size = sizeof(AVPicture);

        ret = av_write_frame( m_formatCtx, &pkt );
    }
    else
    {
        AVPacket pkt;
        av_init_packet(&pkt);

        /* encode the image */
        ret = avcodec_send_frame(m_codecCtx, frame);

        checkError( ret, "Error encoding video frame: ");

        while( ret >= 0 )
        {
            /* If size is zero, it means the image was buffered. */
            ret = avcodec_receive_packet(m_codecCtx, &pkt);

            if (ret == AVERROR(EAGAIN) /* || ret == AVERROR_EOF*/)
            {
                return;
            }
            else if (ret < 0)
            {
qDebug() << "Error during encoding\n";
                  return;
            }

            pkt.stream_index = m_streamOut->index;

            ret = av_write_frame( m_formatCtx, &pkt );
			
            s_frameCount++;
        }
    }

}


void XMovieCode::createFrame( const Image& image, AVFrame* frame )
{
//image.saveToFile("c:\\tmp\\RGBoriginal.png");

    frame->format = STREAM_PIX_FMT;
    frame->width = image.width();
    frame->height = image.height();
    frame->pict_type = AV_PICTURE_TYPE_I;
    frame->display_picture_number = image.uniqueImageNumber();

    int ret = av_image_alloc( frame->data, frame->linesize, frame->width,  frame->height, STREAM_PIX_FMT, 1);

    if (ret < 0)
    {
        return;
    }

    struct SwsContext* sws_ctx = sws_getContext((int)image.width(), (int)image.height(), AV_PIX_FMT_BGR24,
                                                (int)image.width(), (int)image.height(), STREAM_PIX_FMT, 0, NULL, NULL, NULL);

    const uint8_t* rgbData[1] = { (uint8_t* )image.getData() };
    int rgbLineSize[1] = { 3 * image.width() };

    sws_scale(sws_ctx, rgbData, rgbLineSize, 0, image.height(), frame->data, frame->linesize);

//    saveFrameToDisk( frame );
}


void XMovieCode::close()
{
qDebug() << "close ";

    /** reset the framecount */
    s_frameCount = 0 ;

    int ret( 0 );

    /** flush the encoder */
    while( ret >= 0 )
        ret = avcodec_send_frame(m_codecCtx, NULL);
qDebug() << "flushed codec ";

    // Write media trailer
    if( m_formatCtx )
        ret = av_write_trailer( m_formatCtx );

    /* Close the output file. */
    if (!( m_outputFormat->flags & AVFMT_NOFILE))
        ret = avio_close( m_formatCtx->pb);


    /* free the stream */
    avformat_free_context( m_formatCtx );

    fflush( stdout );
}



void XMovieCode::checkError( const int errID, const std::string& message )
{
    if( errID != 0 )
    {
        char error[255];
        av_strerror( errID, error, 255);
        std::string err( message );
        err.append( error );
        Log::printLine( err );
qDebug() << err.c_str();
        return;
    }
}


void XMovieCode::saveFrameToDisk( const AVFrame* frame  )
{
    cv::Mat yuv420p( frame->height + frame->height/2, frame->width, CV_8UC1, frame->data[0]);
    cv::Mat cvmIm;
    cv::cvtColor(yuv420p,cvmIm,CV_YUV420p2BGR);
    std::ostringstream ss;
    ss << "c:\\tmp\\YUVoriginal_" << frame->display_picture_number << ".png";
    cv::imwrite( ss.str().c_str(), cvmIm);
}
