#include "media/mkv_roundtrip.hpp"

#include "media/sei_injector.hpp"

#include <cstring>
#include <span>
#include <vector>

#ifdef SMARTCUT_HAS_FFMPEG
extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/frame.h>
#include <libavutil/hdr_dynamic_metadata.h>
#include <libavutil/mastering_display_metadata.h>
#include <libavutil/mem.h>
#include <libavutil/opt.h>
}
#endif

namespace smartcut::media {
#ifdef SMARTCUT_HAS_FFMPEG
namespace {
std::vector<std::uint8_t> serialize_hdr10plus_scene(AVDynamicHDRPlus* metadata, int scene) {
  auto& params = metadata->params[0];
  const int scene_level = 450 + scene * 100;
  params.maxscl[0] = AVRational{scene_level, 1000};
  params.maxscl[1] = AVRational{scene_level - 50, 1000};
  params.maxscl[2] = AVRational{scene_level - 100, 1000};
  params.average_maxrgb = AVRational{scene_level - 150, 1000};
  params.distribution_maxrgb[0].percentile = AVRational{scene_level - 25, 1000};

  std::uint8_t* t35 = nullptr;
  std::size_t t35_size = 0;
  std::vector<std::uint8_t> sei;
  if (av_dynamic_hdr_plus_to_t35(metadata, &t35, &t35_size) >= 0) {
    sei = make_hevc_prefix_sei_t35(std::span<const std::uint8_t>(t35, t35_size));
  }
  av_free(t35);
  return sei;
}

int inject_hdr10plus(AVPacket* packet,
                     const std::vector<std::vector<std::uint8_t>>& sequence,
                     std::size_t& packet_index) {
  if (packet_index >= sequence.size() || sequence[packet_index].empty()) return -1;
  const auto& sei = sequence[packet_index++];
  const int old_size = packet->size;
  if (av_grow_packet(packet, static_cast<int>(sei.size())) < 0) return -1;
  std::memmove(packet->data + sei.size(), packet->data, static_cast<std::size_t>(old_size));
  std::memcpy(packet->data, sei.data(), sei.size());
  return 0;
}
}
#endif

int write_main10_test_mkv(const std::string& path, int width, int height, int frame_count, bool scene_change) {
#ifndef SMARTCUT_HAS_FFMPEG
  (void)path; (void)width; (void)height; (void)frame_count; (void)scene_change;
  return -1;
#else
  const AVCodec* codec = avcodec_find_encoder_by_name("libx265");
  if (!codec) return -2;
  AVCodecContext* encoder = avcodec_alloc_context3(codec);
  if (!encoder) return -3;
  if (width < 16 || height < 16 || frame_count < 1) { avcodec_free_context(&encoder); return -16; }
  encoder->width = width;
  encoder->height = height;
  encoder->time_base = AVRational{1, 25};
  encoder->framerate = AVRational{25, 1};
  encoder->pix_fmt = AV_PIX_FMT_YUV420P10LE;
  encoder->bit_rate = 1000000;
  encoder->gop_size = 25;
  encoder->max_b_frames = 0;
  encoder->color_range = AVCOL_RANGE_MPEG;
  encoder->color_primaries = AVCOL_PRI_BT2020;
  encoder->color_trc = AVCOL_TRC_SMPTE2084;
  encoder->colorspace = AVCOL_SPC_BT2020_NCL;
  encoder->chroma_sample_location = AVCHROMA_LOC_LEFT;
  av_opt_set(encoder->priv_data, "x265-params",
             "colorprim=bt2020:transfer=smpte2084:colormatrix=bt2020nc:range=limited:"
             "master-display=G(13250,34500)B(7500,3000)R(34000,16000)WP(15635,16450)L(10000000,100):"
             "max-cll=1000,400:hdr10=1:repeat-headers=1",
             0);

  AVFormatContext* output = nullptr;
  if (avformat_alloc_output_context2(&output, nullptr, "matroska", path.c_str()) < 0 || !output) {
    avcodec_free_context(&encoder);
    return -4;
  }
  if (output->oformat->flags & AVFMT_GLOBALHEADER) encoder->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
  AVDictionary* encoder_options = nullptr;
  if (width >= 3840 || height >= 2160) av_dict_set(&encoder_options, "preset", "ultrafast", 0);
  const int open_result = avcodec_open2(encoder, codec, &encoder_options);
  av_dict_free(&encoder_options);
  if (open_result < 0) {
    avformat_free_context(output);
    avcodec_free_context(&encoder);
    return -5;
  }
  AVStream* stream = avformat_new_stream(output, nullptr);
  if (!stream || avcodec_parameters_from_context(stream->codecpar, encoder) < 0) {
    avformat_free_context(output);
    avcodec_free_context(&encoder);
    return -6;
  }
  stream->time_base = encoder->time_base;
  if (!(output->oformat->flags & AVFMT_NOFILE) && avio_open(&output->pb, path.c_str(), AVIO_FLAG_WRITE) < 0) {
    avformat_free_context(output);
    avcodec_free_context(&encoder);
    return -7;
  }
  if (avformat_write_header(output, nullptr) < 0) {
    if (output->pb) avio_closep(&output->pb);
    avformat_free_context(output);
    avcodec_free_context(&encoder);
    return -8;
  }

  AVFrame* frame = av_frame_alloc();
  AVPacket* packet = av_packet_alloc();
  int result = 0;
  std::size_t packet_index = 0;
  std::vector<std::vector<std::uint8_t>> hdr10plus_sequence;
  AVDynamicHDRPlus* hdr10plus = frame ? av_dynamic_hdr_plus_create_side_data(frame) : nullptr;
  if (hdr10plus) {
    hdr10plus->itu_t_t35_country_code = 0xB5;
    hdr10plus->application_version = 0;
    hdr10plus->num_windows = 1;
    hdr10plus->targeted_system_display_maximum_luminance = AVRational{1000, 1};
    auto& params = hdr10plus->params[0];
    params.window_upper_left_corner_x = AVRational{0, 1};
    params.window_upper_left_corner_y = AVRational{0, 1};
    params.window_lower_right_corner_x = AVRational{1, 1};
    params.window_lower_right_corner_y = AVRational{1, 1};
    params.center_of_ellipse_x = static_cast<std::uint16_t>(width / 2);
    params.center_of_ellipse_y = static_cast<std::uint16_t>(height / 2);
    params.semimajor_axis_internal_ellipse = 1;
    params.semimajor_axis_external_ellipse = 1;
    params.semiminor_axis_external_ellipse = 1;
    params.num_distribution_maxrgb_percentiles = 1;
    params.distribution_maxrgb[0].percentage = 50;
    params.fraction_bright_pixels = AVRational{1, 100};
    params.tone_mapping_flag = 1;
    params.knee_point_x = AVRational{1, 2};
    params.knee_point_y = AVRational{1, 2};
    params.color_saturation_mapping_flag = 1;
    params.color_saturation_weight = AVRational{1, 1};
    for (int scene = 0; scene < frame_count; ++scene) {
      hdr10plus_sequence.push_back(serialize_hdr10plus_scene(hdr10plus, scene));
    }
  }

  if (!frame || !packet || hdr10plus_sequence.size() != static_cast<std::size_t>(frame_count)) {
    result = -9;
  } else {
    frame->format = encoder->pix_fmt;
    frame->width = encoder->width;
    frame->height = encoder->height;
    if (av_frame_get_buffer(frame, 32) < 0) {
      result = -10;
    } else {
      auto* mastering = av_mastering_display_metadata_create_side_data(frame);
      if (mastering) {
        mastering->display_primaries[0][0] = AVRational{13250, 100000};
        mastering->display_primaries[0][1] = AVRational{34500, 100000};
        mastering->display_primaries[1][0] = AVRational{7500, 100000};
        mastering->display_primaries[1][1] = AVRational{3000, 100000};
        mastering->display_primaries[2][0] = AVRational{34000, 100000};
        mastering->display_primaries[2][1] = AVRational{16000, 100000};
        mastering->white_point[0] = AVRational{3127, 10000};
        mastering->white_point[1] = AVRational{3290, 10000};
        mastering->min_luminance = AVRational{1, 100};
        mastering->max_luminance = AVRational{1000, 1};
        mastering->has_primaries = 1;
        mastering->has_luminance = 1;
      }
      auto* light = av_content_light_metadata_create_side_data(frame);
      if (light) {
        light->MaxCLL = 1000;
        light->MaxFALL = 400;
      }
      for (int frame_index = 0; frame_index < frame_count && result == 0; ++frame_index) {
        if (av_frame_make_writable(frame) < 0) {
          result = -10;
          break;
        }
        for (int plane = 0; plane < 3; ++plane) {
          const int rows = plane == 0 ? frame->height : (frame->height + 1) / 2;
          std::memset(frame->data[plane], 0x80,
                      static_cast<std::size_t>(frame->linesize[plane]) * static_cast<std::size_t>(rows));
        }
        if (scene_change) {
          const std::uint16_t luma = frame_index < frame_count / 2 ? 64U : 900U;
          for (int row = 0; row < frame->height; ++row) {
            auto* pixels = reinterpret_cast<std::uint16_t*>(frame->data[0] + row * frame->linesize[0]);
            std::fill(pixels, pixels + frame->width, luma);
          }
        }
        frame->pts = frame_index;
        if (avcodec_send_frame(encoder, frame) < 0) {
          result = -11;
          break;
        }
        while (avcodec_receive_packet(encoder, packet) >= 0) {
          if (inject_hdr10plus(packet, hdr10plus_sequence, packet_index) < 0) result = -15;
          packet->stream_index = stream->index;
          av_packet_rescale_ts(packet, encoder->time_base, stream->time_base);
          if (result == 0 && av_interleaved_write_frame(output, packet) < 0) result = -12;
          av_packet_unref(packet);
        }
      }
      avcodec_send_frame(encoder, nullptr);
      while (result == 0 && avcodec_receive_packet(encoder, packet) >= 0) {
        if (inject_hdr10plus(packet, hdr10plus_sequence, packet_index) < 0) result = -15;
        packet->stream_index = stream->index;
        av_packet_rescale_ts(packet, encoder->time_base, stream->time_base);
        if (result == 0 && av_interleaved_write_frame(output, packet) < 0) result = -13;
        av_packet_unref(packet);
      }
      if (result == 0 && packet_index != hdr10plus_sequence.size()) result = -15;
    }
  }

  av_frame_free(&frame);
  av_packet_free(&packet);
  if (result == 0 && av_write_trailer(output) < 0) result = -14;
  if (output->pb) avio_closep(&output->pb);
  avformat_free_context(output);
  avcodec_free_context(&encoder);
  return result;
#endif
}
}
