/*

 * HEIF VVC codec.

 * Copyright (c) 2023 Dirk Farin <dirk.farin@gmail.com>

 *

 * This file is part of libheif.

 *

 * libheif is free software: you can redistribute it and/or modify

 * it under the terms of the GNU Lesser General Public License as

 * published by the Free Software Foundation, either version 3 of

 * the License, or (at your option) any later version.

 *

 * libheif is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public License

 * along with libheif.  If not, see <http://www.gnu.org/licenses/>.

 */

#include "vvc_boxes.h"

#include "hevc_boxes.h"

#include "codecs/decoder.h"

#include "file.h"

#include <cstring>

#include <string>

#include <cassert>

#include <iomanip>

#include <utility>

#include "api_structs.h"

Error Box_vvcC::parse(BitstreamRange& range, const heif_security_limits* limits)

{

  parse_full_box_header(range);

  uint8_t byte;

  auto& c = m_configuration; // abbreviation

  byte = range.read8();

  c.LengthSizeMinusOne = (byte >> 1) & 3;

  c.ptl_present_flag = !!(byte & 1);

  if (c.ptl_present_flag) {

    uint16_t word = range.read16();

    c.ols_idx = (word >> 7) & 0x1FF;

    c.num_sublayers = (word >> 4) & 0x07;

    c.constant_frame_rate = (word >> 2) & 0x03;

    c.chroma_format_idc = word & 0x03;

    byte = range.read8();

    c.bit_depth_minus8 = (byte >> 5) & 0x07;

    // VvcPTLRecord

    auto& ptl = c.native_ptl; // abbreviation

    byte = range.read8();

    ptl.num_bytes_constraint_info = byte & 0x3f;

    if (ptl.num_bytes_constraint_info == 0) {

      return {heif_error_Invalid_input,

              heif_suberror_Invalid_parameter_value,

              "vvcC with num_bytes_constraint_info==0 is not allowed."};

    }

    byte = range.read8();

    ptl.general_profile_idc = (byte >> 1) & 0x7f;

    ptl.general_tier_flag = (byte & 1);

    ptl.general_level_idc = range.read8();

    for (int i = 0; i < ptl.num_bytes_constraint_info; i++) {

      byte = range.read8();

      if (i == 0) {

        ptl.ptl_frame_only_constraint_flag = (byte >> 7) & 1;

        ptl.ptl_multi_layer_enabled_flag = (byte >> 6) & 1;

        byte &= 0x3f;

      }

      ptl.general_constraint_info.push_back(byte);

    }

    if (c.num_sublayers > 1) {

      ptl.ptl_sublayer_level_present_flag.resize(c.num_sublayers - 1);

      byte = range.read8();

      uint8_t mask = 0x80;

      for (int i = c.num_sublayers - 2; i >= 0; i--) {

        ptl.ptl_sublayer_level_present_flag[i] = !!(byte & mask);

        mask >>= 1;

      }

    }

    ptl.sublayer_level_idc.resize(c.num_sublayers);

    if (c.num_sublayers > 0) {

      ptl.sublayer_level_idc[c.num_sublayers - 1] = ptl.general_level_idc;

      for (int i = c.num_sublayers - 2; i >= 0; i--) {

        if (ptl.ptl_sublayer_level_present_flag[i]) {

          ptl.sublayer_level_idc[i] = range.read8();

        }

        else {

          ptl.sublayer_level_idc[i] = ptl.sublayer_level_idc[i + 1];

        }

      }

    }

    uint8_t ptl_num_sub_profiles = range.read8();

    for (int j=0; j < ptl_num_sub_profiles; j++) {

      ptl.general_sub_profile_idc.push_back(range.read32());

    }

    // remaining fields

    c.max_picture_width = range.read16();

    c.max_picture_height = range.read16();

    c.avg_frame_rate = range.read16();

  }

  else {

    return Error{

      heif_error_Unsupported_feature,

      heif_suberror_Unspecified,

      "Reading vvcC configuration with ptl_present_flag=0 is not supported."

    };

  }

  // read NAL arrays

  int nArrays = range.read8();

  for (int i = 0; i < nArrays && !range.error(); i++) {

    byte = range.read8();

    NalArray array;

    array.m_array_completeness = (byte >> 7) & 1;

    array.m_NAL_unit_type = (byte & 0x3F);

    int nUnits = range.read16();

    for (int u = 0; u < nUnits && !range.error(); u++) {

      std::vector<uint8_t> nal_unit;

      int size = range.read16();

      if (!size) {

        // Ignore empty NAL units.

        continue;

      }

      if (range.prepare_read(size)) {

        nal_unit.resize(size);

        bool success = range.get_istream()->read((char*) nal_unit.data(), size);

        if (!success) {

          return Error{heif_error_Invalid_input, heif_suberror_End_of_data, "error while reading hvcC box"};

        }

      }

      array.m_nal_units.push_back(std::move(nal_unit));

    }

    m_nal_array.push_back(std::move(array));

  }

  return range.get_error();

}

