/*

 * HEIF codec.

 * Copyright (c) 2017 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 "hevc_boxes.h"

#include "bitstream.h"

#include "error.h"

#include "file.h"

#include "hevc_dec.h"

#include <cassert>

#include <cmath>

#include <cstring>

#include <iomanip>

#include <string>

#include <utility>

#include <algorithm>

#include "api_structs.h"

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

{

  uint8_t byte;

  configuration_version = range.read8();

  byte = range.read8();

  general_profile_space = (byte >> 6) & 3;

  general_tier_flag = (byte >> 5) & 1;

  general_profile_idc = (byte & 0x1F);

  general_profile_compatibility_flags = range.read32();

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

    byte = range.read8();

    for (int b = 0; b < 8; b++) {

      general_constraint_indicator_flags[i * 8 + b] = (byte >> (7 - b)) & 1;

    }

  }

  general_level_idc = range.read8();

  min_spatial_segmentation_idc = range.read16() & 0x0FFF;

  parallelism_type = range.read8() & 0x03;

  chroma_format = range.read8() & 0x03;

  bit_depth_luma = static_cast<uint8_t>((range.read8() & 0x07) + 8);

  bit_depth_chroma = static_cast<uint8_t>((range.read8() & 0x07) + 8);

  avg_frame_rate = range.read16();

  byte = range.read8();

  constant_frame_rate = (byte >> 6) & 0x03;

  num_temporal_layers = (byte >> 3) & 0x07;

  temporal_id_nested = (byte >> 2) & 1;

  m_length_size = static_cast<uint8_t>((byte & 0x03) + 1);

  int nArrays = range.read8();

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

    byte = range.read8();

    NalArray array;

    array.m_array_completeness = (byte >> 6) & 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));

  }

  range.skip_to_end_of_box();

  return range.get_error();

}

Error HEVCDecoderConfigurationRecord::write(StreamWriter& writer) const

{

  writer.write8(configuration_version);

  writer.write8((uint8_t) (((general_profile_space & 3) << 6) |

                           ((general_tier_flag & 1) << 5) |

                           (general_profile_idc & 0x1F)));

  writer.write32(general_profile_compatibility_flags);

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

    uint8_t byte = 0;

    for (int b = 0; b < 8; b++) {

      if (general_constraint_indicator_flags[i * 8 + b]) {

        byte |= 1;

      }

      byte = (uint8_t) (byte << 1);

    }

    writer.write8(byte);

  }

  writer.write8(general_level_idc);

  writer.write16((min_spatial_segmentation_idc & 0x0FFF) | 0xF000);

  writer.write8(parallelism_type | 0xFC);

  writer.write8(chroma_format | 0xFC);

  writer.write8((uint8_t) ((bit_depth_luma - 8) | 0xF8));

  writer.write8((uint8_t) ((bit_depth_chroma - 8) | 0xF8));

  writer.write16(avg_frame_rate);

  writer.write8((uint8_t) (((constant_frame_rate & 0x03) << 6) |

                           ((num_temporal_layers & 0x07) << 3) |

                           ((temporal_id_nested & 1) << 2) |

                           ((m_length_size - 1) & 0x03)));

  size_t nArrays = m_nal_array.size();

  // There cannot be an overflow because the nal-type is less than 8-bit.

  assert(nArrays <= 0xFF);

  writer.write8(static_cast<uint8_t>(nArrays));

  for (const HEVCDecoderConfigurationRecord::NalArray& array : m_nal_array) {

    writer.write8((uint8_t) (((array.m_array_completeness & 1) << 6) |

                             (array.m_NAL_unit_type & 0x3F)));

    size_t nUnits = array.m_nal_units.size();

    if (nUnits > 0xFFFF) {

      return Error{heif_error_Invalid_input, heif_suberror_Unspecified, "Too many NAL units in hvcC"};

    }

    writer.write16(static_cast<uint16_t>(nUnits));

    for (const std::vector<uint8_t>& nal_unit : array.m_nal_units) {

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

        return Error{heif_error_Invalid_input, heif_suberror_Unspecified, "hvcC NAL unit exceeds maximum size (64kB)"};

      }

      writer.write16(static_cast<uint16_t>(nal_unit.size()));

      writer.write(nal_unit);

    }

  }

  return Error::Ok;

}

bool HEVCDecoderConfigurationRecord::get_general_profile_compatibility_flag(int idx) const

{

  return general_profile_compatibility_flags & (UINT32_C(0x80000000) >> idx);

}

bool HEVCDecoderConfigurationRecord::is_profile_compatible(Profile profile) const

{

  return (general_profile_idc == profile ||

          get_general_profile_compatibility_flag(profile));

}

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

{

  return m_configuration.parse(range, limits);

}

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

{

  std::ostringstream sstr;

  sstr << Box::dump(indent);

  const auto& c = m_configuration; // abbreviation

  sstr << indent << "configuration_version: " << ((int) c.configuration_version) << "\n"

       << indent << "general_profile_space: " << ((int) c.general_profile_space) << "\n"

       << indent << "general_tier_flag: " << c.general_tier_flag << "\n"

       << indent << "general_profile_idc: " << ((int) c.general_profile_idc) << "\n";

  sstr << indent << "general_profile_compatibility_flags: ";

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

    sstr << ((c.general_profile_compatibility_flags >> (31 - i)) & 1);

    if ((i % 8) == 7) sstr << ' ';

    else if ((i % 4) == 3) sstr << '.';

  }

  sstr << "\n";

  sstr << indent << "general_constraint_indicator_flags: ";

  int cnt = 0;

  for (int i = 0; i < HEVCDecoderConfigurationRecord::NUM_CONSTRAINT_INDICATOR_FLAGS; i++) {

    bool b = c.general_constraint_indicator_flags[i];

    sstr << (b ? 1 : 0);

    cnt++;

    if ((cnt % 8) == 0)

      sstr << ' ';

  }

  sstr << "\n";

  sstr << indent << "general_level_idc: " << ((int) c.general_level_idc) << "\n"

       << indent << "min_spatial_segmentation_idc: " << c.min_spatial_segmentation_idc << "\n"

       << indent << "parallelism_type: " << ((int) c.parallelism_type) << "\n"

       << indent << "chroma_format: ";

  switch (c.chroma_format) {

    case 0:

      sstr << "monochrome";

      break;

    case 1:

      sstr << "4:2:0";

      break;

    case 2:

      sstr << "4:2:2";

      break;

    case 3:

      sstr << "4:4:4";

      break;

    default:

      sstr << ((int) c.chroma_format);

      break;

  }

  sstr << "\n"

       << indent << "bit_depth_luma: " << ((int) c.bit_depth_luma) << "\n"

       << indent << "bit_depth_chroma: " << ((int) c.bit_depth_chroma) << "\n"

       << indent << "avg_frame_rate: " << c.avg_frame_rate << "\n"

       << indent << "constant_frame_rate: " << ((int) c.constant_frame_rate) << "\n"

       << indent << "num_temporal_layers: " << ((int) c.num_temporal_layers) << "\n"

       << indent << "temporal_id_nested: " << ((int) c.temporal_id_nested) << "\n"

       << indent << "length_size: " << ((int) c.m_length_size) << "\n";

  for (const auto& array : c.m_nal_array) {

    sstr << indent << "<array>\n";

    indent++;

    sstr << indent << "array_completeness: " << ((int) array.m_array_completeness) << "\n"

         << indent << "NAL_unit_type: " << ((int) array.m_NAL_unit_type) << "\n";

    for (const auto& unit : array.m_nal_units) {

      //sstr << "  unit with " << unit.size() << " bytes of data\n";

      sstr << indent;

      for (uint8_t b : unit) {

        sstr << std::setfill('0') << std::setw(2) << std::hex << ((int) b) << " ";

      }

      sstr << "\n";

      sstr << std::dec;

    }

    indent--;

  }

  return sstr.str();

}

bool Box_hvcC::get_header_nals(std::vector<uint8_t>* dest) const

{

  // Concatenate all header NALs, each prefixed by a 4-byte size.

  for (const auto& array : m_configuration.m_nal_array) {

    for (const auto& unit : array.m_nal_units) {

      // Write 4-byte NALs size

      dest->push_back((unit.size() >> 24) & 0xFF);

      dest->push_back((unit.size() >> 16) & 0xFF);

      dest->push_back((unit.size() >> 8) & 0xFF);

      dest->push_back((unit.size() >> 0) & 0xFF);

      // Copy NAL data

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

    }

  }

  return true;

}

void Box_hvcC::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);

}

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

{

  for (auto& nal_array : m_configuration.m_nal_array) {

    if (nal_array.m_NAL_unit_type == uint8_t(nal[0] >> 1)) {

      // kvazaar may send the same headers multiple times. Filter out the identical copies.

      for (auto& nal_unit : nal_array.m_nal_units) {

        // Note: sometimes kvazaar even sends the same packet twice, but with an extra zero byte.

        //       We detect this by comparing only the common length. This is correct since each NAL

        //       packet must be complete and thus, if a packet is longer than another complete packet,

        //       its extra data must be superfluous.

        //

        // Example:

        //| | | <array>

        //| | | | array_completeness: 1

        //| | | | NAL_unit_type: 34

        //| | | | 44 01 c1 71 82 99 20 00

        //| | | | 44 01 c1 71 82 99 20

        // Check whether packets have similar content.

        const size_t common_length = std::min(nal_unit.size(), nal.size());

        bool similar = true;

        for (size_t i = 0; i < common_length; i++) {

          if (nal_unit[i] != nal[i]) {

            similar = false;

            break;

          }

        }

        if (similar) {

          // If they are similar, keep the smaller one.

          if (nal_unit.size() > nal.size()) {

            nal_unit = std::move(nal);

          }

          // Exit. Do not add a copy of the packet.

          return;

        }

      }

      nal_array.m_nal_units.push_back(std::move(nal));

      return;

    }

  }

  // This is a new NAL type. Add a new NAL array.

  HEVCDecoderConfigurationRecord::NalArray array;

  array.m_array_completeness = 1;

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

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

  m_configuration.m_nal_array.push_back(array);

}

Error Box_hvcC::write(StreamWriter& writer) const

{

  size_t box_start = reserve_box_header_space(writer);

  const auto& c = m_configuration; // abbreviation

  Error err = c.write(writer);

  if (err) {

    return err;

  }

  prepend_header(writer, box_start);

  return Error::Ok;

}

static double read_depth_rep_info_element(BitReader& reader)

{

  uint8_t sign_flag = reader.get_bits8(1);

  int exponent = reader.get_bits(7);

  auto mantissa_len = static_cast<uint8_t>(reader.get_bits8(5) + 1);

  if (mantissa_len < 1 || mantissa_len > 32) {

    // TODO err

  }

  if (exponent == 127) {

    // TODO value unspecified

  }

  uint32_t mantissa = reader.get_bits32(mantissa_len);

  double value;

  //printf("sign:%d exponent:%d mantissa_len:%d mantissa:%d\n",sign_flag,exponent,mantissa_len,mantissa);

  // TODO: this seems to be wrong. 'exponent' is never negative. How to read it correctly?

  if (exponent > 0) {

    value = pow(2.0, exponent - 31) * (1.0 + mantissa / pow(2.0, mantissa_len));

  }

  else {

    value = pow(2.0, -(30 + mantissa_len)) * mantissa;

  }

  if (sign_flag) {

    value = -value;

  }

  return value;

}

static Result<std::shared_ptr<SEIMessage>> read_depth_representation_info(BitReader& reader)

{

  auto msg = std::make_shared<SEIMessage_depth_representation_info>();

  // default values

  msg->version = 1;

  msg->disparity_reference_view = 0;

  msg->depth_nonlinear_representation_model_size = 0;

  msg->depth_nonlinear_representation_model = nullptr;

  // read header

  msg->has_z_near = (uint8_t) reader.get_bits(1);

  msg->has_z_far = (uint8_t) reader.get_bits(1);

  msg->has_d_min = (uint8_t) reader.get_bits(1);

  msg->has_d_max = (uint8_t) reader.get_bits(1);

  uint32_t rep_type;

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

    return Error{heif_error_Invalid_input, heif_suberror_Invalid_parameter_value, "invalid depth representation type in input"};

  }

  if (rep_type > 3) {

    return Error{heif_error_Invalid_input, heif_suberror_Invalid_parameter_value, "input depth representation type out of range"};

  }

  msg->depth_representation_type = (enum heif_depth_representation_type) rep_type;

  //printf("flags: %d %d %d %d\n",msg->has_z_near,msg->has_z_far,msg->has_d_min,msg->has_d_max);

  //printf("type: %d\n",rep_type);

  if (msg->has_d_min || msg->has_d_max) {

    uint32_t ref_view;

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

      return Error{heif_error_Invalid_input, heif_suberror_Invalid_parameter_value, "invalid disparity_reference_view in input"};

    }

    msg->disparity_reference_view = ref_view;

    //printf("ref_view: %d\n",msg->disparity_reference_view);

  }

  if (msg->has_z_near) msg->z_near = read_depth_rep_info_element(reader);

  if (msg->has_z_far) msg->z_far = read_depth_rep_info_element(reader);

  if (msg->has_d_min) msg->d_min = read_depth_rep_info_element(reader);

  if (msg->has_d_max) msg->d_max = read_depth_rep_info_element(reader);

  /*

  printf("z_near: %f\n",msg->z_near);

  printf("z_far: %f\n",msg->z_far);

  printf("dmin: %f\n",msg->d_min);

  printf("dmax: %f\n",msg->d_max);

  */

  if (msg->depth_representation_type == heif_depth_representation_type_nonuniform_disparity) {

    // TODO: load non-uniform response curve

  }

  return {msg};

}

