/*
 * HEIF AVC codec.
 * Copyright (c) 2023 Brad Hards <bradh@frogmouth.net>
 *
 * 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 "avc_boxes.h"
#include <cstdint>
#include <iomanip>
#include <iostream>
#include <vector>
#include "file.h"
#include "context.h"
#include "avc_dec.h"
#include "hevc_boxes.h"
#include <utility>
#include <set>


Error Box_avcC::parse(BitstreamRange& range, const heif_security_limits* limits)
{
  m_configuration.configuration_version = range.read8();
  m_configuration.AVCProfileIndication = range.read8();
  m_configuration.profile_compatibility = range.read8();
  m_configuration.AVCLevelIndication = range.read8();
  uint8_t lengthSizeMinusOneWithReserved = range.read8();
  m_configuration.lengthSize =
      (lengthSizeMinusOneWithReserved & 0b00000011) + 1;

  uint8_t numOfSequenceParameterSets = (range.read8() & 0b00011111);
  for (int i = 0; i < numOfSequenceParameterSets; i++) {
    uint16_t sequenceParameterSetLength = range.read16();
    std::vector<uint8_t> sps(sequenceParameterSetLength);
    range.read(sps.data(), sps.size());
    m_sps.push_back(sps);
  }

  uint8_t numOfPictureParameterSets = range.read8();
  for (int i = 0; i < numOfPictureParameterSets; i++) {
    uint16_t pictureParameterSetLength = range.read16();
    std::vector<uint8_t> pps(pictureParameterSetLength);
    range.read(pps.data(), pps.size());
    m_pps.push_back(pps);
  }

  // See ISO/IEC 14496-15 2017 Section 5.3.3.1.2
  if (range.get_remaining_bytes() > 0) {
    if ((m_configuration.AVCProfileIndication != 66) &&
        (m_configuration.AVCProfileIndication != 77) &&
        (m_configuration.AVCProfileIndication != 88)) {
      m_configuration.chroma_format = (heif_chroma) (range.read8() & 0b00000011);
      m_configuration.bit_depth_luma = 8 + (range.read8() & 0b00000111);
      m_configuration.bit_depth_chroma = 8 + (range.read8() & 0b00000111);
      uint8_t numOfSequenceParameterSetExt = range.read8();
      for (int i = 0; i < numOfSequenceParameterSetExt; i++) {
        uint16_t sequenceParameterSetExtLength = range.read16();
        std::vector<uint8_t> sps_ext(sequenceParameterSetExtLength);
        range.read(sps_ext.data(), sps_ext.size());
        m_sps_ext.push_back(sps_ext);
      }
    }
  }

  return range.get_error();
}

Error Box_avcC::write(StreamWriter& writer) const
{
  size_t box_start = reserve_box_header_space(writer);

  writer.write8(m_configuration.configuration_version);
  writer.write8(m_configuration.AVCProfileIndication);
  writer.write8(m_configuration.profile_compatibility);
  writer.write8(m_configuration.AVCLevelIndication);
  uint8_t lengthSizeMinusOneWithReserved = 0b11111100 | ((m_configuration.lengthSize - 1) & 0b11);
  writer.write8(lengthSizeMinusOneWithReserved);

  if (m_sps.size() > 0b00011111) {
    return {
      heif_error_Encoding_error,
      heif_suberror_Unspecified,
      "Cannot write more than 31 PPS into avcC box."
    };
  }

  uint8_t numSpsWithReserved = 0b11100000 | (m_sps.size() & 0b00011111);
  writer.write8(numSpsWithReserved);
  for (const auto& sps : m_sps) {
    if (sps.size() > 0xFFFF) {
      return {
        heif_error_Encoding_error,
        heif_suberror_Unspecified,
        "Cannot write SPS larger than 65535 bytes into avcC box."
      };
    }
    writer.write16((uint16_t) sps.size());
    writer.write(sps);
  }

  if (m_pps.size() > 0xFF) {
    return {
      heif_error_Encoding_error,
      heif_suberror_Unspecified,
      "Cannot write more than 255 PPS into avcC box."
    };
  }

  writer.write8(m_pps.size() & 0xFF);
  for (const auto& pps : m_pps) {
    if (pps.size() > 0xFFFF) {
      return {
        heif_error_Encoding_error,
        heif_suberror_Unspecified,
        "Cannot write PPS larger than 65535 bytes into avcC box."
      };
    }
    writer.write16((uint16_t) pps.size());
    writer.write(pps);
  }

  if ((m_configuration.AVCProfileIndication != 66) &&
      (m_configuration.AVCProfileIndication != 77) &&
      (m_configuration.AVCProfileIndication != 88)) {
    writer.write8(m_configuration.chroma_format);
    writer.write8(m_configuration.bit_depth_luma - 8);
    writer.write8(m_configuration.bit_depth_chroma - 8);

    if (m_sps_ext.size() > 0xFF) {
      return {
        heif_error_Encoding_error,
        heif_suberror_Unspecified,
        "Cannot write more than 255 SPS-Ext into avcC box."
      };
    }

    writer.write8(m_sps_ext.size() & 0xFF);
    for (const auto& spsext : m_sps_ext) {
      if (spsext.size() > 0xFFFF) {
        return {
          heif_error_Encoding_error,
          heif_suberror_Unspecified,
          "Cannot write SPS-Ext larger than 65535 bytes into avcC box."
        };
      }
      writer.write16((uint16_t) spsext.size());
      writer.write(spsext);
    }
  }

  prepend_header(writer, box_start);

  return Error::Ok;
}

std::string Box_avcC::dump(Indent& indent) const
{
  std::ostringstream sstr;
  sstr << Box::dump(indent);
  sstr << indent << "configuration_version: " << ((int) m_configuration.configuration_version) << "\n"
      << indent << "AVCProfileIndication: " << ((int) m_configuration.AVCProfileIndication) << " (" << profileIndicationAsText() << ")\n"
      << indent << "profile_compatibility: " << ((int) m_configuration.profile_compatibility) << "\n"
      << indent << "AVCLevelIndication: " << ((int) m_configuration.AVCLevelIndication) << "\n"
      << indent << "Chroma format: ";

  switch (m_configuration.chroma_format) {
    case heif_chroma_monochrome:
      sstr << "4:0:0\n";
      break;
    case heif_chroma_420:
      sstr << "4:2:0\n";
      break;
    case heif_chroma_422:
      sstr << "4:2:2\n";
      break;
    case heif_chroma_444:
      sstr << "4:4:4\n";
      break;
    default:
      sstr << "unsupported\n";
      break;
  }

  sstr << indent << "Bit depth luma: " << ((int) m_configuration.bit_depth_luma) << "\n"
      << indent << "Bit depth chroma: " << ((int) m_configuration.bit_depth_chroma) << "\n";

  for (const auto& sps : m_sps) {
    sstr << indent << "SPS: ";
    for (uint8_t b : sps) {
      sstr << std::setfill('0') << std::setw(2) << std::hex << ((int) b) << " ";
    }
    sstr << "\n";
    sstr << std::dec;
  }

  for (const auto& spsext : m_sps_ext) {
    sstr << indent << "SPS-EXT: ";
    for (uint8_t b : spsext) {
      sstr << std::setfill('0') << std::setw(2) << std::hex << ((int) b) << " ";
    }
    sstr << "\n";
    sstr << std::dec;
  }

  for (const auto& pps : m_pps) {
    sstr << indent << "PPS: ";
    for (uint8_t b : pps) {
      sstr << std::setfill('0') << std::setw(2) << std::hex << ((int) b) << " ";
    }
    sstr << "\n";
    sstr << std::dec;
  }

  return sstr.str();
}

std::string Box_avcC::profileIndicationAsText() const
{
  // See ISO/IEC 14496-10:2022 Annex A
  switch (m_configuration.AVCProfileIndication) {
    case 44:
      return "CALVC 4:4:4";
    case 66:
      return "Constrained Baseline";
    case 77:
      return "Main";
    case 88:
      return "Extended";
    case 100:
      return "High variant";
    case 110:
      return "High 10";
    case 122:
      return "High 4:2:2";
    case 244:
      return "High 4:4:4";
    default:
      return "Unknown";
  }
}


void Box_avcC::get_header_nals(std::vector<uint8_t>& data) const
{
  for (const auto& sps : m_sps) {
    data.push_back((sps.size() >> 24) & 0xFF);
    data.push_back((sps.size() >> 16) & 0xFF);
    data.push_back((sps.size() >> 8) & 0xFF);
    data.push_back((sps.size() >> 0) & 0xFF);

    data.insert(data.end(), sps.begin(), sps.end());
  }

  for (const auto& spsext : m_sps_ext) {
    data.push_back((spsext.size() >> 24) & 0xFF);
    data.push_back((spsext.size() >> 16) & 0xFF);
    data.push_back((spsext.size() >> 8) & 0xFF);
    data.push_back((spsext.size() >> 0) & 0xFF);

    data.insert(data.end(), spsext.begin(), spsext.end());
  }

  for (const auto& pps : m_pps) {
    data.push_back((pps.size() >> 24) & 0xFF);
    data.push_back((pps.size() >> 16) & 0xFF);
    data.push_back((pps.size() >> 8) & 0xFF);
    data.push_back((pps.size() >> 0) & 0xFF);

    data.insert(data.end(), pps.begin(), pps.end());
  }
}


void Box_avcC::append_sps_nal(const uint8_t* data, size_t size)
{
  std::vector<uint8_t> vec(data, data + size);
  m_sps.emplace_back(std::move(vec));
}

void Box_avcC::append_sps_ext_nal(const uint8_t* data, size_t size)
{
  std::vector<uint8_t> vec(data, data + size);
  m_sps_ext.emplace_back(std::move(vec));
}

void Box_avcC::append_pps_nal(const uint8_t* data, size_t size)
{
  std::vector<uint8_t> vec(data, data + size);
  m_pps.emplace_back(std::move(vec));
}

static bool skip_scaling_list(BitReader& reader, int sizeOfScalingList)
{
  int lastScale = 8;
  int nextScale = 8;

  // TODO: it seems that this can be simplified by exiting the loop
  //       as soon as nextScale==0.

#if 0
  // original version
  for (int j = 0; j < sizeOfScalingList; j++) {
    if (nextScale != 0) {
      int32_t delta_scale;
      if (!reader.get_svlc(&delta_scale)) {
        return false;
      }
      nextScale = (lastScale + delta_scale + 256) % 256;
    }

    lastScale = (nextScale == 0) ? lastScale : nextScale;
  }
#else
  // fast version
  for (int j = 0; j < sizeOfScalingList; j++) {
    int32_t delta_scale;
    if (!reader.get_svlc(&delta_scale)) {
      return false;
    }
    nextScale = (lastScale + delta_scale + 256) % 256;

    if (nextScale == 0) {
      break;
    }

    lastScale = nextScale;
  }
#endif

  return true;
}


Error parse_sps_for_avcC_configuration(const uint8_t* sps, size_t size,
                                       Box_avcC::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, size);

  // skip NAL header
  reader.skip_bits(8);

  config->configuration_version = 1;
  config->AVCProfileIndication = reader.get_bits8(8);
  config->profile_compatibility = reader.get_bits8(8);
  config->AVCLevelIndication = reader.get_bits8(8);
  config->lengthSize = 4;

  Error invalidUVLC{
    heif_error_Invalid_input,
    heif_suberror_Unspecified,
    "Invalid variable length code in AVC SPS header"
  };

  uint32_t value;
  if (!reader.get_uvlc(&value)) { return invalidUVLC; } // SPS ID

  if (std::set<int>{100, 110, 122, 244, 44, 83, 86, 118, 128, 138, 139, 134, 135}.contains(config->AVCProfileIndication)) {
    if (!reader.get_uvlc(&value)) {
      return invalidUVLC;
    }

    if (value > heif_chroma_444) {
      return {
        heif_error_Invalid_input,
        heif_suberror_Unspecified,
        "Invalid chroma format in AVC SPS header"
      };
    }

    config->chroma_format = static_cast<heif_chroma>(value);
    if (config->chroma_format == heif_chroma_444) {
      reader.skip_bits(1);
    }

    if (!reader.get_uvlc(&value)) {
      return invalidUVLC;
    }
    config->bit_depth_luma = static_cast<uint8_t>(8 + value);

    if (!reader.get_uvlc(&value)) {
      return invalidUVLC;
    }
    config->bit_depth_chroma = static_cast<uint8_t>(8 + value);

    reader.skip_bits(1);
    int seq_scaling_matrix_present_flag = reader.get_bits(1);
    if (seq_scaling_matrix_present_flag) {
      for (int i = 0; i < ((config->chroma_format != heif_chroma_444) ? 8 : 12); i++) {
        int scaling_list_present_flag = reader.get_bits(1);
        if (scaling_list_present_flag) {
          if (!skip_scaling_list(reader, i < 6 ? 16 : 64)) {
            return invalidUVLC;
          }
        }
      }
    }
  }
  else {
    config->chroma_format = heif_chroma_420;
    config->bit_depth_luma = 8;
    config->bit_depth_chroma = 8;
  }

  if (!reader.get_uvlc(&value)) { return invalidUVLC; } // log2_max_frame_num_minus4
  uint32_t pic_order_cnt_type;
  if (!reader.get_uvlc(&pic_order_cnt_type)) { return invalidUVLC; }
  if (pic_order_cnt_type == 0) {
    if (!reader.get_uvlc(&value)) { return invalidUVLC; }
  }
  else if (pic_order_cnt_type == 1) {
    reader.get_bits(1);
    int32_t svalue;
    if (!reader.get_svlc(&svalue) ||
        !reader.get_svlc(&svalue)) { return invalidUVLC; }
    uint32_t num_ref_franes_in_pic_order_cnt_cycle;
    if (!reader.get_uvlc(&num_ref_franes_in_pic_order_cnt_cycle)) { return invalidUVLC; }
    for (uint32_t i = 0; i < num_ref_franes_in_pic_order_cnt_cycle; i++) {
      if (!reader.get_uvlc(&value)) { return invalidUVLC; }
    }
  }

  if (!reader.get_uvlc(&value)) { return invalidUVLC; } // num_ref_frames
  reader.skip_bits(1);

  uint32_t pic_width_in_mbs_minus1;
  uint32_t pic_height_in_map_units_minus1;
  if (!reader.get_uvlc(&pic_width_in_mbs_minus1) ||
      !reader.get_uvlc(&pic_height_in_map_units_minus1)) {
    return invalidUVLC;
  }

  uint32_t frame_mbs_only_flag = reader.get_bits(1);
  if (!frame_mbs_only_flag) {
    reader.skip_bits(1); // mb_adaptive_frame_field_flag
  }
  reader.skip_bits(1); // direct_8x8_inference_flag

  // for interlaced content, a map unit covers two macroblock rows

  uint64_t width64 = (static_cast<uint64_t>(pic_width_in_mbs_minus1) + 1) * 16;
  uint64_t height64 = (static_cast<uint64_t>(pic_height_in_map_units_minus1) + 1) * 16 * (2 - frame_mbs_only_flag);

  if (width64 > UINT32_MAX || height64 > UINT32_MAX) {
    return {heif_error_Invalid_input,
            heif_suberror_Invalid_image_size,
            "AVC SPS image size too large"};
  }

  *width = static_cast<uint32_t>(width64);
  *height = static_cast<uint32_t>(height64);

  if (coded_size) {
    coded_size->width = *width;
    coded_size->height = *height;
  }

  uint32_t frame_cropping_flag = reader.get_bits(1);
  if (frame_cropping_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;
    }

    // The crop offsets are not in luma samples, but in crop units, which depend on the
    // chroma format and the frame_mbs_only_flag (CropUnitX/CropUnitY, H.264 7.4.2.1.1).

    uint32_t crop_unit_x = (config->chroma_format == heif_chroma_420 ||
                            config->chroma_format == heif_chroma_422) ? 2 : 1;
    uint32_t crop_unit_y = ((config->chroma_format == heif_chroma_420) ? 2 : 1) * (2 - frame_mbs_only_flag);

    uint64_t crop_horizontal = (static_cast<uint64_t>(left) + right) * crop_unit_x;
    uint64_t crop_vertical = (static_cast<uint64_t>(top) + bottom) * crop_unit_y;
    if (crop_horizontal >= *width || crop_vertical >= *height) {
      return {heif_error_Invalid_input,
              heif_suberror_Invalid_image_size,
              "AVC SPS cropping exceeds image size"};
    }

    *width -= static_cast<uint32_t>(crop_horizontal);
    *height -= static_cast<uint32_t>(crop_vertical);
  }

  return {};
}
