#include "media/boundary_fixture.hpp"

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <string>

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

namespace smartcut::media {
#ifdef SMARTCUT_HAS_FFMPEG
namespace {

const char* profile_for(AVCodecID codec, AVPixelFormat format) {
  const auto* descriptor = av_pix_fmt_desc_get(format);
  if (!descriptor) return nullptr;
  const int depth = descriptor->comp[0].depth;
  const bool chroma_444 = descriptor->log2_chroma_w == 0 && descriptor->log2_chroma_h == 0;
  const bool chroma_422 = descriptor->log2_chroma_w == 1 && descriptor->log2_chroma_h == 0;
  if (codec == AV_CODEC_ID_H264) {
    if (chroma_444) return "high444";
    if (chroma_422) return "high422";
    return depth > 8 ? "high10" : "high";
  }
  if (codec == AV_CODEC_ID_HEVC) {
    if (chroma_444) return depth > 8 ? "main444-10" : "main444-8";
    if (chroma_422) return "main422-10";
    return depth > 8 ? "main10" : "main";
  }
  return nullptr;
}

void fill_frame(AVFrame* frame, int index) {
  const auto* descriptor = av_pix_fmt_desc_get(static_cast<AVPixelFormat>(frame->format));
  if (!descriptor) return;
  for (int plane = 0; plane < 4 && frame->data[plane]; ++plane) {
    const int rows = plane == 0 || plane == 3
                         ? frame->height
                         : AV_CEIL_RSHIFT(frame->height, descriptor->log2_chroma_h);
    std::memset(frame->data[plane], plane == 0 ? 32 + index * 7 : 0x80,
                static_cast<std::size_t>(frame->linesize[plane]) * static_cast<std::size_t>(rows));
  }
}

}  // namespace
#endif

int write_boundary_format_fixture(const std::string& path,
                                  const std::string& codec_name,
                                  const std::string& pixel_format_name,
                                  const std::string& transfer,
                                  bool captions,
                                  int frame_count) {
#ifndef SMARTCUT_HAS_FFMPEG
  (void)path; (void)codec_name; (void)pixel_format_name; (void)transfer; (void)captions;
  (void)frame_count;
  return -1;
#else
  const AVCodecID codec_id = codec_name == "h264" ? AV_CODEC_ID_H264
                              : codec_name == "hevc" ? AV_CODEC_ID_HEVC
                                                       : AV_CODEC_ID_NONE;
  const AVPixelFormat pixel_format = av_get_pix_fmt(pixel_format_name.c_str());
  if (codec_id == AV_CODEC_ID_NONE || pixel_format == AV_PIX_FMT_NONE || frame_count < 3) return -2;
  const char* encoder_name = codec_id == AV_CODEC_ID_H264 ? "libx264" : "libx265";
  const AVCodec* codec = avcodec_find_encoder_by_name(encoder_name);
  AVCodecContext* encoder = codec ? avcodec_alloc_context3(codec) : nullptr;
  AVFormatContext* output = nullptr;
  AVFrame* frame = av_frame_alloc();
  AVPacket* packet = av_packet_alloc();
  bool completed = false;
  auto cleanup = [&] {
    av_frame_free(&frame);
    av_packet_free(&packet);
    avcodec_free_context(&encoder);
    if (output && output->pb) avio_closep(&output->pb);
    avformat_free_context(output);
  };
  if (!encoder || !frame || !packet) { cleanup(); return -3; }
  encoder->width = 96;
  encoder->height = 64;
  encoder->time_base = AVRational{1, 25};
  encoder->framerate = AVRational{25, 1};
  encoder->pix_fmt = pixel_format;
  encoder->bit_rate = 1'000'000;
  encoder->gop_size = 8;
  encoder->max_b_frames = 0;
  encoder->sample_aspect_ratio = AVRational{1, 1};
  encoder->color_range = AVCOL_RANGE_MPEG;
  encoder->color_primaries = transfer == "hlg" || transfer == "pq" ? AVCOL_PRI_BT2020 : AVCOL_PRI_BT709;
  encoder->color_trc = transfer == "hlg" ? AVCOL_TRC_ARIB_STD_B67
                         : transfer == "pq" ? AVCOL_TRC_SMPTE2084
                                              : AVCOL_TRC_BT709;
  encoder->colorspace = transfer == "hlg" || transfer == "pq" ? AVCOL_SPC_BT2020_NCL : AVCOL_SPC_BT709;
  encoder->chroma_sample_location = AVCHROMA_LOC_LEFT;
  if (const auto* profile = profile_for(codec_id, pixel_format))
    av_opt_set(encoder->priv_data, "profile", profile, 0);
  av_opt_set(encoder->priv_data, "preset", codec_id == AV_CODEC_ID_HEVC ? "ultrafast" : "veryfast", 0);
  av_opt_set(encoder->priv_data, "a53cc", captions ? "1" : "0", 0);
  if (codec_id == AV_CODEC_ID_HEVC) {
    const std::string params = "repeat-headers=1:colorprim=" +
        std::to_string(static_cast<int>(encoder->color_primaries)) + ":transfer=" +
        std::to_string(static_cast<int>(encoder->color_trc)) + ":colormatrix=" +
        std::to_string(static_cast<int>(encoder->colorspace)) + ":range=limited";
    av_opt_set(encoder->priv_data, "x265-params", params.c_str(), 0);
  }
  if (avformat_alloc_output_context2(&output, nullptr, "matroska", path.c_str()) < 0 || !output) {
    cleanup(); return -4;
  }
  if (output->oformat->flags & AVFMT_GLOBALHEADER) encoder->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
  if (avcodec_open2(encoder, codec, nullptr) < 0) { cleanup(); return -5; }
  AVStream* stream = avformat_new_stream(output, nullptr);
  if (!stream || avcodec_parameters_from_context(stream->codecpar, encoder) < 0) {
    cleanup(); return -6;
  }
  stream->time_base = encoder->time_base;
  if (avio_open(&output->pb, path.c_str(), AVIO_FLAG_WRITE) < 0 ||
      avformat_write_header(output, nullptr) < 0) {
    cleanup(); return -7;
  }
  frame->format = encoder->pix_fmt;
  frame->width = encoder->width;
  frame->height = encoder->height;
  frame->sample_aspect_ratio = encoder->sample_aspect_ratio;
  frame->color_range = encoder->color_range;
  frame->color_primaries = encoder->color_primaries;
  frame->color_trc = encoder->color_trc;
  frame->colorspace = encoder->colorspace;
  frame->chroma_location = encoder->chroma_sample_location;
  if (av_frame_get_buffer(frame, 32) < 0) { cleanup(); return -8; }
  AVFrameSideData* caption_data = captions
      ? av_frame_new_side_data(frame, AV_FRAME_DATA_A53_CC, 3)
      : nullptr;
  if (captions && !caption_data) { cleanup(); return -9; }
  auto drain = [&]() {
    while (true) {
      const int received = avcodec_receive_packet(encoder, packet);
      if (received == AVERROR(EAGAIN) || received == AVERROR_EOF) return 0;
      if (received < 0) return received;
      packet->stream_index = stream->index;
      av_packet_rescale_ts(packet, encoder->time_base, stream->time_base);
      const int written = av_interleaved_write_frame(output, packet);
      av_packet_unref(packet);
      if (written < 0) return written;
    }
  };
  int result = 0;
  for (int index = 0; index < frame_count && result >= 0; ++index) {
    if (av_frame_make_writable(frame) < 0) { result = -10; break; }
    fill_frame(frame, index);
    if (caption_data) {
      caption_data->data[0] = 0xfc;
      caption_data->data[1] = static_cast<std::uint8_t>(0x10 + index);
      caption_data->data[2] = static_cast<std::uint8_t>(0x20 + index);
    }
    frame->pts = index;
    result = avcodec_send_frame(encoder, frame);
    if (result >= 0) result = drain();
  }
  if (result >= 0) {
    result = avcodec_send_frame(encoder, nullptr);
    if (result >= 0) result = drain();
  }
  if (result >= 0) result = av_write_trailer(output);
  completed = result >= 0;
  cleanup();
  return completed ? 0 : -11;
#endif
}

}  // namespace smartcut::media
