/*
 *
 * SAMPLE CODE FOR MP4-to-MPEG-TS CONTAINER CONVERSION
 *
 * This is the basis for a python plugin for doing on-the-fly segmentation of
 *
 *
 */

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

#include <libavformat/avformat.h>
#include <libavutil/mathematics.h>  // For av_rescale_q

typedef struct
{
    const char *input_filename;
    const char *output_playlist;
    const char *output_prefix;
    double segment_length;
} config_info;

typedef struct
{
    double duration;
    char *filename;
} segment_info;

static AVStream *add_output_stream(AVFormatContext *output_format_context, AVStream *input_stream)
{
    AVCodecContext *input_codec_context;
    AVCodecContext *output_codec_context;
    AVStream *output_stream;

    output_stream = avformat_new_stream(output_format_context, NULL);
    if (!output_stream) {
        fprintf(stderr, "Could not allocate stream\n");
        exit(1);
    }

    input_codec_context = input_stream->codec;
    output_codec_context = output_stream->codec;

    output_codec_context->codec_id = input_codec_context->codec_id;
    output_codec_context->codec_type = input_codec_context->codec_type;
    output_codec_context->codec_tag = input_codec_context->codec_tag;
    output_codec_context->bit_rate = input_codec_context->bit_rate;
    output_codec_context->extradata = input_codec_context->extradata;
    output_codec_context->extradata_size = input_codec_context->extradata_size;

    // Some codecs put the wrong ticks_per_frame in (in the thousands instead of tens)
    if(av_q2d(input_codec_context->time_base) * input_codec_context->ticks_per_frame > av_q2d(input_stream->time_base) && av_q2d(input_stream->time_base) < 1.0/1000) {
        output_codec_context->time_base = input_codec_context->time_base;
        output_codec_context->time_base.num *= input_codec_context->ticks_per_frame;
    }
    else {
        output_codec_context->time_base = input_stream->time_base;
    }

    switch (input_codec_context->codec_type) {
        case AVMEDIA_TYPE_AUDIO:
            output_codec_context->channel_layout = input_codec_context->channel_layout;
            output_codec_context->sample_rate = input_codec_context->sample_rate;
            output_codec_context->channels = input_codec_context->channels;
            output_codec_context->frame_size = input_codec_context->frame_size;
            if ((input_codec_context->block_align == 1 && input_codec_context->codec_id == AV_CODEC_ID_MP3) || input_codec_context->codec_id == AV_CODEC_ID_AC3) {
                output_codec_context->block_align = 0;
            }
            else {
                output_codec_context->block_align = input_codec_context->block_align;
            }
            fprintf(stdout, "Audio time bases:\n");
            break;
        case AVMEDIA_TYPE_VIDEO:
            output_codec_context->pix_fmt = input_codec_context->pix_fmt;
            output_codec_context->width = input_codec_context->width;
            output_codec_context->height = input_codec_context->height;
            output_codec_context->has_b_frames = input_codec_context->has_b_frames;

            if (output_format_context->oformat->flags & AVFMT_GLOBALHEADER) {
                output_codec_context->flags |= CODEC_FLAG_GLOBAL_HEADER;
            }
            fprintf(stdout, "Video time bases:\n");
            break;
        default:
            break;
    }

    return output_stream;
}

