#include <libavfilter/avfilter.h>
#include <libavfilter/buffersink.h>
#include <libavfilter/buffersrc.h>
#include <libavutil/channel_layout.h>
#include <libavutil/frame.h>
#include <stdio.h>
#include <stdlib.h>

typedef struct FilteringContext {
    const char *desc;
    AVFilterGraph *graph;
    int nb_inputs;
    AVFilterContext **inputs;
    int nb_outputs;
    AVFilterContext **outputs;
} FilteringContext;

static AVFrame *get_dummy_frame(int width, int height, int frame_index)
{
    int ret = 0;
    AVFrame *frame = av_frame_alloc();
    if (!frame)
        return NULL;

    frame->format = AV_PIX_FMT_YUV420P;
    frame->width = width;
    frame->height = height;

    if ((ret = av_frame_get_buffer(frame, 0)) < 0) {
        printf("Failed to allocate frame buffer: %s\n", av_err2str(ret));
        av_frame_free(&frame);
        return NULL;
    }

    for (int y = 0; y < height; y++)
        for (int x = 0; x < width; x++)
            frame->data[0][y * frame->linesize[0] + x] = x + y + frame_index * 3;

    for (int y = 0; y < height / 2; y++) {
        for (int x = 0; x < width / 2; x++) {
            frame->data[1][y * frame->linesize[1] + x] = 128 + y + frame_index * 2;
            frame->data[2][y * frame->linesize[2] + x] = 64 + x + frame_index * 5;
        }
    }

    return frame;
}

static int init_input_filter(FilteringContext *fc, AVFilterInOut *in)
{
    int ret = 0;
    char args[512];
    const AVFilter *buffersrc;
    AVFilterContext *buffersrc_ctx = NULL;
    enum AVMediaType type = avfilter_pad_get_type(in->filter_ctx->input_pads, in->pad_idx);

    if (type == AVMEDIA_TYPE_VIDEO) {
        buffersrc = avfilter_get_by_name("buffer");
        snprintf(args, sizeof(args),
                 "video_size=%dx%d:pix_fmt=%d:time_base=%d/%d:pixel_aspect=%d/%d",
                 320, 240, AV_PIX_FMT_YUV420P, 1, 1, 1, 1);
    } else if (type == AVMEDIA_TYPE_AUDIO) {
        buffersrc = avfilter_get_by_name("abuffer");
        const int64_t channel_layout = AV_CH_LAYOUT_STEREO;
        snprintf(args, sizeof(args),
                 "time_base=%d/%d:sample_rate=%d:sample_fmt=%s:channel_layout=0x%"PRIx64,
                 1, 1, 44100, av_get_sample_fmt_name(AV_SAMPLE_FMT_S16), channel_layout);
    } else {
        printf("Only video and audio filters are supported\n");
        return AVERROR(EINVAL);
    }

    if ((ret = avfilter_graph_create_filter(&buffersrc_ctx, buffersrc, "in",
            args, NULL, fc->graph)) < 0) {
        printf("Failed to create buffer source filter: %s\n", av_err2str(ret));
        return ret;
    }

    fc->nb_inputs++;
    fc->inputs = av_realloc_array(fc->inputs, fc->nb_inputs, sizeof(AVFilterContext*));
    fc->inputs[fc->nb_inputs - 1] = buffersrc_ctx;

    if ((ret = avfilter_link(buffersrc_ctx, 0, in->filter_ctx, in->pad_idx)) < 0) {
        printf("Could not link buffer source to graph: %s\n", av_err2str(ret));
        return ret;
    }

    return ret;
}

static int init_output_filter(FilteringContext *fc, AVFilterInOut *out)
{
    int ret = 0;
    const AVFilter *buffersink;
    AVFilterContext *buffersink_ctx = NULL;
    enum AVMediaType type = avfilter_pad_get_type(out->filter_ctx->input_pads, out->pad_idx);

    if (type == AVMEDIA_TYPE_VIDEO) {
        buffersink = avfilter_get_by_name("buffersink");
    } else if (type == AVMEDIA_TYPE_AUDIO) {
        buffersink = avfilter_get_by_name("abuffersink");
    } else {
        printf("Only video and audio filters are supported\n");
        return AVERROR(EINVAL);
    }

    if ((ret = avfilter_graph_create_filter(&buffersink_ctx, buffersink, "out",
            NULL, NULL, fc->graph)) < 0) {
        printf("Failed to create buffer sink filter: %s\n", av_err2str(ret));
        return ret;
    }

    fc->nb_outputs++;
    fc->outputs = av_realloc_array(fc->outputs, fc->nb_outputs, sizeof(AVFilterContext*));
    fc->outputs[fc->nb_outputs - 1] = buffersink_ctx;

    if ((ret = avfilter_link(out->filter_ctx, out->pad_idx, buffersink_ctx, 0)) < 0) {
        printf("Could not link buffer source to graph: %s\n", av_err2str(ret));
        return ret;
    }

    return ret;
}

