/*
 *  dumpmovie     dump a movie into YUV PNG files
 *  Copyright (c) 2010 Sam Hocevar <sam@hocevar.net>
 *                All Rights Reserved
 *
 *  $Id$
 *
 *  This program is free software. It comes without any warranty, to
 *  the extent permitted by applicable law. You can redistribute it
 *  and/or modify it under the terms of the Do What The Fuck You Want
 *  To Public License, Version 2, as published by Sam Hocevar. See
 *  http://sam.zoy.org/wtfpl/COPYING for more details.
 */

#include "config.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

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

#include <pipi.h>

int main(int argc, char *argv[])
{
    char fmtstr[1024];
    AVPacket packet;
    AVFormatContext *fmt;
    AVCodecContext *ctx;
    AVCodec *codec;
    AVFrame *frame;
    struct SwsContext *sws = NULL;
    pipi_image_t *image = NULL;
    pipi_pixels_t *p;
    uint8_t *dst[4], *data = NULL;
    char *parser;
    int stream, pitch[4], i, k = 0, nframes = -1;
    double skip_seconds = 0.0;

    if (argc < 2)
        return EXIT_FAILURE;

    if (argc > 2)
        skip_seconds = atof(argv[2]);
    if (argc > 3)
        nframes = atoi(argv[3]);

    /* Ensure our linear YUV values do not get gamma-corrected */
    pipi_set_gamma(1.0);

    parser = strrchr(argv[1], '/');
    strcpy(fmtstr, parser ? parser + 1 : argv[1]);
    parser = strrchr(fmtstr, '.');
    if(parser)
        *parser = '\0';
    strcat(fmtstr, "-%06i.png");

    av_register_all();
    if(av_open_input_file(&fmt, argv[1], NULL, 0, NULL) != 0)
        return EXIT_FAILURE;
    if(av_find_stream_info(fmt) < 0 )
        return EXIT_FAILURE;

    stream = -1;
    for(i = 0; stream == -1 && i < (int)fmt->nb_streams; i++)
        if(fmt->streams[i]->codec->codec_type == CODEC_TYPE_VIDEO)
            stream = i;
    if(stream == -1)
        return EXIT_FAILURE;
    ctx = fmt->streams[stream]->codec;

    codec = avcodec_find_decoder(ctx->codec_id);
    if(codec == NULL)
        return EXIT_FAILURE;
    if(avcodec_open(ctx, codec) < 0)
        return EXIT_FAILURE;

    skip_seconds /= av_q2d(fmt->streams[stream]->time_base);
    av_seek_frame(fmt, stream, (int)(skip_seconds + 0.5), SEEK_SET);
    //avformat_seek_file(fmt, stream, skip_bytes, skip_bytes,
    //                   skip_bytes, AVSEEK_FLAG_BYTE);

    frame = avcodec_alloc_frame();

    for (k = 0; k < nframes || nframes == -1; /* k incremented below */)
    {
        int finished, ret, x, y;

        ret = av_read_frame(fmt, &packet);

        if(ret < 0)
            break;

        if(packet.stream_index != stream)
        {
            av_free_packet(&packet);
            continue;
        }

        avcodec_decode_video(ctx, frame, &finished, packet.data, packet.size);
        if(!finished)
        {
            av_free_packet(&packet);
            continue;
        }

        if(!sws)
        {
            sws = sws_getContext(ctx->width, ctx->height, ctx->pix_fmt,
                                 ctx->width, ctx->height, PIX_FMT_YUV444P,
                                 SWS_FAST_BILINEAR, NULL, NULL, NULL);
            pitch[0] = ctx->width;
            pitch[1] = ctx->width;
            pitch[2] = ctx->width;
            pitch[3] = 0;
            dst[0] = malloc(ctx->width * ctx->height * 3);
            dst[1] = dst[0] + ctx->width * ctx->height;
            dst[2] = dst[1] + ctx->width * ctx->height;
            dst[3] = NULL;
            image = pipi_new(ctx->width, ctx->height);
        }

        {
            char buf[1024];
            sprintf(buf, fmtstr, k);
            printf("saving in %s\n", buf);

            sws_scale(sws, (uint8_t const **)frame->data, frame->linesize, 0,
                      ctx->height, dst, pitch);

            p = pipi_get_pixels(image, PIPI_PIXELS_RGBA_U8);
            data = (uint8_t *)p->pixels;

            for (y = 0; y < ctx->height; y++)
            {
                int off = y * ctx->width;

                for (x = 0; x < ctx->width; x++, off++)
                {
                    /* Reorder components to store YUVA */
                    data[4 * off] = dst[0][off];
                    data[4 * off + 1] = dst[2][off];
                    data[4 * off + 2] = dst[1][off];
                    data[4 * off + 3] = 0xff;
                }
            }

            pipi_save(image, buf);
        }

        av_free_packet(&packet);

        k++;
    }

    return EXIT_SUCCESS;
}