bool Box_vvcC::get_headers(std::vector<uint8_t>* dest) const

{

  for (const auto& nal_array : m_nal_array) {

    for (const auto& nal : nal_array.m_nal_units) {

      assert(nal.size() <= 0xFFFF);

      auto size = static_cast<uint16_t>(nal.size());

      dest->push_back(0);

      dest->push_back(0);

      dest->push_back(static_cast<uint8_t>(size >> 8));

      dest->push_back(static_cast<uint8_t>(size & 0xFF));

      dest->insert(dest->end(), nal.begin(), nal.end());

    }

  }

  return true;

}

void Box_vvcC::append_nal_data(const std::vector<uint8_t>& nal)

{

  assert(nal.size()>=2);

  uint8_t nal_type = (nal[1] >> 3) & 0x1F;

  // insert into existing array if it exists

  for (auto& nalarray : m_nal_array) {

    if (nalarray.m_NAL_unit_type == nal_type) {

      nalarray.m_nal_units.push_back(nal);

      return;

    }

  }

  // generate new NAL array

  NalArray array;

  array.m_array_completeness = true;

  array.m_NAL_unit_type = uint8_t((nal[1] >> 3) & 0x1F);

  array.m_nal_units.push_back(nal);

  m_nal_array.push_back(array);

}

const std::vector<uint8_t>* Box_vvcC::get_first_nal_of_type(uint8_t nal_type) const

{

  for (const auto& arr : m_nal_array) {

    if (arr.m_NAL_unit_type == nal_type && !arr.m_nal_units.empty()) {

      return &arr.m_nal_units[0];

    }

  }

  return nullptr;

}

void Box_vvcC::append_nal_data(const uint8_t* data, size_t size)

{

  std::vector<uint8_t> nal;

  nal.resize(size);

  memcpy(nal.data(), data, size);

  append_nal_data(nal);

}

Error Box_vvcC::write(StreamWriter& writer) const

