#include "media/access_unit_index.hpp"

#include <sstream>

#include "media/nal_inventory.hpp"

namespace smartcut::media {
namespace {
bool is_vcl(unsigned type) { return type <= 31U; }
bool is_closed_random_access(unsigned type) { return type >= 16U && type <= 20U; }
bool first_slice(const NalUnit& unit, std::span<const std::uint8_t> bytes) {
  std::size_t start = unit.offset;
  if (start + 4 <= bytes.size() && bytes[start] == 0 && bytes[start + 1] == 0 && bytes[start + 2] == 0 && bytes[start + 3] == 1) start += 4;
  else if (start + 3 <= bytes.size() && bytes[start] == 0 && bytes[start + 1] == 0 && bytes[start + 2] == 1) start += 3;
  return start + 2 < bytes.size() && (bytes[start + 2] & 0x80U) != 0;
}
}  // namespace

HevcAccessUnitIndex index_hevc_annex_b(std::span<const std::uint8_t> bytes) {
  HevcAccessUnitIndex result{bytes.size(), {}};
  const auto inventory = inventory_annex_b(Codec::hevc, bytes);
  bool have_current = false;
  for (const auto& unit : inventory.units) {
    if (is_vcl(unit.type) && first_slice(unit, bytes)) {
      result.access_units.push_back({unit.offset, unit.offset + unit.size, 1, unit.random_access, 0,
                                     unit.random_access ? unit.type : 0U,
                                     unit.random_access && is_closed_random_access(unit.type), unit.type});
      have_current = true;
      continue;
    }
    if (!have_current) continue;
    auto& current = result.access_units.back();
    current.last_nal_end = unit.offset + unit.size;
    ++current.nal_count;
    current.random_access = current.random_access || unit.random_access;
    if (unit.random_access) {
      current.random_access_nal_type = unit.type;
      current.closed_random_access = is_closed_random_access(unit.type);
    }
    current.rpu_count += unit.dolby_vision_rpu ? 1U : 0U;
  }
  return result;
}

std::string HevcAccessUnitIndex::to_json() const {
  std::ostringstream output;
  output << "{\"schema\":\"hevc-access-unit-index-v2\",\"input_bytes\":" << input_bytes << ",\"access_units\":[";
  for (std::size_t i = 0; i < access_units.size(); ++i) {
    if (i != 0) output << ',';
    const auto& unit = access_units[i];
    output << "{\"index\":" << i << ",\"first_nal_offset\":" << unit.first_nal_offset
           << ",\"last_nal_end\":" << unit.last_nal_end << ",\"nal_count\":" << unit.nal_count
           << ",\"random_access\":" << (unit.random_access ? "true" : "false")
           << ",\"rpu_count\":" << unit.rpu_count
           << ",\"random_access_nal_type\":" << unit.random_access_nal_type
           << ",\"closed_random_access\":" << (unit.closed_random_access ? "true" : "false")
           << ",\"vcl_nal_type\":" << unit.vcl_nal_type << '}';
  }
  output << "]}";
  return output.str();
}
}  // namespace smartcut::media
