#include "vid.h"

#include <unistd.h>
#include <signal.h>
#include <stdio.h>

#include <sys/time.h>

static AVFormatContext *outvid_format_context;
static AVStream *outvid_stream;
static AVCodec *outvid_codec;
static int outvid_width, outvid_height;

static int inimg_stride[8];
static uint8_t **inimg_data;
static char *inimg_data_plane_1;
static AVFrame *outvid_frame;

static struct SwsContext *sws;

static struct timeval time_current, time_previous;

static const int outvid_stride[8] = {
    640, 320, 320, 0, 0, 0, 0, 0
};

static int have_video = 0, encode_video = 0;

static int next_pts;

static int add_stream(AVFormatContext *, AVCodec **, AVStream **,
        enum AVCodecID, int width, int height); 

static AVFrame *alloc_picture(
        enum AVPixelFormat pix_fmt,
        int width, int height);

static int open_video(AVFormatContext *oc, AVCodec *codec, AVStream *stream, AVFrame **frame); //AVFrame **tmp_frame

static int get_video_frame();
static int write_video_frame();

static struct sigaction sigIntHandler;
int running = 1;
void my_handler(int s)
{
    printf("Caught signal %d\n", s);
    //exit(1);
    running = 0;
}
 
void setup_sighandler()
{
    sigIntHandler.sa_handler = my_handler;
    sigemptyset(&sigIntHandler.sa_mask);
    sigIntHandler.sa_flags = 0;
    sigaction(SIGINT, &sigIntHandler, NULL);
}
int video_main(const char *filename)
{
    AVDictionary *opt = NULL;
    AVOutputFormat *fmt = NULL;
    int outvid_width = VID_WIDTH;
    int outvid_height = VID_HEIGHT;
    int retval = 0;

    av_register_all();
    avcodec_register_all();
    avformat_network_init();

    av_log_set_level(AV_LOG_DEBUG);
    setup_sighandler();

    if (avformat_alloc_output_context2(&outvid_format_context, NULL, NULL, filename) < 0)
    {
        fprintf(stderr, "Could not dedice the output format "
                    "from file extension. using "
                    VID_DEFAULT_CONTAINER
                    "\n");
        if(avformat_alloc_output_context2(&outvid_format_context, NULL, VID_DEFAULT_CONTAINER, filename) < 0)
        {
            fprintf(stderr, "Could not create the output context\n");
            exit(1);
        }
    }
    fmt = outvid_format_context->oformat;
    if (fmt->video_codec != AV_CODEC_ID_NONE)
    {
        printf("codec: %s\n",
            (avcodec_descriptor_get(fmt->video_codec))->name);
        add_stream(outvid_format_context, &outvid_codec, &outvid_stream, AV_CODEC_ID_H264, outvid_width, outvid_height);
        have_video = 1;
        encode_video = 1;
    }
    if (have_video)
    {
        open_video(outvid_format_context, outvid_codec, outvid_stream, &outvid_frame);
    }

    if (!(fmt->flags & AVFMT_NOFILE))
    {
        retval = avio_open(&outvid_format_context->pb, filename, AVIO_FLAG_WRITE);
        if (retval < 0)
        {
            fprintf(stderr, "Could not open '%s': %s\n", filename,
                av_err2str(retval));
            exit(1);
        }
    }
    retval = avformat_write_header(outvid_format_context, &opt);
    if (retval < 0)
    {
        fprintf(stderr, "Error occurred when opening output file: %d\n",
            av_err2str(retval));
        exit(1);
    }

    inimg_data = calloc(8, sizeof(uint8_t*));

    if (wank_init() < 0)
        exit(0);
    gettimeofday(&time_previous, NULL);
    while (encode_video && running)
    {
        encode_video = !write_video_frame();
    }
    av_write_trailer(outvid_format_context);
    exit(0);
}

static int add_stream(AVFormatContext *oc, AVCodec **codec, AVStream **stream,
        enum AVCodecID id, int width, int height)
{
    printf("add_stream()\n");

