#include <stdio.h>
#include <assert.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>

int main(int argc, char *argv[]) {
    int i;
    int r;
    int video_s_ind = -1, audio_s_ind = -1;
    AVFormatContext *avInputFmtCtx = NULL;
    AVFormatContext *avOutputFmtCtx = NULL;
    AVCodec *avInputVideoDecoder = NULL;
    AVCodec *avOutputVideoEncoder = NULL;
    AVCodecContext *avInputVideoDecoderCtx = NULL;
    AVCodecContext *avOutputVideoEncoderCtx = NULL;

    if (argc < 3) {
        fprintf(stderr, "Usage: %s <http stream URL> <output file>\n", argv[0]);
        return 1;
    }

    av_register_all();
    av_log_set_level(AV_LOG_DEBUG);

    r = avformat_open_input(&avInputFmtCtx, argv[1], NULL, NULL);
    assert(r == 0);

    avInputFmtCtx->max_analyze_duration = 10 * 1000 * 1000;
    r = avformat_find_stream_info(avInputFmtCtx, NULL);
    assert(r >= 0);

    av_dump_format(avInputFmtCtx, 0, argv[1], 0);

    avOutputFmtCtx = avformat_alloc_context();
    assert(avOutputFmtCtx);

    avOutputFmtCtx->oformat = av_guess_format("mpegts", NULL, NULL);
    assert(avOutputFmtCtx->oformat);

    for (i = 0; i < avInputFmtCtx->nb_streams; i++) {
        AVStream *inputStream;
        AVStream *newStream;
        inputStream = avInputFmtCtx->streams[i];
        newStream = av_new_stream(avOutputFmtCtx, inputStream->id);
        assert(newStream);
        newStream->sample_aspect_ratio = (AVRational){1, 1};

        newStream->codec->codec_type = inputStream->codec->codec_type;
        avcodec_get_context_defaults2(newStream->codec, newStream->codec->codec_type);
        newStream->codec->codec_id = inputStream->codec->codec_id;
        newStream->codec->codec_tag = 0;
        newStream->codec->bit_rate = inputStream->codec->bit_rate;
        newStream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
        newStream->codec->time_base = inputStream->codec->time_base;
        newStream->time_base = inputStream->time_base;

        /* copy extradata if exists */
        newStream->codec->extradata_size = inputStream->codec->extradata_size;
        if (inputStream->codec->extradata_size > 0) {
            newStream->codec->extradata = av_mallocz(inputStream->codec->extradata_size);
            if (!newStream->codec->extradata) {
                printf("failed to alloc space for extradata\n");
                assert(0);
            }
            memcpy(newStream->codec->extradata, inputStream->codec->extradata,
                    inputStream->codec->extradata_size);
        }

        switch (inputStream->codec->codec_type) {
            case AVMEDIA_TYPE_VIDEO:
                video_s_ind = i;
                newStream->codec->width = inputStream->codec->width;
                newStream->codec->height = inputStream->codec->height;

                printf("input frame rate = %f\n", av_q2d(inputStream->r_frame_rate));
                newStream->r_frame_rate = inputStream->r_frame_rate;
                newStream->avg_frame_rate = inputStream->avg_frame_rate;
                break;

            case AVMEDIA_TYPE_AUDIO:
                audio_s_ind = i;
                newStream->codec->channels = inputStream->codec->channels;
                newStream->codec->sample_rate =
                    inputStream->codec->sample_rate;
                newStream->codec->frame_size =
                    inputStream->codec->frame_size;
                newStream->codec->block_align =
                    inputStream->codec->block_align;
                break;
            default:
                printf("stream of unknown type %d\n", inputStream->codec->codec_type);
                assert(0);
        }
    }
    if (audio_s_ind == -1) {
        printf("Audio stream absent, fail\n");
        assert(0);
    }
    if (video_s_ind == -1) {
        printf("Video stream absent, fail\n");
        assert(0);
    }

    avInputVideoDecoder = avcodec_find_decoder(avInputFmtCtx->streams[video_s_ind]->codec->codec_id);
    assert(avInputVideoDecoder);
    avInputVideoDecoderCtx = avInputFmtCtx->streams[video_s_ind]->codec;
    r = avcodec_open2(avInputVideoDecoderCtx, avInputVideoDecoder, NULL);
    assert(r == 0);

    avOutputVideoEncoder = avcodec_find_encoder(avOutputFmtCtx->streams[video_s_ind]->codec->codec_id);
    assert(avOutputVideoEncoder);
    avOutputVideoEncoderCtx = avOutputFmtCtx->streams[video_s_ind]->codec;
    r = avcodec_copy_context(avOutputVideoEncoderCtx, avInputVideoDecoderCtx);
    assert(r == 0);

    avOutputVideoEncoderCtx->qmin = -1;
    avOutputVideoEncoderCtx->qmax = -1;

    r = avcodec_open2(avOutputVideoEncoderCtx, avOutputVideoEncoder, NULL);
    assert(r == 0);

    r = avio_open(&avOutputFmtCtx->pb, argv[2], AVIO_FLAG_WRITE);
    assert(r == 0);
    if (!avOutputFmtCtx->pb) {
        printf("av_alloc_put_byte failed\n");
        assert(0);
    }

    av_dump_format(avOutputFmtCtx, 0, avOutputFmtCtx->filename, 1);

    r = avformat_write_header(avOutputFmtCtx, NULL);
    if (r) {
        printf("Could not write header for output file");
        assert(0);
    }

    int64_t start_pts = 0;

    for (i = 0; i < 1000; i++) {
        AVPacket pkt;
        r = av_read_frame(avInputFmtCtx, &pkt);
        assert(r == 0);

        if (start_pts == 0) {
            start_pts = pkt.pts;
        }
        pkt.pts -= start_pts - 2 * 90000; /* add 2 secs to avoid negative pts in another AVStream */
        pkt.dts = pkt.pts;

        if (pkt.stream_index == video_s_ind) {
            AVFrame *frame;
            int got_picture_ptr;
            uint8_t video_outbuf[1024 * 1024];
            int video_outbuf_size = sizeof(video_outbuf);

            frame = avcodec_alloc_frame();

            r = avcodec_decode_video2(avInputVideoDecoderCtx, frame, &got_picture_ptr, &pkt);
            assert(r >= 0);

            frame->pts = pkt.pts;
            r = avpicture_alloc((AVPicture*)frame,
                    avOutputVideoEncoderCtx->pix_fmt,
                    avOutputVideoEncoderCtx->width,
                    avOutputVideoEncoderCtx->height);
            assert(r == 0);

            r = avcodec_encode_video(avOutputVideoEncoderCtx, video_outbuf, video_outbuf_size, frame);
            assert(r >= 0);
            if (r > 0) {
                video_outbuf_size = r;
                AVPacket x_pkt; // transcoded packet
                av_init_packet(&x_pkt);
                x_pkt.stream_index = video_s_ind;
                x_pkt.pts = avOutputVideoEncoderCtx->coded_frame->pts;
                x_pkt.dts = x_pkt.pts;
                x_pkt.data = video_outbuf;
                x_pkt.size = video_outbuf_size;
                if(avOutputVideoEncoderCtx->coded_frame->key_frame)
                    x_pkt.flags |= AV_PKT_FLAG_KEY;
                printf("writing transcoded video frame with pts %"PRId64", size %d\n",
                        x_pkt.pts, video_outbuf_size);
                r = av_interleaved_write_frame(avOutputFmtCtx, &x_pkt);
                assert(r == 0);
            }
            av_free(frame);
        } else {
            printf("writing audio frame with pts %"PRId64"\n",
                    pkt.pts);
            r = av_interleaved_write_frame(avOutputFmtCtx, &pkt);
            assert(r == 0);
        }
    }

    r = av_write_trailer(avOutputFmtCtx);

    // TODO close contexts
    return 0;
}