{

  size_t box_start = reserve_box_header_space(writer);

  const auto& c = m_configuration;

  uint8_t byte;

  byte = uint8_t(0xF8 | (c.LengthSizeMinusOne<<1) | (c.ptl_present_flag ? 1 : 0));

  writer.write8(byte);

  if (c.ptl_present_flag) {

    assert(c.ols_idx <= 0x1FF);

    assert(c.num_sublayers <= 7);

    assert(c.constant_frame_rate <= 3);

    assert(c.chroma_format_idc <= 3);

    assert(c.bit_depth_minus8 <= 7);

    auto word = uint16_t((c.ols_idx << 7) | (c.num_sublayers << 4) | (c.constant_frame_rate << 2) | (c.chroma_format_idc));

    writer.write16(word);

    writer.write8(uint8_t((c.bit_depth_minus8<<5) | 0x1F));

    const auto& ptl = c.native_ptl;

    assert(ptl.general_profile_idc <= 0x7F);

    writer.write8(ptl.num_bytes_constraint_info & 0x3f);

    writer.write8(static_cast<uint8_t>((ptl.general_profile_idc<<1) | ptl.general_tier_flag));

    writer.write8(ptl.general_level_idc);

    for (int i=0;i<ptl.num_bytes_constraint_info;i++) {

      if (i==0) {

        assert(ptl.ptl_frame_only_constraint_flag <= 1);

        assert(ptl.ptl_multi_layer_enabled_flag <= 1);

        assert(ptl.general_constraint_info[0] <= 0x3F);

        byte = static_cast<uint8_t>((ptl.ptl_frame_only_constraint_flag << 7) | (ptl.ptl_multi_layer_enabled_flag << 6) | ptl.general_constraint_info[0]);

      }

      else {

        byte = ptl.general_constraint_info[i];

      }

      writer.write8(byte);

    }

    byte = 0;

    if (c.num_sublayers > 1) {

      uint8_t mask = 0x80;

      for (int i = c.num_sublayers - 2; i >= 0; i--) {

        if (ptl.ptl_sublayer_level_present_flag[i]) {

          byte |= mask;

        }

        mask >>= 1;

      }

      writer.write8(byte);

    }

    for (int i=c.num_sublayers-2; i >= 0; i--) {

      if (ptl.ptl_sublayer_level_present_flag[i]) {

        writer.write8(ptl.sublayer_level_idc[i]);

      }

    }

    assert(ptl.general_sub_profile_idc.size() <= 0xFF);

    byte = static_cast<uint8_t>(ptl.general_sub_profile_idc.size());

    writer.write8(byte);

    for (int j=0; j < byte; j++) {

      writer.write32(ptl.general_sub_profile_idc[j]);

    }

    writer.write16(c.max_picture_width);

    writer.write16(c.max_picture_height);

    writer.write16(c.avg_frame_rate);

  }

  // --- write configuration NALs

  if (m_nal_array.size() > 255) {

    return {heif_error_Encoding_error, heif_suberror_Unspecified, "Too many VVC NAL arrays."};

  }

  writer.write8((uint8_t)m_nal_array.size());

  for (const NalArray& nal_array : m_nal_array) {

    uint8_t v2 = (nal_array.m_array_completeness ? 0x80 : 0);

    v2 |= nal_array.m_NAL_unit_type;

    writer.write8(v2);

    if (nal_array.m_nal_units.size() > 0xFFFF) {

      return {heif_error_Encoding_error, heif_suberror_Unspecified, "Too many VVC NAL units."};

    }

    writer.write16((uint16_t)nal_array.m_nal_units.size());

    for (const auto& nal : nal_array.m_nal_units) {

      if (nal.size() > 0xFFFF) {

        return {heif_error_Encoding_error, heif_suberror_Unspecified, "VVC NAL too large."};

      }

      writer.write16((uint16_t)nal.size());

      writer.write(nal);

    }

  }

  prepend_header(writer, box_start);

  return Error::Ok;

}

static const char* vvc_chroma_names[4] = {"mono", "4:2:0", "4:2:2", "4:4:4"};

const char* NAL_name(uint8_t nal_type)

{

  switch (nal_type) {

    case 12: return "OPI";

    case 13: return "DCI";

    case 14: return "VPS";

    case 15: return "SPS";

    case 16: return "PPS";

    case 17: return "PREFIX_APS";

    case 18: return "SUFFIX_APS";

    case 19: return "PH";

    default: return "?";

  }

}

std::string Box_vvcC::dump(Indent& indent) const