// aux subtypes: 00 00 00 11 / 00 00 00 0d / 4e 01 / b1 09 / 35 1e 78 c8 01 03 c5 d0 20

Error decode_hevc_aux_sei_messages(const std::vector<uint8_t>& data,

                                   std::vector<std::shared_ptr<SEIMessage>>& msgs)

{

  // TODO: we probably do not need a full BitReader just for the array size.

  // Read this and the NAL size directly on the array data.

  BitReader reader(data.data(), (int) data.size());

  if (reader.get_bits_remaining() < 32) {

    return {heif_error_Invalid_input,

            heif_suberror_End_of_data,

            "HEVC SEI NAL too short"};

  }

  uint32_t len = reader.get_bits32(32);

  if (len > data.size() - 4) {

    // ERROR: read past end of data

  }

  while (reader.get_current_byte_index() < (int) len) {

    int currPos = reader.get_current_byte_index();

    BitReader sei_reader(data.data() + currPos, (int) data.size() - currPos);

    if (sei_reader.get_bits_remaining() < 32+8) {

      return {heif_error_Invalid_input,

              heif_suberror_End_of_data,

              "HEVC SEI NAL too short"};

    }

    uint32_t nal_size = sei_reader.get_bits32(32);

    (void) nal_size;

    auto nal_type = static_cast<uint8_t>(sei_reader.get_bits8(8) >> 1);

    sei_reader.skip_bits(8);

    // SEI

    if (nal_type == 39 ||

        nal_type == 40) {

      if (sei_reader.get_bits_remaining() < 16) {

        return {heif_error_Invalid_input,

                heif_suberror_End_of_data,

                "HEVC SEI NAL too short"};

      }

      // TODO: loading of multi-byte sei headers

      uint8_t payload_id = sei_reader.get_bits8(8);

      uint8_t payload_size = sei_reader.get_bits8(8);

      (void) payload_size;

      if (payload_id == 177) {

        // depth_representation_info

        Result<std::shared_ptr<SEIMessage>> seiResult = read_depth_representation_info(sei_reader);

        if (!seiResult) {

          return seiResult.error();

        }

        msgs.push_back(*seiResult);

      }

    }

    break; // TODO: read next SEI

  }

  return Error::Ok;

}