static int write_frame_fixed(AVPacket *packet, AVFormatContext *input_context, AVFormatContext *output_context, AVBitStreamFilterContext *bsf)
{
    int si = packet->stream_index;
    if (output_context->streams[si]->codec->codec_type == AVMEDIA_TYPE_VIDEO)
    {
        if (input_context->streams[si]->codec->codec_id == AV_CODEC_ID_H264)
        {
            int a = av_bitstream_filter_filter(bsf, output_context->streams[si]->codec, NULL,
                                               &packet->data, &packet->size,
                                               packet->data, packet->size,
                                               packet->flags & AV_PKT_FLAG_KEY);

            if (a == 0)
            {
                fprintf(stderr, "Eh?!\n");
            }
            else if (a < 0)
            {
                fprintf(stderr, "%s failed for stream %d, codec %s, error %d\n", bsf->filter->name, packet->stream_index, output_context->streams[si]->codec->codec->name, a);
                exit(1);
            }
        }
    }

    packet->pts = av_rescale_q(packet->pts, input_context->streams[si]->time_base, output_context->streams[si]->time_base);
    packet->dts = av_rescale_q(packet->dts, input_context->streams[si]->time_base, output_context->streams[si]->time_base);

    return av_interleaved_write_frame(output_context, packet);
}