{

  std::ostringstream sstr;

  sstr << FullBox::dump(indent);

  const auto& c = m_configuration; // abbreviation

  sstr << indent << "NAL length size: " << ((int) c.LengthSizeMinusOne + 1) << "\n";

  if (c.ptl_present_flag) {

    const auto& ptl = c.native_ptl;

    sstr << indent << "ols-index: " << c.ols_idx << "\n"

         << indent << "num sublayers: " << ((int) c.num_sublayers) << "\n"

         << indent << "constant frame rate: " << (c.constant_frame_rate == 1 ? "constant" : (c.constant_frame_rate == 2 ? "multi-layer" : "unknown")) << "\n"

         << indent << "chroma-format: " << vvc_chroma_names[c.chroma_format_idc] << "\n"

         << indent << "bit-depth: " << ((int) c.bit_depth_minus8 + 8) << "\n"

         << indent << "max picture width:  " << c.max_picture_width << "\n"

         << indent << "max picture height: " << c.max_picture_height << "\n";

    sstr << indent << "general profile: " << ((int)ptl.general_profile_idc) << "\n"

         << indent << "tier flag: " << ((int)ptl.general_tier_flag) << "\n"

         << indent << "general level:" << ((int)ptl.general_level_idc) << "\n"

         << indent << "ptl frame only constraint flag: " << ((int)ptl.ptl_frame_only_constraint_flag) << "\n"

         << indent << "ptl multi layer enabled flag: " << ((int)ptl.ptl_multi_layer_enabled_flag) << "\n";

  }

  sstr << indent << "num of arrays: " << m_nal_array.size() << "\n";

  sstr << indent << "config NALs:\n";

  for (const auto& nal_array : m_nal_array) {

    indent++;

    sstr << indent << "NAL type: " << ((int)nal_array.m_NAL_unit_type) << " (" << NAL_name(nal_array.m_NAL_unit_type) << ")\n";

    sstr << indent << "array completeness: " << ((int)nal_array.m_array_completeness) << "\n";

    for (const auto& nal : nal_array.m_nal_units) {

      indent++;

      std::string ind = indent.get_string();

      sstr << write_raw_data_as_hex(nal.data(), nal.size(), ind, ind);

      indent--;

    }

    indent--;

  }

  return sstr.str();

}

Error parse_sps_for_vvcC_configuration(const uint8_t* sps, size_t size,

                                       Box_vvcC::configuration* config,

                                       uint32_t* width, uint32_t* height,

                                       ImageSize* coded_size)