    AVCodecContext *c;
    if (!(*codec = avcodec_find_encoder(id)))
    {
        fprintf(stderr, "Could not find encoder for %s\n", avcodec_get_name(id));
        exit(1);
    }
    if (!(*stream = avformat_new_stream(oc, *codec)))
    {
        fprintf(stderr, "Could not allocate stream\n");
        exit(1);
    }
    (*stream)->id = oc->nb_streams-1;
    c = (*stream)->codec;
    switch ((*codec)->type)
    {
        case AVMEDIA_TYPE_VIDEO:
            c->codec_id = id;

            c->bit_rate = 4500000;
            c->rc_buffer_size = c->bit_rate;
            c->rc_min_rate = c->rc_max_rate = c->bit_rate;
            c->width = VID_WIDTH;
            c->height = VID_HEIGHT;
            (*stream)->time_base = (AVRational){1, FPS};
            c->gop_size = 60;
            //c->keyint_min = 600;
            c->pix_fmt = AV_PIX_FMT_YUV420P;
            if (c->codec_id == AV_CODEC_ID_H264)
            {
                av_opt_set(c->priv_data, "preset", "fast", 0);
                av_opt_set(c->priv_data, "tune", "zerolatency", 0);
            }
            if (c->codec_id == AV_CODEC_ID_MPEG2VIDEO)
            {

                c->max_b_frames = 2;
            }
            if (c->codec_id == AV_CODEC_ID_MPEG1VIDEO)
            {
                c->mb_decision = 2;
            }

            break;

        default:

            break;
    }

    return 0;
}

static int open_video(AVFormatContext *oc, AVCodec *codec, AVStream *stream, AVFrame **frame) //AVFrame **tmp_frame
{
    int ret;
    AVCodecContext *c = stream->codec;
    if ((ret = avcodec_open2(c, codec, NULL)) < 0)
    {
        fprintf(stderr, "Could not open video codec: %s\n", av_err2str(ret));
        exit(1);
    }
    if (!(*frame = alloc_picture(c->pix_fmt, c->width, c->height)))
    {
        fprintf(stderr, "Could not allocate video frame\n");
        exit(1);
    }
    if (c->pix_fmt != AV_PIX_FMT_YUV420P)
    {
        fprintf(stderr, "pix_fmt not YUV420P!\n");
        exit(1);
    }


    return 0;
}

static AVFrame *alloc_picture(
        enum AVPixelFormat pix_fmt,
        int width, int height
    )
{
    int ret;
    AVFrame *result = av_frame_alloc();
    if (!result) return NULL;
    result->format = pix_fmt;
    result->width = width;
    result->height = height;
    if ((ret = av_frame_get_buffer(result, 32)) < 0)
    {
        fprintf(stderr, "Could not allocate frame data\n");
        exit(1);
    }
    return result;
}

static int write_video_frame()
{

    int ret;
    static int i = 0;
    AVCodecContext *c;
    int got_packet;
    AVPacket pkt = { 0 };


    c = outvid_stream->codec;
    gettimeofday(&time_current, NULL);
    if ( i == 0 )
    {
        get_video_frame();
    }
    else
    {
    }
    i = (i+1) % (2);
    if (!outvid_frame) return 0;
    outvid_frame->pts = next_pts++;
    av_init_packet(&pkt);
    if ((ret = avcodec_encode_video2(c, &pkt, outvid_frame, &got_packet)) < 0)
    {
        fprintf(stderr, "Error encoding video frame: %s\n",
            av_err2str(ret));
        exit(1);
    }
    if (got_packet)
    {
        pkt.stream_index = outvid_stream->index;
        if (pkt.pts != AV_NOPTS_VALUE)
        {
            pkt.pts =
                av_rescale_q(pkt.pts, c->time_base, outvid_stream->time_base);
        }
        else
        {
            pkt.pts = 0;
        }
        if (pkt.dts != AV_NOPTS_VALUE)
        {
            pkt.dts =
                av_rescale_q(pkt.dts, c->time_base, outvid_stream->time_base);
        }
        if ((ret = av_write_frame(outvid_format_context, &pkt)) < 0)
        {
            fprintf(stderr, "Error while writing video frame: %s\n",
                av_err2str(ret));
            exit(1);
        }
    }
    return (outvid_frame || got_packet) ? 0 : 1;
}

static int get_video_frame()
{
    static int prevW = -1, prevH = -1;
    AVCodecContext *c = outvid_stream->codec;
    int newW, newH;
    wank_get_image(&inimg_data_plane_1, &newW, &newH);
    inimg_data[0] = (uint8_t*)inimg_data_plane_1;
    if (!sws || newW != prevW || newH != prevH)
    {
        sws = sws_getContext(
            newW, newH, AV_PIX_FMT_BGRA,
            c->width, c->height, c->pix_fmt,
            SWS_BICUBIC,
            NULL, NULL, NULL
        );
        inimg_stride[0] = newW * 4;
        prevW = newW, prevH = newH;
    } 
    sws_scale(
        sws,
        (const uint8_t * const *)inimg_data, inimg_stride,
        0, prevH,
        outvid_frame->data, outvid_frame->linesize
    );
    free(inimg_data_plane_1);
}