// Used for AVC, HEVC, and VVC.

std::vector<uint8_t> remove_start_code_emulation(const uint8_t* sps, size_t size)

{

  std::vector<uint8_t> out_data;

  for (size_t i = 0; i < size; i++) {

    if (i + 2 < size &&

        sps[i] == 0 &&

        sps[i + 1] == 0 &&

        sps[i + 2] == 3) {

      out_data.push_back(0);

      out_data.push_back(0);

      i += 2;

    }

    else {

      out_data.push_back(sps[i]);

    }

  }

  return out_data;

}

Error parse_sps_for_hvcC_configuration(const uint8_t* sps, size_t size,

                                       HEVCDecoderConfigurationRecord* 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 VPS ID

  reader.skip_bits(4);

  uint8_t nMaxSubLayersMinus1 = reader.get_bits8(3);

  config->temporal_id_nested = reader.get_bits8(1);

  // --- profile_tier_level ---

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

  config->general_tier_flag = reader.get_bits8(1);

  config->general_profile_idc = reader.get_bits8(5);

  config->general_profile_compatibility_flags = reader.get_bits32(32);

  reader.skip_bits(16); // skip reserved bits

  reader.skip_bits(16); // skip reserved bits

  reader.skip_bits(16); // skip reserved bits

  config->general_level_idc = reader.get_bits8(8);

  std::vector<bool> layer_profile_present(nMaxSubLayersMinus1);

  std::vector<bool> layer_level_present(nMaxSubLayersMinus1);

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

    layer_profile_present[i] = reader.get_bits(1);

    layer_level_present[i] = reader.get_bits(1);

  }

  if (nMaxSubLayersMinus1 > 0) {

    for (int i = nMaxSubLayersMinus1; i < 8; i++) {

      reader.skip_bits(2);

    }

  }

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

    if (layer_profile_present[i]) {

      reader.skip_bits(2 + 1 + 5);

      reader.skip_bits(32);

      reader.skip_bits(16);

    }

    if (layer_level_present[i]) {

      reader.skip_bits(8);

    }

  }

  // --- SPS continued ---

  Error invalidUVLC{

    heif_error_Invalid_input,

    heif_suberror_Invalid_parameter_value,

    "Invalid variable length code in HEVC SPS header"

  };

  uint32_t dummy, value;

  if (!reader.get_uvlc(&dummy) || // skip seq_parameter_seq_id

      !reader.get_uvlc(&value)) {

    return invalidUVLC;

  }

  config->chroma_format = (uint8_t) value;

  if (config->chroma_format == 3) {

    reader.skip_bits(1);

  }

  if (!reader.get_uvlc(width) ||

      !reader.get_uvlc(height)) {

    return invalidUVLC;

  }

  if (coded_size) {

    coded_size->width = *width;

    coded_size->height = *height;

  }

  bool conformance_window = reader.get_bits(1);

  if (conformance_window) {

    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;

    }

    //printf("conformance borders: %u %u %u %u\n",left,right,top,bottom);

    uint32_t subH = 1, subV = 1;

    if (config->chroma_format == 1) {

      subV = 2;

      subH = 2;

    }

    if (config->chroma_format == 2) { subH = 2; }

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

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

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

      return Error{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;

  }

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

    return invalidUVLC;

  }

  if (value > 8) {

    return Error{heif_error_Invalid_input,

                 heif_suberror_Invalid_parameter_value,

                 "SPS bit_depth_luma_minus8 out of range"};

  }

  config->bit_depth_luma = (uint8_t) (value + 8);

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

    return invalidUVLC;

  }

  if (value > 8) {

    return Error{heif_error_Invalid_input,

                 heif_suberror_Invalid_parameter_value,

                 "SPS bit_depth_chroma_minus8 out of range"};

  }

  config->bit_depth_chroma = (uint8_t) (value + 8);

  // --- init static configuration fields ---

  config->configuration_version = 1;

  config->min_spatial_segmentation_idc = 0; // TODO: get this value from the VUI, 0 should be safe

  config->parallelism_type = 0; // TODO, 0 should be safe

  config->avg_frame_rate = 0; // makes no sense for HEIF (TODO)

  config->constant_frame_rate = 0; // makes no sense for HEIF (TODO)

  config->num_temporal_layers = 1; // makes no sense for HEIF

  return Error::Ok;

}