int main(int argc, char **argv)
{
    if (argc != 5)
    {
        fprintf(stderr, "Usage: %s <input filename> <playlist filename> <output prefix> <segment length>\n", argv[0]);
        exit(1);
    }

    config_info config;
    memset(&config, 0, sizeof(config_info));

    config.input_filename = argv[1];
    config.output_playlist = argv[2];
    config.output_prefix = argv[3];
    config.segment_length = atoi(argv[4]);

    printf("Segment length: %f", config.segment_length);

    char *output_filename = (char*) malloc(sizeof(char) * (strlen(config.output_prefix) + 10));
    if (!output_filename)
    {
        fprintf(stderr, "Segmenter error: Could not allocate space for output filenames with prefix %s\n", config.output_prefix);
        exit(1);
    }

    avcodec_register_all();
    av_register_all();

    // Start context from file, guess format type based on filename
    AVFormatContext *input_context = NULL;
    int ret = avformat_open_input(&input_context, config.input_filename, NULL, NULL);
    if (ret != 0)
    {
        fprintf(stderr, "Segmenter error: Could not open input file, make sure it is an MP4 file: %d\n", ret);
        exit(1);
    }

    if (avformat_find_stream_info(input_context, NULL) < 0)
    {
        fprintf(stderr, "Segmenter error Could not read stream information.\n");
        exit(1);
    }

    AVOutputFormat *output_format = av_guess_format("mpegts", NULL, NULL);
    if (!output_format)
    {
        fprintf(stderr, "Segmenter error: Could not find MPEG-TS muxer.\n");
        exit(1);
    }

    AVFormatContext *output_context = avformat_alloc_context();
    if (!output_context)
    {
        fprintf(stderr, "Segmenter error: Could not allocate output context.\n");
        exit(1);
    }
    output_context->oformat = output_format;

    int video_index = -1;
    int audio_index = -1;

    AVStream *video_stream;
    AVStream *audio_stream;

    int i;

    for (i = 0; i < input_context->nb_streams && (video_index < 0 || audio_index < 0); i++)
    {
        switch (input_context->streams[i]->codec->codec_type)
        {
            case AVMEDIA_TYPE_VIDEO:
                video_index = i;
                input_context->streams[i]->discard = AVDISCARD_NONE;
                video_stream = add_output_stream(output_context, input_context->streams[i]);
                break;

            case AVMEDIA_TYPE_AUDIO:
                audio_index = i;
                input_context->streams[i]->discard = AVDISCARD_NONE;
                audio_stream = add_output_stream(output_context, input_context->streams[i]);
                break;

            default:
                input_context->streams[i]->discard = AVDISCARD_ALL;
                break;
        }
    }

    if (video_index == -1 || audio_index == -1)
    {
        fprintf(stderr, "Segmenter error: Video or audio stream not found.\n");
        exit(1);
    }

    av_dump_format(input_context, 0, config.input_filename, 0);
    av_dump_format(output_context, 0, config.output_prefix, 1);

    AVCodec *codec = avcodec_find_decoder(video_stream->codec->codec_id);
    if (!codec)
    {
        fprintf(stderr, "Segmenter error: Could not find video decoder, key frames will not be honored.\n");
    }

    if (avcodec_open2(video_stream->codec, codec, NULL) < 0)
    {
        fprintf(stderr, "Segmenter error: Could not open video decoder, key frames will not be honored.\n");
    }

    AVCodec *codec2 = avcodec_find_decoder(audio_stream->codec->codec_id);
    if (!codec2)
    {
        fprintf(stderr, "Segmenter error: Could not find audio decoder, key frames will not be honored.\n");
    }

    if (avcodec_open2(audio_stream->codec, codec2, NULL) < 0)
    {
        fprintf(stderr, "Segmenter error: Could not open audio decoder, key frames will not be honored.\n");
    }

    // Generate output filename
    unsigned int output_index = 1;
    snprintf(output_filename, strlen(config.output_prefix) + 10, "%s-%05u.ts", config.output_prefix, output_index++);

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


    if (avformat_write_header(output_context, NULL))
    {
        fprintf(stderr, "Segmenter error: Could not write mpegts header to first output file.\n");
        exit(1);
    }

    // Required when muxing h264 into MPEG-TS
    AVBitStreamFilterContext *bsf = av_bitstream_filter_init("h264_mp4toannexb");
    if (!bsf)
    {
        fprintf(stderr, "Error initializing the h264_mp4toannexb filter!");
        exit(1);
    }

    double prev_segment_time = 0;
    int decode_done;
    int packet_c = 0;
    int64_t ppts = 0;
    double tbase;

    av_seek_frame(input_context, video_index, seek_to , 0);
    do
    {
        double segment_time;
        AVPacket packet;
        av_init_packet(&packet);
        decode_done = av_read_frame(input_context, &packet);
        if (decode_done < 0)
        {
            break;
        }

        if (av_dup_packet(&packet) < 0)
        {
            fprintf(stderr, "Segmenter error: Could not duplicate packet (av_dup_packet).\n");
            av_free_packet(&packet);
            break;
        }

        ppts = packet.pts;
        tbase = av_q2d(input_context->streams[packet.stream_index]->time_base);

        if (prev_segment_time == 0)
        {
            prev_segment_time = ppts * tbase;
        }

        // Segmentation time!
        if (packet.stream_index == video_index && (packet.flags & AV_PKT_FLAG_KEY))
        {
            segment_time = ppts * tbase;
            //printf("Video keyframe at %f\n", segment_time);
        }
        else
        {
            segment_time = prev_segment_time;
        }

        if (segment_time - prev_segment_time >= config.segment_length)
        {
            fprintf(stdout, "Segment from %f to %f (%f seconds, %d packets).\n", prev_segment_time, segment_time, (segment_time - prev_segment_time), packet_c);
            avio_flush(output_context->pb);
            avio_close(output_context->pb);

            packet_c = 0;

            snprintf(output_filename, strlen(config.output_prefix) + 10, "%s-%05u.ts", config.output_prefix, output_index++);
            if (avio_open(&output_context->pb, output_filename, AVIO_FLAG_WRITE) < 0)
            {
                fprintf(stderr, "Segmenter error: Could not open '%s'\n", output_filename);
                exit(1);
            }

            prev_segment_time = segment_time;
        }
        ret = write_frame_fixed(&packet, input_context, output_context, bsf);
        packet_c++;
        if (ret < 0)
        {
            fprintf(stderr, "Segmenter error: Could not write frame of stream: %d\n", ret);
            exit(1);
        }
        else if (ret > 0)
        {
            fprintf(stderr, "Segmenter info: End of stream requested.\n");
            av_free_packet(&packet);
            break;
        }
        av_free_packet(&packet);
    } while (!decode_done);

    av_write_trailer(output_context);
    avcodec_close(video_stream->codec);

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

    av_bitstream_filter_close(bsf);

    avio_close(output_context->pb);
    av_free(output_context);

    fprintf(stdout, "SUCCESS!!!\n");
    exit(0);
}