// Adapted from
// https://github.com/leandromoreira/ffmpeg-libav-tutorial/blob/master/2_remuxing.c
// based on https://ffmpeg.org/doxygen/trunk/remuxing_8c-example.html
// FIXME: looks ugly, further refactor needed :(
#include <libavformat/avformat.h>
#include <libavutil/timestamp.h>
#include <stdio.h>

#include "../logger.h"
#include "ffmpeg.h"

int merge_av(const char *videofn, const char *audiofn, const char *outfn) {
  AVFormatContext *input1_format_context = NULL, *input2_format_context = NULL,
                  *output_format_context = NULL;
  AVPacket packet;
  int ret, i;
  int stream_index = 0;
  int *streams_list = NULL;
  int number_of_streams = 0;

  if ((ret = avformat_open_input(&input1_format_context, videofn, NULL, NULL)) <
      0) {
    append_log("Could not open input file '%s'\n", videofn);
    goto end;
  }
  if ((ret = avformat_open_input(&input2_format_context, audiofn, NULL, NULL)) <
      0) {
    append_log("Could not open input file '%s'\n", audiofn);
    goto end;
  }
  if ((ret = avformat_find_stream_info(input1_format_context, NULL)) < 0) {
    append_log("Failed to retrieve input stream information\n");
    goto end;
  }
  if ((ret = avformat_find_stream_info(input2_format_context, NULL)) < 0) {
    append_log("Failed to retrieve input stream information\n");
    goto end;
  }

  avformat_alloc_output_context2(&output_format_context, NULL, NULL, outfn);
  if (!output_format_context) {
    append_log("Could not create output context\n");
    ret = AVERROR_UNKNOWN;
    goto end;
  }

  number_of_streams =
      input1_format_context->nb_streams + input2_format_context->nb_streams;
  streams_list = av_malloc_array(number_of_streams, sizeof(*streams_list));

  if (!streams_list) {
    ret = AVERROR(ENOMEM);
    goto end;
  }

  for (i = 0; i < input1_format_context->nb_streams; i++) {
    AVStream *out_stream;
    AVStream *in_stream = input1_format_context->streams[i];
    AVCodecParameters *in_codecpar = in_stream->codecpar;
    if (in_codecpar->codec_type != AVMEDIA_TYPE_AUDIO &&
        in_codecpar->codec_type != AVMEDIA_TYPE_VIDEO &&
        in_codecpar->codec_type != AVMEDIA_TYPE_SUBTITLE) {
      streams_list[i] = -1;
      continue;
    }
    streams_list[i] = stream_index++;
    out_stream = avformat_new_stream(output_format_context, NULL);
    if (!out_stream) {
      append_log("Failed allocating output stream\n");
      ret = AVERROR_UNKNOWN;
      goto end;
    }
    ret = avcodec_parameters_copy(out_stream->codecpar, in_codecpar);
    if (ret < 0) {
      append_log("Failed to copy codec parameters\n");
      goto end;
    }
  }
  for (i = 0; i < input2_format_context->nb_streams; i++) {
    AVStream *out_stream;
    AVStream *in_stream = input2_format_context->streams[i];
    AVCodecParameters *in_codecpar = in_stream->codecpar;
    if (in_codecpar->codec_type != AVMEDIA_TYPE_AUDIO &&
        in_codecpar->codec_type != AVMEDIA_TYPE_VIDEO &&
        in_codecpar->codec_type != AVMEDIA_TYPE_SUBTITLE) {
      streams_list[i + stream_index] = -1;
      continue;
    }
    streams_list[i + stream_index] = stream_index;
    stream_index++;
    out_stream = avformat_new_stream(output_format_context, NULL);
    if (!out_stream) {
      append_log("Failed allocating output stream\n");
      ret = AVERROR_UNKNOWN;
      goto end;
    }
    ret = avcodec_parameters_copy(out_stream->codecpar, in_codecpar);
    if (ret < 0) {
      append_log("Failed to copy codec parameters\n");
      goto end;
    }
  }
#ifdef DEBUG
  av_dump_format(input1_format_context, 0, videofn, 0);
  av_dump_format(input2_format_context, 0, audiofn, 0);
  av_dump_format(output_format_context, 0, outfn, 1);
#endif
  if (!(output_format_context->oformat->flags & AVFMT_NOFILE)) {
    ret = avio_open(&output_format_context->pb, outfn, AVIO_FLAG_WRITE);
    if (ret < 0) {
      append_log("Could not open output file '%s'", outfn);
      goto end;
    }
  }
  AVDictionary *opts = NULL;
  /* Support fragmented MP4
   * https://developer.mozilla.org/en-US/docs/Web/API/Media_Source_Extensions_API/Transcoding_assets_for_MSE
   */
  av_dict_set(&opts, "movflags", "frag_keyframe+empty_moov+default_base_moof",
              0);
  ret = avformat_write_header(output_format_context, &opts);
  if (ret < 0) {
    append_log("Error occurred when opening output file\n");
    goto end;
  }
  while (1) {
    AVStream *in_stream, *out_stream;
    ret = av_read_frame(input1_format_context, &packet);
    if (ret < 0)
      break;
    in_stream = input1_format_context->streams[packet.stream_index];
    if (packet.stream_index >= number_of_streams ||
        streams_list[packet.stream_index] < 0) {
      av_packet_unref(&packet);
      continue;
    }
    packet.stream_index = streams_list[packet.stream_index];
    out_stream = output_format_context->streams[packet.stream_index];
    /* copy packet */
    packet.pts = av_rescale_q_rnd(packet.pts, in_stream->time_base,
                                  out_stream->time_base,
                                  AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
    packet.dts = av_rescale_q_rnd(packet.dts, in_stream->time_base,
                                  out_stream->time_base,
                                  AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
    packet.duration = av_rescale_q(packet.duration, in_stream->time_base,
                                   out_stream->time_base);
    packet.pos = -1;

    ret = av_interleaved_write_frame(output_format_context, &packet);
    if (ret < 0) {
      append_log("Error muxing packet\n");
      break;
    }
    av_packet_unref(&packet);
  }
  while (1) {
    AVStream *in_stream, *out_stream;
    ret = av_read_frame(input2_format_context, &packet);
    if (ret < 0)
      break;
    in_stream = input2_format_context->streams[packet.stream_index];
    if (packet.stream_index >= number_of_streams ||
        streams_list[packet.stream_index + input1_format_context->nb_streams] <
            0) {
      av_packet_unref(&packet);
      continue;
    }
    packet.stream_index =
        streams_list[packet.stream_index + input1_format_context->nb_streams];
    out_stream = output_format_context->streams[packet.stream_index];
    /* copy packet */
    packet.pts = av_rescale_q_rnd(packet.pts, in_stream->time_base,
                                  out_stream->time_base,
                                  AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
    packet.dts = av_rescale_q_rnd(packet.dts, in_stream->time_base,
                                  out_stream->time_base,
                                  AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
    packet.duration = av_rescale_q(packet.duration, in_stream->time_base,
                                   out_stream->time_base);
    packet.pos = -1;

    ret = av_interleaved_write_frame(output_format_context, &packet);
    if (ret < 0) {
      append_log("Error muxing packet\n");
      break;
    }
    av_packet_unref(&packet);
  }
  av_write_trailer(output_format_context);
end:
  avformat_close_input(&input1_format_context);
  avformat_close_input(&input2_format_context);
  /* close output */
  if (output_format_context &&
      !(output_format_context->oformat->flags & AVFMT_NOFILE))
    avio_closep(&output_format_context->pb);
  avformat_free_context(output_format_context);
  av_freep(&streams_list);
  if (ret < 0 && ret != AVERROR_EOF) {
    append_log("Error occurred: %s\n", av_err2str(ret));
    return 1;
  }

  // Delete seperate files
  if (remove(videofn) != 0) {
    append_log("Error deleting partial file %s\n", videofn);
  }
  if (remove(audiofn) != 0) {
    append_log("Error deleting partial file %s\n", audiofn);
  }
  return 0;
}