segmenter.c

// segmenter.c
// Copyright (C) 2012  Iliya Grushevskiy <iliya.gr@gmail.com>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#include "config.h"
#include "segmenter.h"
#include "util.h"

#include <math.h>
#include <stdio.h>
#include <limits.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#ifndef AV_INPUT_BUFFER_PADDING_SIZE
#define AV_INPUT_BUFFER_PADDING_SIZE FF_INPUT_BUFFER_PADDING_SIZE
#endif

#define max(a,b) (((a) > (b)) ? (a) : (b))

static const char* kExtensionAAC    = "aac";
static const char* kExtensionMP3    = "mp3";
static const char* kExtensionMPEGTS = "ts";

static const char* kFormatADTS      = "adts";
static const char* kFormatMP3       = "mp3";
static const char* kFormatMPEGTS    = "mpegts";

static const size_t kAvgSegmentsCount = 128;

static AVStream* copy_stream(AVFormatContext *context, AVStream *source_stream) {
    AVStream *output_stream = avformat_new_stream(context, source_stream->codec->codec);
    
    AVCodecContext *output_context = output_stream->codec, *source_context = source_stream->codec;
    
    output_context->codec_id   = source_context->codec_id;
    output_context->codec_type = source_context->codec_type;
    
    if (!output_context->codec_tag) {
        if (!context->oformat->codec_tag || 
            av_codec_get_id(context->oformat->codec_tag, source_context->codec_tag) ==  output_context->codec_id ||
            av_codec_get_tag(context->oformat->codec_tag, source_context->codec_id) <=0) {
            output_context->codec_tag = source_context->codec_tag;
        }
    }
    
    
    output_context->bit_rate = source_context->bit_rate;
    output_context->rc_max_rate = source_context->rc_max_rate;
    output_context->rc_buffer_size = source_context->rc_buffer_size;
    output_context->field_order = source_context->field_order;
    output_context->extradata = av_mallocz(source_context->extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
    
    if (!output_context->extradata) {
        exit(-1);
    }
    
    memcpy(output_context->extradata, source_context->extradata, source_context->extradata_size);
    
    output_context->extradata_size = source_context->extradata_size;
    
    output_context->time_base = source_context->time_base;
    
    switch (output_context->codec_type) {
        case AVMEDIA_TYPE_AUDIO:
            output_context->channel_layout = source_context->channel_layout;
            output_context->sample_rate = source_context->sample_rate;
            output_context->channels = source_context->channels;
            output_context->frame_size = source_context->frame_size;
            output_context->audio_service_type = source_context->audio_service_type;
            output_context->block_align = source_context->block_align;
            break;
            
        case AVMEDIA_TYPE_VIDEO:    
            output_context->pix_fmt = source_context->pix_fmt;
            output_context->width = source_context->width;
            output_context->height = source_context->height;
            output_context->has_b_frames = source_context->has_b_frames;
            output_context->sample_aspect_ratio = source_context->sample_aspect_ratio;
            output_stream->sample_aspect_ratio  = output_context->sample_aspect_ratio;
        default:
            break;
    }
    
    return output_stream;
}


static int load_decoder(AVCodecContext *context) {
    AVCodec *codec = avcodec_find_decoder(context->codec_id);
    
    if (!codec || avcodec_open2(context, codec, NULL)) {
        return -1;
    }
    
    return 0;
}

/**
 * @brief allocate segmenter context and set default values
 * @param output segmenter context
 * @return 0 on success, negative error code on failure
 */
int segmenter_alloc_context(SegmenterContext** context) {
    
    SegmenterContext *_context = (SegmenterContext*)malloc(sizeof(SegmenterContext));
    
    if (!_context) {
        return SGERROR(SGERROR_MEM_ALLOC);
    }
    
    _context->output           = NULL;
    
    _context->buf              = NULL;
    _context->buf_size         = 0;
    
    _context->file_base_name   = NULL;
    _context->media_base_name  = NULL;
    
    _context->segment_file_sequence = 0;
    _context->segment_sequence      = 0;
    _context->segment_index         = 0;
    _context->segment_duration      = 0;
    
    _context->duration         = 0;
    
    _context->target_duration  = 0;
    _context->max_duration     = 0;
    
    _context->max_bitrate      = 0;
    _context->avg_bitrate      = 0;
    
    _context->_pts             = 0;
    _context->_dts             = 0;
     
    _context->eof              = 0;
    
    _context->durations_size   = kAvgSegmentsCount;
    _context->durations        = (double*)malloc(sizeof(double) * _context->durations_size);
    
    if (!_context->durations) {
        free(_context);
        return SGERROR(SGERROR_MEM_ALLOC);
    }
    
    _context->bfilter = av_bitstream_filter_init("h264_mp4toannexb");
    
    *context = _context;
    
    return 0;
}

/**
 * @brief free segmenter context
 * @param context segmenter context
 */
void segmenter_free_context(SegmenterContext* context) {
    
    if (context->bfilter) {
        av_bitstream_filter_close(context->bfilter);
    }
    
    if (context->output) {
        avformat_free_context(context->output);
    }
    
    if (context->buf) {
        free(context->buf);
    }
    
    if (context->durations) {
        free(context->durations);
    }
    
    free(context);
}

/**
 * @brief initialize segmenter with source context
 * @param context segmenter context
 * @param source input source context
 * @param filter stream filter
 * @return 0 on success, negative error code on failure
 */
int segmenter_init(SegmenterContext *context, AVFormatContext *source, char* file_base_name, char* media_base_name, 
                    double target_duration, int media_filter) {
    int i;
    int video_index = -1, audio_index = -1;
    
    for (i=0; i < source->nb_streams && (video_index < 0 || audio_index < 0); i++) {
        AVStream *_stream = source->streams[i];
        
        switch (_stream->codec->codec_type) {
            case AVMEDIA_TYPE_VIDEO:
                
                if ((media_filter & MediaTypeVideo) && !load_decoder(_stream->codec)) {
                    video_index = i;
                } 
                
                break;
                
            case AVMEDIA_TYPE_AUDIO:
                
                if ((media_filter & MediaTypeAudio) && !load_decoder(_stream->codec)) {
                    audio_index = i;
                } 
                
                break;    
                
            default:
                break;
        }
    }
    
    if (video_index < 0 && audio_index < 0) {
        return SGERROR(SGERROR_NO_STREAM);
    }
    
    context->source_video_index = video_index;
    context->source_audio_index = audio_index;
    
    AVOutputFormat *oformat;
    
    if (video_index < 0) {
        switch (source->streams[audio_index]->codec->codec_id) {
            case AV_CODEC_ID_AAC:
                oformat = av_guess_format(kFormatADTS, NULL, NULL);
                context->extension = kExtensionAAC;
                break;
            case AV_CODEC_ID_MP3:
                oformat = av_guess_format(kFormatMP3, NULL, NULL);
                context->extension = kExtensionMP3;
            default:
                oformat = av_guess_format(kFormatMPEGTS, NULL, NULL);
                context->extension = kExtensionMPEGTS;
                break;
        }
    } else {
        oformat = av_guess_format(kFormatMPEGTS, NULL, NULL);
        context->extension = kExtensionMPEGTS;
    }
    
    if (!oformat) {
        return SGERROR(SGERROR_UNSUPPORTED_FORMAT);
    }
    
    if (avformat_alloc_output_context2(&context->output, oformat, NULL, NULL)) {
        return SGERROR(SGERROR_MEM_ALLOC);
    }
    
    context->video = video_index >= 0 ? copy_stream(context->output, source->streams[video_index]) : NULL;
    context->audio = audio_index >= 0 ? copy_stream(context->output, source->streams[audio_index]) : NULL;
    
    context->file_base_name  = file_base_name;
    context->media_base_name = media_base_name;
    context->target_duration = target_duration;
    
    context->buf_size = snprintf(NULL, 0, "%s/%s%u.%s", context->file_base_name, context->media_base_name, UINT_MAX, context->extension) + 1;
    
    if (!(context->buf = (char*)malloc(context->buf_size * sizeof(char)))) {
        return SGERROR(SGERROR_MEM_ALLOC);
    }
    
    return 0;
}

static int set_segment_duration(SegmenterContext *context, unsigned int index, double duration) {
    
    if (index - context->segment_sequence >= context->durations_size) {
        context->durations_size += kAvgSegmentsCount;
        context->durations = (double*)realloc(context->durations, context->durations_size * sizeof(double));
        
        if (!context->durations) {
            return SGERROR(SGERROR_MEM_ALLOC);
        }
    }
    
    context->durations[index - context->segment_sequence] = duration;
    context->max_duration = max(context->max_duration, duration);
    
    return 0;
}

static inline double segment_duration(SegmenterContext *context, unsigned int segment) {
    return context->durations[segment - context->segment_sequence];
}


/**
 * @brief starts next segment
 * @param context segmenter context
 * @return 0 on success, negative error code on failure
 */
static int start_segment(SegmenterContext *context) {

    snprintf(context->buf, context->buf_size, "%s/%s%u.%s", context->file_base_name, context->media_base_name, context->segment_index, context->extension);
    
    if (avio_open(&context->output->pb, context->buf, AVIO_FLAG_WRITE)) {
        return SGERROR(SGERROR_FILE_WRITE);
    }
    
    if (context->segment_index == 0) {
        return avformat_write_header(context->output, NULL);
    }
    
    return 0;
}

/**
 * @brief finish segment
 * @param context segmenter context
 * @return 0 on success, negative error code on failure
 */
static int finish_segment(SegmenterContext *context) {
    AVFormatContext *output = context->output;
    size_t size;
    int    ret;
    
    avio_flush(output->pb);
    size = avio_size(output->pb);
    avio_close(output->pb);
    
    context->max_bitrate = max(context->max_bitrate, size * 8 / context->segment_duration);
    context->avg_bitrate = (context->avg_bitrate * context->segment_index + (size * 8 / context->segment_duration)) / (context->segment_index + 1);
    
    if ((ret = set_segment_duration(context, context->segment_index, context->segment_duration))) {
        return ret;
    }
    
    context->segment_index++;
    context->segment_duration = 0;
    
    return 0;
}

/**
 * @brief open first output segment
 * @param context segmenter context
 * @return 0 on success, negative error code on failure
 */
int segmenter_open(SegmenterContext* context) {
    return start_segment(context);
}

/**
 * @brief close last output segment
 * @param context segmenter context
 * @return 0 on success, negative error code on failure
 */
int segmenter_close(SegmenterContext* context) {
    context->eof = 1;
    
    return finish_segment(context);
}

/**
 * @brief delete segmenets out from current sequence
 * @param context segmenter context
 * @return 0 on success, negative error code on failure
 */
static int clear_segments(SegmenterContext *context) {
    
    unsigned int i;
    
    for (i = context->segment_file_sequence; i < context->segment_sequence; i++) {
        snprintf(context->buf, context->buf_size, "%s/%s%u.%s", context->file_base_name, context->media_base_name, i, context->extension);
        unlink(context->buf);
    }
    
    context->segment_file_sequence = context->segment_sequence;
        
    return 0;
}

/**
 * @brief set segmenter sequence index
 * @param context segmenter context
 * @param sequence new sequence value
 * @param del whether to delete out sequence segments from disk
 */
int segmenter_set_sequence(SegmenterContext *context, unsigned int sequence, int del) {
    double duration = 0;
    
    int i;
    
    if (sequence <= context->segment_sequence || sequence >= context->segment_index) {
        return 0;
    }
    
    for (i=sequence; i < context->segment_index; i++) {
        duration += segment_duration(context, i);
    }
    
    i = sequence;
    
    while (duration < context->target_duration * 3 && i > context->segment_sequence) {
        duration += segment_duration(context, --i);
    }
        
    sequence = i;
    
    context->max_duration = 0;
    
    for (i = sequence; i < context->segment_index; i++) {
        context->durations[i - sequence] = context->durations[i - context->segment_sequence];
        context->max_duration = max(context->max_duration, context->durations[i - sequence]);
    }
    
    context->segment_sequence = sequence;
    
    if (del) {
        clear_segments(context);
    }
    
    return 0;
}

/**
 * @brief write packet from input to output
 * @param context segmenter context
 * @param source source input
 * @param pkt packet to write
 * @return 0 on success, negative error code on failure
 */
int segmenter_write_pkt(SegmenterContext* context, AVFormatContext *source, AVPacket *pkt) {
    
    AVPacket opkt;
    AVStream *stream, *output_stream;
    
    stream = source->streams[pkt->stream_index];
    
    if (pkt->stream_index != context->source_audio_index && pkt->stream_index != context->source_video_index) {
        av_packet_unref(pkt);
        return 0;
    }
    
    av_init_packet(&opkt);
    
    output_stream = pkt->stream_index == context->source_audio_index ? context->audio : context->video;
    
    opkt.stream_index = output_stream->index;
    
    if (pkt->pts != AV_NOPTS_VALUE) {
        opkt.pts = av_rescale_q(pkt->pts, stream->time_base, output_stream->time_base);
    } else {
        opkt.pts = AV_NOPTS_VALUE;
    }
    
    if (pkt->dts == AV_NOPTS_VALUE) {
        opkt.dts = context->_dts;
    } else {
        opkt.dts = av_rescale_q(pkt->dts, stream->time_base, output_stream->time_base);
        context->_dts = opkt.dts;
    }
    
    opkt.duration = av_rescale_q(pkt->duration, stream->time_base, output_stream->time_base);
    opkt.flags    = pkt->flags;
    opkt.data     = pkt->data;
    opkt.size     = pkt->size;
    
    if (pkt->stream_index == context->source_video_index && context->bfilter) {
        AVPacket _pkt = opkt;
        
        av_bitstream_filter_filter(context->bfilter, context->video->codec, NULL, &_pkt.data, &_pkt.size, pkt->data, pkt->size, pkt->flags & AV_PKT_FLAG_KEY);
        
        av_packet_unref(&opkt);
        opkt = _pkt;
        // opkt.destruct = av_destruct_packet;
    }  
    
    context->duration = opkt.pts * av_q2d(output_stream->time_base);
    
    if (context->source_video_index < 0 || (output_stream == context->video && (opkt.flags & AV_PKT_FLAG_KEY))) {
        context->segment_duration = (opkt.pts - context->_pts) * av_q2d(output_stream->time_base);
    }
    
    if (context->segment_duration >= context->target_duration) {
        int ret;
        
        context->_pts = opkt.pts;
        
        if((ret = finish_segment(context))) {
            return ret;
        }
        
        if ((ret = start_segment(context))) {
            return ret;
        }
    }
    
    
    av_write_frame(context->output , &opkt);
    av_packet_unref(&opkt);
    
    return 0;
}


static FILE* segmenter_open_playlist(SegmenterContext *context, char *index_file) {
    int length;
    char *filename;
    
    length = snprintf(NULL, 0, "%s/%s", context->file_base_name, index_file);
    
    if (!(filename = (char*)alloca(sizeof(char)*(length+1)))) {
        return NULL;
    }
    
    snprintf(filename, length + 1, "%s/%s", context->file_base_name, index_file);
    
    return fopen(filename, "w");
}

static void segmenter_close_playlist(FILE *out) {
    fclose(out);
}

/**
 * @brief write stream index
 * @param context segmenter context
 * @param index_file index file base name
 * @return 0 on success, negative error code on error 
 */
int segmenter_write_playlist(SegmenterContext *context, IndexType type, char* base_url, char *index_file) {
    
    if (type == IndexTypeVOD && !context->eof) {
        return 0;
    }
    
    FILE *out = segmenter_open_playlist(context, index_file);
    
    if (!out) {
        return SGERROR(SGERROR_FILE_WRITE);
    }
        
    fprintf(out, "#EXTM3U\n"
                 "#EXT-X-TARGETDURATION:%ld\n"
                 "#EXT-X-VERSION:3\n"
                 "#EXT-X-MEDIA-SEQUENCE:%u\n", lround(context->max_duration), context->segment_sequence);
    
    switch (type) {
        case IndexTypeVOD:
            fprintf(out, "#EXT-X-PLAYLIST-TYPE:VOD\n");
            break;
        case IndexTypeEvent:
            fprintf(out, "#EXT-X-PLAYLIST-TYPE:EVENT\n");
            break;
        default:
            break;
    }
    
    int i;
    for (i = context->segment_sequence; i < context->segment_index; i++) {
        fprintf(out, "#EXTINF:%ld,\n"
                     "%s%s%u.%s\n", lround(segment_duration(context, i)), base_url, context->media_base_name, i, context->extension);
    }
    
    if ((type == IndexTypeEvent || type == IndexTypeVOD) && context->eof) {
        fprintf(out, "#EXT-X-ENDLIST");
    }
    
    segmenter_close_playlist(out);
    
    return 0;
}