static int init_filters(FilteringContext *fc)
{
    int ret = 0;
    AVFilterInOut *inputs = NULL;
    AVFilterInOut *outputs = NULL;

    fc->graph = avfilter_graph_alloc();
    if (!fc->graph)
        return -1;

    printf("Parsing: %s\n", fc->desc);
    if ((ret = avfilter_graph_parse2(fc->graph, fc->desc, &inputs, &outputs)) < 0) {
        printf("Failed to parse filter graph: %s\n", av_err2str(ret));
        goto end;
    }

    for (AVFilterInOut *current = inputs; current; current = current->next) {
        if ((ret = init_input_filter(fc, current)) < 0) {
            printf("Failed to create input for graph\n");
            goto end;
        }
    }
    for (AVFilterInOut *current = outputs; current; current = current->next) {
        if ((ret = init_output_filter(fc, current)) < 0) {
            printf("Failed to create output for graph\n");
            goto end;
        }
    }

    if ((ret = avfilter_graph_config(fc->graph, NULL)) < 0) {
        printf("Failed to configure graph: %s\n", av_err2str(ret));
        goto end;
    }

end:
    avfilter_inout_free(&inputs);
    avfilter_inout_free(&outputs);

    return ret;
}

static int pull_frame(AVFilterContext *sink_ctx, AVFrame **frame)
{
    int ret = 0;
    while (ret >= 0) {
        if (!*frame)
            *frame = av_frame_alloc();
        if (!*frame)
            return AVERROR(ENOMEM);

        ret = av_buffersink_get_frame_flags(sink_ctx, *frame, AV_BUFFERSINK_FLAG_NO_REQUEST);
        if (ret < 0) {
            if (ret != AVERROR(EAGAIN) && ret != AVERROR_EOF) {
                printf("Failed to pull from filter chain: %s (%d)\n", av_err2str(ret), ret);
            }
            break;
        }
    }

    return ret;
}

static int push_frame(AVFilterContext *src_ctx, AVFrame *frame)
{
    int ret = av_buffersrc_add_frame_flags(src_ctx, frame, AV_BUFFERSRC_FLAG_PUSH);
    if (ret < 0) {
        if (ret != AVERROR_EOF)
            printf("Failed to feed filter chain: %s (%d)\n", av_err2str(ret), ret);
        return ret;
    }

    return 0;
}

int main(int argc, char *argv[])
{
    int ret = 0;
    FilteringContext *fc = NULL;
    const char *filterspec = "[in1] scale=iw/4:ih/4 [mid1]; [in2] [mid1] overlay=main_w-overlay_w-10:main_h-overlay_h-10:shortest=1 [out1]";

    if (argc > 1 && argv[1]) {
        int level = atoi(argv[1]);
        av_log_set_level(level);
    }

    fc = av_mallocz(sizeof(*fc));
    if (!fc)
        goto end;
    fc->desc = av_strdup(filterspec);

    if ((ret = init_filters(fc)) < 0)
        goto end;

    //printf("Graph:\n%s\n", avfilter_graph_dump(fc->graph, NULL));

    for (int frame_index = 0; frame_index < 25; ++frame_index) {
        for (int i = 0; i < fc->nb_inputs; ++i) {
            AVFrame *input = get_dummy_frame(320, 240, frame_index);
            input->pts = frame_index;
            if ((ret = push_frame(fc->inputs[i], input)) < 0) {
                printf("Error pushing frame: %s\n", av_err2str(ret));
                av_frame_free(&input);
                break;
            }
            av_frame_free(&input);
        }

        for (int i = 0; i < fc->nb_outputs; ++i) {
            AVFrame *output;
            if ((ret = pull_frame(fc->outputs[i], &output)) < 0) {
                printf("Error pulling frame: %s\n", av_err2str(ret));
                av_frame_free(&output);
                break;
            }
        }
    }

    for (int i = 0; i < fc->nb_inputs; ++i) {
        if ((ret = push_frame(fc->inputs[i], NULL)) < 0) {
            printf("Error flushing filters: %s\n", av_err2str(ret));
        }
    }

end:
    if (fc) {
        avfilter_graph_free(&fc->graph);
        av_free(fc->inputs);
        av_free(fc->outputs);
        av_freep(&fc->desc);
    }
    av_free(fc);

    return (ret < 0 ? 1 : 0);
}