// tutorial01.c
// Code based on a tutorial by Martin Bohme (boehme@inb.uni-luebeckREMOVETHIS.de)
// Tested on Gentoo, CVS version 5/01/07 compiled with GCC 4.1.1

// A small sample program that shows how to use libavformat and libavcodec to
// read video from a file.
//
// Use
//
// gcc -O3 -g -Wall -o dumpframe dumpframe.c -I/usr/local/include/ -L/usr/local/lib/ -lavformat -lavcodec -lz
//
// to build (assuming libavformat and libavcodec are correctly installed
// your system).
//
// Run using
//
// tutorial01 myvideofile.mpg
//
// to write the first five frames from "myvideofile.mpg" to disk in PPM
// format.

#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <getopt.h>


static int frame_start = 0;
static int frame_number = 0xffff;
static int frame_saved = 0;
static int frame_current = 0;
static int frame_step = 1;
static int frame_single = 0; /* write each frame to single files */

static char *src_fn = NULL;
static char *dst_fn = NULL;
static int srcFd = 0;
static int outFd = 0;
static int is_vp8 = 0;

static int dump_frame(int pos, int size)
{
    char fn[1024];
    
    if ((frame_start > frame_current) || (frame_saved >= frame_number)) {
        //printf("==Dumped Frame #%d: Skiped and Not Happend==\n",  frame_current);
        return 0;
    }
    

    if (frame_single) /* dump each frame into a single file */
        sprintf(fn, "%s.%d", dst_fn, frame_saved);
    else
        sprintf(fn, "%s", dst_fn);

    if (outFd == 0) {
        outFd = open(fn, O_TRUNC|O_RDWR|O_CREAT, S_IROTH);
        if (outFd <= 0) {
            printf("Open file %s error (%s)\n", fn, strerror(errno));
            return 0;
        }
    }
    
    if (is_vp8) { /* include the IVF header */
        if (frame_current == 0) {
            pos = pos - 32;
            size = size + 32 + 12;
        } else {
            //pos = pos - 12;
            size = size + 12;
        }
    }
    
    if (outFd > 0 && srcFd > 0 && pos >=0 && size > 0) {
        char *p = malloc(size+4096);
                    
        lseek(srcFd, pos, SEEK_SET);
        read(srcFd, p, size);
        write(outFd, p, size);
        if (frame_single) {
            close(outFd);
            outFd = 0;
        }
        
        free(p);
        
        printf("=====Dumped Frame #%d to File %s (start:%d, size:%d)=====\n",
               frame_current, fn, pos, size);
    } else
        printf("=====Fail to dumped Frame #%d to File %s for Some Reason=====\n",  frame_current, fn);

    frame_saved++;    
    frame_start += frame_step; /* pointer to next frame location*/
    
    return 0;
}

static int process_frame(AVCodecContext *pCodecCtx,
                         AVFrame *pFrame, AVPacket *packet)
{
    int ret, frameFinished;
    
    avcodec_get_frame_defaults(pFrame);            
    /// Decode video frame
    //avcodec_decode_video(pCodecCtx, pFrame, &frameFinished,packet.data, packet.size);
    ret = avcodec_decode_video2(pCodecCtx, pFrame, &frameFinished, packet);

    if (ret < 0)
        return ret;
            
#define FFMIN(a, b) ((a) > (b) ? (b) : (a))
    ret = FFMIN(ret, packet->size);/* guard against bogus return values */
    packet->data += ret;
    packet->size -= ret;
            
    // Did we get a video frame?
    if (frameFinished == 0)
        return 0;
            
    int pos = -1, size = -1;
                
    if (av_frame_get_pkt_pos (pFrame) != -1) {
        pos = (int)av_frame_get_pkt_pos(pFrame);
    } else {
        printf("pkt_pos:N/A\n");
    }
                    
    if (av_frame_get_pkt_size(pFrame) != -1) {
        size = av_frame_get_pkt_size(pFrame);
    } else {
        printf("pkt_size:N/A");
    }

    printf("=====Frame #%d (pkt_pos=%d, pkt_size=%d)====\n", frame_current, pos, size);
    
    dump_frame(pos, size);
    
    frame_current++;
    
    return frameFinished;
}

static void help()
{
    printf("dumpframe -s <frame_start> -n <frame_count> -step <number> -single -v <video> -o <dest file>\n");
    printf("Example: dumpframe -step 2 -v test.vp8 -o out.ivf \n");
    printf("         dump frame 0/2/4/6/... to out.ivf\n");
    printf("Example: dumpframe -s 2 -step 4 -v test.vp8 -o out.ivf \n");
    printf("         dump frame 2/6/10... to out.ivf\n");
    printf("Example: dumpframe  -v /tmp/test.ivf -o out.ivf\n");
    printf("         dump all frame of test.ivf to out.ivf, like a direct file copy\n");
    printf("Example: dumpframe -v test.264 -o out.264 -single\n");
    printf("         dump frame 0 of test.264 into out.264.0\n");
    printf("         dump frame 1 of test.264 into out.264.1\n");
    printf("         dump frame 2 of test.264 into out.264.2\n");
    printf("         .............\n");
}

