#include "media/rpu_injector.hpp"

#ifdef SMARTCUT_HAS_LIBDOVI
extern "C" {
#include <libdovi/rpu_parser.h>
}
#endif

namespace smartcut::media {
std::vector<std::uint8_t> make_hevc_rpu_nal(std::span<const std::uint8_t> raw_rpu) {
  std::vector<std::uint8_t> nal{0,0,0,1,0x7c,0x01}; nal.insert(nal.end(), raw_rpu.begin(), raw_rpu.end()); return nal;
}
std::vector<std::uint8_t> prepend_hevc_rpu(std::span<const std::uint8_t> access_unit, std::span<const std::uint8_t> raw_rpu) {
  auto result=make_hevc_rpu_nal(raw_rpu); result.insert(result.end(), access_unit.begin(), access_unit.end()); return result;
}

bool rewrite_dovi_rpu_as_scene_refresh(std::span<const std::uint8_t> annex_b_nal,
                                      std::vector<std::uint8_t>& rewritten,
                                      std::string& error) {
#ifndef SMARTCUT_HAS_LIBDOVI
  (void)annex_b_nal;
  rewritten.clear();
  error = "libdovi support is not enabled";
  return false;
#else
  rewritten.clear();
  error.clear();
  std::size_t offset = 0;
  if (annex_b_nal.size() >= 4 && annex_b_nal[0] == 0 && annex_b_nal[1] == 0 && annex_b_nal[2] == 0 && annex_b_nal[3] == 1) offset = 4;
  else if (annex_b_nal.size() >= 3 && annex_b_nal[0] == 0 && annex_b_nal[1] == 0 && annex_b_nal[2] == 1) offset = 3;
  if (annex_b_nal.size() < offset + 3 || ((annex_b_nal[offset] >> 1U) & 0x3fU) != 62U) {
    error = "input is not a Dolby Vision HEVC UNSPEC62 NAL";
    return false;
  }

  DoviRpuOpaque* rpu = dovi_parse_unspec62_nalu(annex_b_nal.data() + offset, annex_b_nal.size() - offset);
  if (!rpu) {
    error = "libdovi did not allocate an RPU parser";
    return false;
  }
  const auto release_rpu = [&] { dovi_rpu_free(rpu); };
  if (const char* parse_error = dovi_rpu_get_error(rpu)) {
    error = parse_error;
    release_rpu();
    return false;
  }
  if (dovi_rpu_set_scene_refresh_flag(rpu, true) != 0) {
    error = dovi_rpu_get_error(rpu) ? dovi_rpu_get_error(rpu) : "libdovi scene-refresh rewrite failed";
    release_rpu();
    return false;
  }
  const DoviData* encoded = dovi_write_unspec62_nalu(rpu);
  if (!encoded || !encoded->data || encoded->len == 0) {
    error = dovi_rpu_get_error(rpu) ? dovi_rpu_get_error(rpu) : "libdovi RPU serialization failed";
    if (encoded) dovi_data_free(encoded);
    release_rpu();
    return false;
  }
  rewritten = {0, 0, 0, 1};
  rewritten.insert(rewritten.end(), encoded->data, encoded->data + encoded->len);
  dovi_data_free(encoded);
  release_rpu();
  return true;
#endif
}
}
