#include <stdio.h>

#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/opt.h>
#include <libswscale/swscale.h>

int main (int argc, char **argv)
{
    if (argc !=3)
    {
        fprintf(stderr, "Error: invalid arguments\n");
        fprintf(stderr, "\tCorrect usage: %s <input-filename> <output-filename>\n", argv[0]);
        exit(1);
    }

    // initialize all formats and codecs
    av_register_all();

    /// Set up input
    const char *input_filename = argv[1];
    // must be initialized to NULL or otherwise avformat_open_input() will segfault
    AVFormatContext *input_context = NULL;

    // Get input file context
    if ( avformat_open_input(&input_context, input_filename, NULL, NULL) != 0 )
    {
        fprintf(stderr, "Error: Couldn't open input file.\n");
        exit(1);
    }

    // Get streams information
    if( avformat_find_stream_info(input_context, NULL) < 0 )
    {
        fprintf(stderr, "Error: Couldn't find stream information.\n");
        exit(1);
    }

    // debug: dump informations about input to stdout
    av_dump_format(input_context, 0, input_filename, 0);


    // Find the first video stream
    int i, video_stream;
    video_stream=-1;

    for(i=0; i<input_context->nb_streams; i++)
    {
        if(input_context->streams[i]->codec->codec_type==AVMEDIA_TYPE_VIDEO)
        {
            video_stream=i;
            break;
        }
    }
    if(video_stream==-1)
    {
        fprintf(stderr, "Error: Couldn't find a video stream.\n");
        exit(1);
    }

    // Get codec context for selected stream
    AVCodecContext *iv_codec_ctx = input_context->streams[video_stream]->codec;
    // Find the decoder for the video stream
    AVCodec *v_decoder  = avcodec_find_decoder ( iv_codec_ctx->codec_id );
    if(v_decoder==NULL)
    {
        fprintf(stderr, "Error: Video decoder not found.\n");
        exit(1);
    }

    // Inform the codec that we can handle truncated bitstreams -- i.e.,
    // bitstreams where frame boundaries can fall in the middle of packets
    if(v_decoder->capabilities & CODEC_CAP_TRUNCATED)
        iv_codec_ctx->flags|=CODEC_FLAG_TRUNCATED;

    // Open decoder with desired setup
    if ( avcodec_open(iv_codec_ctx, v_decoder) < 0 )
    {
        fprintf(stderr, "Error: Couldn't open video decoder.\n");
        exit(1);
    }

    /// Set up output
    const char *output_filename = argv[2];
    AVFormatContext *output_context;

    // Alloc and create output file context
    //   tell avcodec that it will contain h264 video stream
    avformat_alloc_output_context2(&output_context, NULL, "h264", NULL);
    if (!output_context)
    {
        fprintf(stderr, "Error: Couldn't create context for output file.\n");
        exit(1);
    }

    //open output file
    if ( avio_open(&output_context->pb, output_filename, AVIO_FLAG_WRITE) < 0 )
    {
        fprintf(stderr, "Error: Could not open '%s'\n", output_filename);
        exit(1);
    }

    // find coder for h264
    AVCodec *v_coder = avcodec_find_encoder(CODEC_ID_H264);
    if (!v_coder)
    {
        fprintf(stderr, "Error: Could not found h264 codec.\n");
        exit(1);
    }

    // add new stream encoded by this coder to ouput context
    AVStream *ov_stream = avformat_new_stream(output_context, v_coder);
    if (!ov_stream)
    {
        fprintf(stderr, "Error: Could not alloc stream.\n");
        exit(1);
    }

    // set up encoding parameters (codec context)
    AVCodecContext *ov_codec_ctx = ov_stream->codec;
    ov_codec_ctx->codec_id = CODEC_ID_H264;
    ov_codec_ctx->codec_type = AVMEDIA_TYPE_VIDEO;
    /* put sample parameters */
    ov_codec_ctx->bit_rate = 400000;
    /* resolution must be a multiple of two */
    ov_codec_ctx->width = 352;
    ov_codec_ctx->height = 288;
    /* frames per second */
    ov_codec_ctx->time_base= //(AVRational){1,25};
        iv_codec_ctx->time_base;
    ov_codec_ctx->gop_size = 10; /* emit one intra frame every ten frames */
    ov_codec_ctx->max_b_frames=2;
    ov_codec_ctx->pix_fmt = PIX_FMT_YUV420P;

    // set h264 preset
    if ( av_opt_set(ov_codec_ctx->priv_data, "preset", "slow", 0) )
    {
        fprintf(stderr, "Error: Could not set h264 preset.\n");
        exit(1);
    }

    // debug: dump informations about output to stdout
    av_dump_format(output_context, 0, output_filename, 1);

    // open h264 codec
    if (avcodec_open(ov_codec_ctx, v_coder) < 0)
    {
        fprintf(stderr, "Error: Could not open h264 codec.\n");
        exit(1);
    }

    // create output buffer for video
    int video_outbuf_size = 200000;
    uint8_t *video_outbuf = av_malloc(video_outbuf_size);

    // allocate output picture
    AVFrame *output_picture = avcodec_alloc_frame();
    int size = avpicture_get_size(ov_codec_ctx->pix_fmt, ov_codec_ctx->width, ov_codec_ctx->height);
    uint8_t * output_picture_buf = av_malloc(size);
    if (!output_picture || !output_picture_buf)
    {
        fprintf(stderr, "Error: Could not allocate output picture.\n");
        exit(1);
    }
    avpicture_fill((AVPicture *)output_picture, output_picture_buf,
                   ov_codec_ctx->pix_fmt, ov_codec_ctx->width, ov_codec_ctx->height);

    // create picture conversion context
    struct SwsContext *img_convert_ctx =
            sws_getContext( iv_codec_ctx->width, iv_codec_ctx->height, iv_codec_ctx->pix_fmt,
                            ov_codec_ctx->width, ov_codec_ctx->height, ov_codec_ctx->pix_fmt,
                            SWS_BICUBIC, NULL, NULL, NULL);



    AVFrame *input_picture = avcodec_alloc_frame();
    int got_picture;
    int len, out_size;
    int picture_counter = 0;
    AVPacket packet;
    AVPacket output_packet;
    //av_init_packet(&packet);
    uint8_t * packet_start;


    /* write the output stream header, if any */
    av_write_header(output_context);

    while ( av_read_frame(input_context, &packet) >= 0 )
    // grab packet from input file
    {
        if ( packet.stream_index == video_stream )
        // care only for packets from desired video stream
        {
            // store ptr to packet data start
            packet_start = packet.data;
            while (packet.size > 0)
            // decode data from packet
            {
                len = avcodec_decode_video2( iv_codec_ctx, input_picture, &got_picture, &packet);
                //printf("[video packet]%d bytes decoded from packet, got_picture = %d\n", len, got_picture);
                if (len<0)
                {
                    fprintf(stderr, "Error: While decoding picture.\n");
                    exit(1);
                }
                if (got_picture)
                {
                    // handle picture
                    picture_counter++;
                    // convert img to output img format
                    sws_scale(img_convert_ctx, &input_picture->data, input_picture->linesize,
                        0, ov_codec_ctx->height, &output_picture->data, output_picture->linesize);

                    // send image for encoding
                    out_size = avcodec_encode_video(ov_codec_ctx, video_outbuf, video_outbuf_size, output_picture);
                    printf("encoding frame %3d (size=%5d)\n", picture_counter, out_size);
                    // encoder returned some data, lets write them into output file
                    if (out_size>0)
                    {
                        av_init_packet(&output_packet);
                        // todo: co je tohle ??
                        if (ov_codec_ctx->coded_frame->pts != AV_NOPTS_VALUE)
                            output_packet.pts=
                                av_rescale_q(ov_codec_ctx->coded_frame->pts, ov_codec_ctx->time_base, ov_stream->time_base);
                        printf("pts = %lf,   ", (double)ov_stream->pts.val * ov_stream->time_base.num / ov_stream->time_base.den);
                        printf("%d %d %d\n", ov_codec_ctx->coded_frame->pts, ov_codec_ctx->time_base, ov_stream->time_base);
                        //output_packet.dts = AV_NOPTS_VALUE;
                        if(ov_codec_ctx->coded_frame->key_frame)
                            output_packet.flags |= AV_PKT_FLAG_KEY;
                        output_packet.stream_index = ov_stream->index;
                        output_packet.data = video_outbuf;
                        output_packet.size = out_size;

                        /* write the compressed frame in the media file */
                        if ( av_interleaved_write_frame(output_context, &output_packet)!= 0 )
                        {
                            fprintf(stderr, "Error: while writing video frame.\n");
                            exit(1);
                        }

                    }
                }
                packet.size -= len;
                packet.data += len;
            }
            // restore ptr to packet data start so it can get released correctly by AV
            packet.data = packet_start;
        }
        av_free_packet(&packet);    // free packet
    }

    printf("video frames grabbed = %d \n", picture_counter);

    // write output trailer if any
    av_write_trailer(output_context);

    // clean resources

    // free output
    avcodec_close(ov_stream->codec);
    av_free(output_picture);
    av_free(output_picture->data[0]);
    av_free(video_outbuf);

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

    // free input
    avcodec_close(iv_codec_ctx);
    av_close_input_file(input_context);
    av_free(input_picture);

    return 0;
}