{

  // remove start-code emulation bytes from SPS header stream

  std::vector<uint8_t> sps_no_emul = remove_start_code_emulation(sps, size);

  sps = sps_no_emul.data();

  size = sps_no_emul.size();

  BitReader reader(sps, (int) size);

  // skip NAL header

  reader.skip_bits(2 * 8);

  // skip SPS ID

  reader.skip_bits(4);

  // skip VPS ID

  reader.skip_bits(4);

  config->ols_idx = 0;

  config->num_sublayers = reader.get_bits8(3) + 1;

  config->chroma_format_idc = reader.get_bits8(2);

  reader.skip_bits(2);

  bool sps_ptl_dpb_hrd_params_present_flag = reader.get_bits(1);

  if (sps_ptl_dpb_hrd_params_present_flag) {

    // profile_tier_level( 1, sps_max_sublayers_minus1 )

    auto& ptl = config->native_ptl;

    if (true /*profileTierPresentFlag*/) {

      ptl.general_profile_idc = reader.get_bits8(7);

      ptl.general_tier_flag = reader.get_bits8(1);

    }

    ptl.general_level_idc = reader.get_bits8(8);

    ptl.ptl_frame_only_constraint_flag = reader.get_bits8(1);

    ptl.ptl_multi_layer_enabled_flag = reader.get_bits8(1);

    if (true /* profileTierPresentFlag*/ ) {

      // general_constraints_info()

      bool gci_present_flag = reader.get_bits(1);

      if (gci_present_flag) {

        // TODO

        return {heif_error_Unsupported_feature,

                heif_suberror_Unsupported_data_version,

                "VVC SPS with general_constraints_info is not supported yet"};

      }

      else {

        ptl.num_bytes_constraint_info = 1;

        ptl.general_constraint_info.push_back(0);

      }

      reader.skip_to_byte_boundary();

    }

    ptl.ptl_sublayer_level_present_flag.resize(config->num_sublayers);

    for (int i = config->num_sublayers-2; i >= 0; i--) {

      ptl.ptl_sublayer_level_present_flag[i] = reader.get_bits(1);

    }

    reader.skip_to_byte_boundary();

    ptl.sublayer_level_idc.resize(config->num_sublayers);

    for (int i = config->num_sublayers-2; i >= 0; i--) {

      if (ptl.ptl_sublayer_level_present_flag[i]) {

        ptl.sublayer_level_idc[i] = reader.get_bits8(8);

      }

    }

    if (true /*profileTierPresentFlag*/) {

      int ptl_num_sub_profiles = reader.get_bits(8);

      ptl.general_sub_profile_idc.resize(ptl_num_sub_profiles);

      for (int i = 0; i < ptl_num_sub_profiles; i++) {

        ptl.general_sub_profile_idc[i] = reader.get_bits(32);

      }

    }

  }

  reader.skip_bits(1); // sps_gdr_enabled_flag

  bool sps_ref_pic_resampling_enabled_flag = reader.get_bits(1);

  if (sps_ref_pic_resampling_enabled_flag) {

    reader.skip_bits(1); // sps_res_change_in_clvs_allowed_flag

  }

  Error invalidUVLC{

    heif_error_Invalid_input,

    heif_suberror_Invalid_parameter_value,

    "Invalid variable length code in VVC SPS header"

  };

  uint32_t sps_pic_width_max_in_luma_samples;

  uint32_t sps_pic_height_max_in_luma_samples;

  if (!reader.get_uvlc(&sps_pic_width_max_in_luma_samples) ||

      !reader.get_uvlc(&sps_pic_height_max_in_luma_samples)) {

    return invalidUVLC;

  }

  *width = sps_pic_width_max_in_luma_samples;

  *height = sps_pic_height_max_in_luma_samples;

  if (coded_size) {

    coded_size->width = *width;

    coded_size->height = *height;

  }

  if (sps_pic_width_max_in_luma_samples > 0xFFFF ||

      sps_pic_height_max_in_luma_samples > 0xFFFF) {

    return {heif_error_Encoding_error,

            heif_suberror_Invalid_parameter_value,

            "SPS max picture width or height exceeds maximum (65535)"};

  }

  config->max_picture_width = static_cast<uint16_t>(sps_pic_width_max_in_luma_samples);

  config->max_picture_height = static_cast<uint16_t>(sps_pic_height_max_in_luma_samples);

  int sps_conformance_window_flag = reader.get_bits(1);

  if (sps_conformance_window_flag) {

    uint32_t left, right, top, bottom;

    if (!reader.get_uvlc(&left) ||

        !reader.get_uvlc(&right) ||

        !reader.get_uvlc(&top) ||

        !reader.get_uvlc(&bottom)) {

      return invalidUVLC;

    }

    // SubWidthC / SubHeightC per ITU-T H.266 Table 5-1, indexed by chroma_format_idc.

    uint32_t subWidthC = 1, subHeightC = 1;

    switch (config->chroma_format_idc) {

      case 1: subWidthC = 2; subHeightC = 2; break;  // 4:2:0

      case 2: subWidthC = 2; subHeightC = 1; break;  // 4:2:2

      default: break;                                // mono / 4:4:4

    }

    const uint64_t crop_w = (uint64_t)subWidthC  * ((uint64_t)left + (uint64_t)right);

    const uint64_t crop_h = (uint64_t)subHeightC * ((uint64_t)top  + (uint64_t)bottom);

    if (crop_w > *width || crop_h > *height) {

      return {heif_error_Invalid_input,

              heif_suberror_Invalid_parameter_value,

              "SPS conformance window exceeds image dimensions"};

    }

    *width  -= (uint32_t)crop_w;

    *height -= (uint32_t)crop_h;

  }

  bool sps_subpic_info_present_flag = reader.get_bits(1);

  if (sps_subpic_info_present_flag) {

    // TODO

    return {heif_error_Unsupported_feature,

            heif_suberror_Unsupported_data_version,

            "VVC SPS with subpicture info is not supported yet"};

  }

  uint32_t bitDepth_minus8;

  if (!reader.get_uvlc(&bitDepth_minus8)) {

    return invalidUVLC;

  }

  if (bitDepth_minus8 > 0xFF - 8) {

    return {heif_error_Encoding_error, heif_suberror_Unspecified, "VCC bit depth out of range."};

  }

  config->bit_depth_minus8 = static_cast<uint8_t>(bitDepth_minus8);

  config->constant_frame_rate = 1; // is constant (TODO: where do we get this from)

  return Error::Ok;

}