#pragma mark - Main function
int main(int argc, char *argv[]) {
    AVFormatContext *pFormatCtx=NULL;
    int i, videoStreamIdx;
    AVCodecContext *pCodecCtx;
    AVCodec *pCodec;
    AVFrame *pFrame;
    AVFrame *pFrameRGB;
    AVPacket packets;
    int frameFinished;
    int numBytes;
    uint8_t *buffer;
    static struct SwsContext *img_convert_ctx;
    int long_index;
    char c;

    const struct option long_opts[] = {
        {"single", no_argument, NULL, 0 },
        {"step", required_argument, NULL, 1 }
    };
    
    while ((c =getopt_long_only(argc,argv,"s:n:v:o:h?",long_opts,&long_index)) != EOF) {
        switch (c) {
        case 's':
            frame_start = atoi(optarg);
            break;
        case 'n':
            frame_number = atoi(optarg);
            break;
        case 'v':
            src_fn = strdup(optarg);
            break;
        case 'o':
            dst_fn = strdup(optarg);
            break;
        case 0:
            frame_single = 1;
            break;
        case 1:
            frame_step = atoi(optarg);
            break;
        
        case 'h':
        case ':':
        case '?':
        default:
            help();
            
            return -1;
        }
    }

    if (src_fn == NULL || dst_fn == NULL) {
        help();
        return 0;
    }
        
    // Register all formats and codecs
    av_register_all();
    

    /// Open video file
    //if(av_open_input_file(&pFormatCtx, argv[1], NULL, 0, NULL)!=0) // Deprecated
    if(avformat_open_input(&pFormatCtx, src_fn, NULL, NULL) != 0)
        return -1; // Couldn't open file

    srcFd = open(src_fn, O_RDONLY);
    if (srcFd <=0) {
        printf("Open source file failed\n");
        return 1;
    }
        
    /// Retrieve stream information
    //if(av_find_stream_info(pFormatCtx)<0) // Deprecated
    if(avformat_find_stream_info(pFormatCtx, NULL) < 0)
        return -1; // Couldn't find stream information
    

    /// Dump information about file onto standard error
    //av_dump_format(pFormatCtx, 0, src_fn, 0);
    

    /// Find the first video stream
    videoStreamIdx=-1;
    for(i=0; i<pFormatCtx->nb_streams; i++)
        if(pFormatCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO) { //CODEC_TYPE_VIDEO
            videoStreamIdx=i;
            break;
        }
    if(videoStreamIdx==-1)
        return -1; // Didn't find a video stream
    

    /// Get a pointer to the codec context for the video stream
    pCodecCtx = pFormatCtx->streams[videoStreamIdx]->codec;


    /// Find the decoder for the video stream
    pCodec = avcodec_find_decoder( pCodecCtx->codec_id);
    if(pCodec==NULL) {
        fprintf(stderr, "Unsupported codec!\n");
        return -1; // Codec not found
    }
    is_vp8 = strstr(pCodec->name, "vp8") != NULL;
    
    /// Open codec
    //if( avcodec_open(pCodecCtx, pCodec) < 0 ) -- Deprecated
    if( avcodec_open2(pCodecCtx, pCodec, NULL) < 0 )
        return -1; // Could not open codec
    
    /// Allocate video frame
    pFrame = avcodec_alloc_frame();
    

    /// Allocate an AVFrame structure
    pFrameRGB = avcodec_alloc_frame();
    if(pFrameRGB==NULL)
        return -1;
    

    /// Determine required buffer size and allocate buffer
    numBytes = avpicture_get_size(PIX_FMT_RGB24,
                                  pCodecCtx->width,
                                  pCodecCtx->height);
    
    buffer = (uint8_t *) av_malloc(numBytes*sizeof(uint8_t));


    /// Assign appropriate parts of buffer to image planes in pFrameRGB
    // Note that pFrameRGB is an AVFrame, but AVFrame is a superset
    // of AVPicture
    avpicture_fill((AVPicture *)pFrameRGB, buffer, PIX_FMT_RGB24,
                   pCodecCtx->width, pCodecCtx->height);


    int w = pCodecCtx->width;
    int h = pCodecCtx->height;
    /*
      img_convert_ctx = sws_getContext(w, h, pCodecCtx->pix_fmt,
      w, h, PIX_FMT_RGB24,
      SWS_BICUBIC, NULL, NULL, NULL);
    */
    
    // Read frames and save first five frames to disk
    i=0;
    while(av_read_frame(pFormatCtx, &packets)>=0) {
        AVPacket packet = packets;
        int ret;

        // Is this a packet from the video stream?
        if(packets.stream_index!=videoStreamIdx)
            continue;
        
        while (packet.size && process_frame(pCodecCtx, pFrame, &packet));
        
        // Free the packet that was allocated by av_read_frame
        av_free_packet(&packet);

        if (frame_saved >= frame_number)
            break;
    }

    av_init_packet(&packets);
    packets.data = NULL;
    packets.size = 0;
    //Flush remaining frames that are cached in the decoder
    // Is this a packet from the video stream?
    if(packets.stream_index==videoStreamIdx) {
        int ret = 1;
        while (ret > 0) {
            ret = process_frame(pCodecCtx, pFrame, &packets);
            if (frame_saved >= frame_number)
                break;
        }
    }

    // Free the RGB image
    av_free(buffer);
    av_free(pFrameRGB);
    
    // Free the YUV frame
    av_free(pFrame);
    
    // Close the codec
    avcodec_close(pCodecCtx);
    
    // Close the video file
    avformat_close_input(&pFormatCtx);
    
    //*/
    return 0;
}