// Copyright 2020 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
////////////////////////////////////////////////////////////////////////////////

#include "src/piex_cr3.h"

#include <array>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <unordered_set>

#include "src/binary_parse/range_checked_byte_ptr.h"
#include "src/piex_types.h"
#include "src/tiff_directory/tiff_directory.h"
#include "src/tiff_parser.h"

namespace piex {
namespace {

constexpr size_t kUuidSize = 16;
using Uuid = std::array<std::uint8_t, kUuidSize>;
// Uuid of uuid box under the moov box.
constexpr Uuid kUuidMoov = {0x85, 0xc0, 0xb6, 0x87, 0x82, 0x0f, 0x11, 0xe0,
                            0x81, 0x11, 0xf4, 0xce, 0x46, 0x2b, 0x6a, 0x48};

// Uuid of uuid box containing PRVW box.
constexpr Uuid kUuidPrvw = {0xea, 0xf4, 0x2b, 0x5e, 0x1c, 0x98, 0x4b, 0x88,
                            0xb9, 0xfb, 0xb7, 0xdc, 0x40, 0x6e, 0x4d, 0x16};

constexpr size_t kTagSize = 4;
using BoxTag = std::array<char, kTagSize>;

constexpr BoxTag NewTag(const char s[kTagSize + 1]) {
  return BoxTag{s[0], s[1], s[2], s[3]};
}

constexpr BoxTag kUuidTag = NewTag("uuid");
constexpr BoxTag kPrvwTag = NewTag("PRVW");
constexpr BoxTag kThmbTag = NewTag("THMB");
constexpr BoxTag kCmt1Tag = NewTag("CMT1");
constexpr BoxTag kCmt2Tag = NewTag("CMT2");
constexpr BoxTag kStblTag = NewTag("stbl");
constexpr BoxTag kStsdTag = NewTag("stsd");
constexpr BoxTag kCrawTag = NewTag("CRAW");
constexpr BoxTag kStszTag = NewTag("stsz");
constexpr BoxTag kCo64Tag = NewTag("co64");
constexpr BoxTag kMdatTag = NewTag("mdat");

// Convenience class for a box.
class Box {
 public:
  Box()
      : is_valid_(false),
        tag_(BoxTag()),
        offset_(0),
        header_offset_(0),
        next_box_offset_(0) {}
  Box(const BoxTag& tag, size_t offset, size_t header_length, size_t length)
      : is_valid_(true),
        tag_(tag),
        offset_(offset),
        header_offset_(offset + header_length),
        next_box_offset_(offset + length) {}

  bool IsValid() const { return is_valid_ && next_box_offset_ > offset_; }
  const BoxTag& tag() const { return tag_; }

  // Returns offset from start of file.
  size_t offset() const { return offset_; }
  // Returns offset from start of file, including box's header.
  size_t header_offset() const { return header_offset_; }
  // Returns offset from start of file of the next box, accounting for size of
  // this box.
  size_t next_box_offset() const { return next_box_offset_; }

 private:
  bool is_valid_;
  BoxTag tag_;
  size_t offset_;
  size_t header_offset_;
  size_t next_box_offset_;
};

struct ProcessData {
  PreviewImageData* preview_image_data = nullptr;
  Image mdat_image;
  Image prvw_image;
};

// Wraps Get16u w/ assumption that CR3 is always big endian, based on
// ISO/IEC 14496-12 specification that all box fields are big endian.
bool Get16u(StreamInterface* stream, size_t offset, std::uint16_t* value) {
  return Get16u(stream, offset, tiff_directory::kBigEndian, value);
}

// Wraps Get32u w/ assumption that CR3 is always big endian, based on
// ISO/IEC 14496-12 specification that all box fields are big endian.
bool Get32u(StreamInterface* stream, size_t offset, std::uint32_t* value) {
  return Get32u(stream, offset, tiff_directory::kBigEndian, value);
}

// Always big endian, based on ISO/IEC 14496-12 specification that all box
// fields are big endian.
bool Get64u(StreamInterface* stream, size_t offset, std::uint64_t* value) {
  std::uint8_t data[8];
  if (stream->GetData(offset, 8, data) == kOk) {
    *value = (data[0] * 0x1000000u) | (data[1] * 0x10000u) |
             (data[2] * 0x100u) | data[3];
    *value <<= 32;
    *value = (data[4] * 0x1000000u) | (data[5] * 0x10000u) |
             (data[6] * 0x100u) | data[7];
    return true;
  } else {
    return false;
  }
}

// Jpeg box offsets based on the box tag. The expected layout is as follows:
//        Byte Offset Type     Meaning
//                  0 [long]   size of box
//                  4 [char[]] box tag
//       offset.width [short]  width of jpeg
//      offset.height [short]  height of jpeg
//   offset.jpeg_size [long]   number of bytes in jpeg
//   offset.jpeg_data [byte[]] start of jpeg data
struct JpegBoxOffset {
  size_t width = 0;
  size_t height = 0;
  size_t jpeg_size = 0;
  size_t jpeg_data = 0;
};

// Processes box w/ JPEG data. Box must be PRVW and THMB boxes.
bool ProcessJpegBox(StreamInterface* stream, const Box& box, Image* image) {
  static constexpr JpegBoxOffset kPrvwJpegOffsets{14, 16, 20, 24};
  static constexpr JpegBoxOffset kThmbJpegOffsets{12, 14, 16, 24};
  if (box.tag() != kPrvwTag && box.tag() != kThmbTag) {
    return false;
  }
  const JpegBoxOffset& offsets =
      box.tag() == kPrvwTag ? kPrvwJpegOffsets : kThmbJpegOffsets;
  uint16_t width, height;
  uint32_t jpeg_size;
  if (!Get16u(stream, box.offset() + offsets.width, &width)) {
    return false;
  }
  if (!Get16u(stream, box.offset() + offsets.height, &height)) {
    return false;
  }
  if (!Get32u(stream, box.offset() + offsets.jpeg_size, &jpeg_size)) {
    return false;
  }
  image->format = Image::kJpegCompressed;
  image->width = width;
  image->height = height;
  image->offset = box.offset() + offsets.jpeg_data;
  image->length = jpeg_size;
  return true;
}

// Parses the Exif IFD0 tags at tiff_offset.
bool ParseExifIfd0(StreamInterface* stream, size_t tiff_offset,
                   PreviewImageData* preview_image_data) {
  static const TagSet kIfd0TagSet = {kTiffTagModel, kTiffTagMake,
                                     kTiffTagOrientation, kTiffTagImageWidth,
                                     kTiffTagImageLength};
  TiffContent content;
  TiffParser(stream, tiff_offset).Parse(kIfd0TagSet, 1, &content);
  if (content.tiff_directory.size() != 1) {
    return false;
  }

  content.tiff_directory[0].Get(kTiffTagModel, &preview_image_data->model);
  content.tiff_directory[0].Get(kTiffTagMake, &preview_image_data->maker);
  content.tiff_directory[0].Get(kTiffTagOrientation,
                                &preview_image_data->exif_orientation);
  content.tiff_directory[0].Get(kTiffTagImageWidth,
                                &preview_image_data->full_width);
  content.tiff_directory[0].Get(kTiffTagImageLength,
                                &preview_image_data->full_height);
  return true;
}

// Parses the Exif Exif IFD tags at tiff_offset.
bool ParseExifExifIfd(StreamInterface* stream, size_t tiff_offset,
                      PreviewImageData* preview_image_data) {
  static const TagSet kExifIfdTagSet = {kExifTagDateTimeOriginal,
                                        kExifTagExposureTime, kExifTagFnumber,
                                        kExifTagFocalLength, kExifTagIsoSpeed};
  TiffContent content;
  TiffParser(stream, tiff_offset).Parse(kExifIfdTagSet, 1, &content);
  if (content.tiff_directory.size() != 1) {
    return false;
  }

  content.tiff_directory[0].Get(kExifTagDateTimeOriginal,
                                &preview_image_data->date_time);
  GetRational(kExifTagExposureTime, content.tiff_directory[0], 1,
              &preview_image_data->exposure_time);
  GetRational(kExifTagFnumber, content.tiff_directory[0], 1,
              &preview_image_data->fnumber);
  GetRational(kExifTagFocalLength, content.tiff_directory[0], 1,
              &preview_image_data->focal_length);
  content.tiff_directory[0].Get(kExifTagIsoSpeed, &preview_image_data->iso);
  return true;
}

// Returns the next box or an invalid box.
//
// Based on ISO/IEC 14496-12: boxes start with a header: size and type. The size
// can be compact (32-bits) or extended (64-bit, e.g. mdat box).
// The type can be compact (32 bits) or extended (full UUID, e.g. uuid boxes).
// values are stored after the compact size/type.
//
// Fields in a box are big-endian.
Box GetNextBox(StreamInterface* stream, size_t offset) {
  uint32_t length_32;
  if (!Get32u(stream, offset, &length_32)) {
    return Box();
  }
  BoxTag tag;
  Error status = stream->GetData(offset + sizeof(length_32), kTagSize,
                                 reinterpret_cast<std::uint8_t*>(tag.data()));
  if (status != kOk) {
    return Box();
  }
  size_t length;
  size_t header_offset = sizeof(length_32) + sizeof(tag);
  if (length_32 == 1) {
    // Magic number of 1 implies extended size.
    uint64_t length_64 = 0;
    if (!Get64u(stream, offset + header_offset, &length_64)) {
      return Box();
    }
    length = length_64;
    header_offset += sizeof(length_64);
  } else {
    // Compact size.
    length = length_32;
  }
  return Box(tag, offset, header_offset, length);
}

// Searches for the next box with the given tag.
Box GetNextBoxWithTag(StreamInterface* stream, size_t offset,
                      const BoxTag& expected_tag) {
  while (true) {
    Box box = GetNextBox(stream, offset);
    if (!box.IsValid() || box.tag() == expected_tag) {
      return box;
    }
    offset = box.next_box_offset();
  }
}

// Returns the width, height, and content type from the CRAW box.
bool ProcessCrawBox(StreamInterface* stream, const Box& craw_box,
                    uint16_t* width, uint16_t* height, uint16_t* content_type) {
  constexpr size_t kWidthOffset = 32;
  if (!Get16u(stream, craw_box.offset() + kWidthOffset, width)) {
    return false;
  }

  constexpr size_t kHeightOffset = 34;
  if (!Get16u(stream, craw_box.offset() + kHeightOffset, height)) {
    return false;
  }

  constexpr size_t kTypeOffset = 86;
  if (!Get16u(stream, craw_box.offset() + kTypeOffset, content_type)) {
    return false;
  }
  return true;
}

// stsz box offset:
//        Byte Offset Type     Meaning
//                  0 [long]   size of box
//                  4 [char[]] box tag
//                  8 [long]   version/flags
//                 12 [long]   sample size
//                 16 [long]   number of entries in sample table
//                 20 [long[]] sample table if samples size is 0
bool ProcessStszBox(StreamInterface* stream, const Box& stsz_box,
                    uint32_t* image_size) {
  uint32_t sample_size;
  if (!Get32u(stream, stsz_box.offset() + 12, &sample_size)) {
    return false;
  }
  if (sample_size > 0) {
    *image_size = sample_size;
    return true;
  }
  // sample_size of 0 implies the data is in the sample table. We expect only
  // one entry. This is true of Canon EOS RP Cr3 files.
  uint32_t count;
  if (!Get32u(stream, stsz_box.offset() + 16, &count)) {
    return false;
  }
  if (count != 1) {
    // Expect at most one entry in the table.
    return false;
  }
  return Get32u(stream, stsz_box.offset() + 20, image_size);
}

// co64 box offsets:
//        Byte Offset Type     Meaning
//                  0 [long]   size of box
//                  4 [char[]] box tag
//                  8 [long]   version
//                 12 [long]   count (expect to be value 1)
//                 16 [long]   offset of image data in mdat
bool ProcessCo64(StreamInterface* stream, const Box& co64_box,
                 uint32_t* image_offset) {
  uint32_t count = 0;
  if (!Get32u(stream, co64_box.header_offset() + 4, &count)) {
    return false;
  }
  if (count != 1) {
    return false;
  }
  return Get32u(stream, co64_box.header_offset() + 8, image_offset);
}

// Process the stbl box. Expected box layout:
// stbl
//   stsd
//     CRAW  (embedded image (JPEG) information)
//   (0 or more skipped boxes)
//   stsz (embedded image byte size)
//   (0 or more skipped boxes)
//   co64 (offset of embedded image, relative to mdat box)
bool ProcessStblBox(StreamInterface* stream, const Box& stbl_box,
                    ProcessData* data) {
  Box stsd_box = GetNextBoxWithTag(stream, stbl_box.header_offset(), kStsdTag);
  if (!stsd_box.IsValid()) {
    return false;
  }
  // This is either CRAW or CTMD. Skip when CTMD.
  Box craw_box = GetNextBox(stream, stsd_box.header_offset() + 8);
  if (!craw_box.IsValid()) {
    return false;
  }
  if (craw_box.tag() != kCrawTag) {
    return true;
  }
  // CRAW contains info about the full-size image embedded in the mdat box.
  // The image is either JPEG or HEVC.
  uint16_t image_width = 0;
  uint16_t image_height = 0;
  uint16_t content_type = 0;
  if (!ProcessCrawBox(stream, craw_box, &image_width, &image_height,
                      &content_type)) {
    return false;
  }
  // Only continue if JPEG or HEVC content.
  constexpr uint16_t kJpegContentType = 3;
  constexpr uint16_t kHevcContentType = 4;
  if (content_type != kJpegContentType && content_type != kHevcContentType) {
    return true;
  }

  // Skip until we find stsz, contains the size (# of bytes) of image data.
  Box stsz_box =
      GetNextBoxWithTag(stream, stsd_box.next_box_offset(), kStszTag);
  if (!stsz_box.IsValid()) {
    return false;
  }
  uint32_t image_size;
  if (!ProcessStszBox(stream, stsz_box, &image_size)) {
    return false;
  }

  // Skip until we find co64, contains the offset of image data.
  Box co64_box =
      GetNextBoxWithTag(stream, stsz_box.next_box_offset(), kCo64Tag);
  if (!co64_box.IsValid()) {
    return false;
  }

  uint32_t image_offset = 0;
  if (!ProcessCo64(stream, co64_box, &image_offset)) {
    return false;
  }

  data->mdat_image.format = content_type == kJpegContentType
                                ? Image::kJpegCompressed
                                : Image::kHevcCompressed;
  data->mdat_image.width = image_width;
  data->mdat_image.height = image_height;
  data->mdat_image.length = image_size;
  // This offset is relative to the position of the mdat box. The value will
  // be updated once mdat's offset is found.
  data->mdat_image.offset = image_offset;
  return true;
}

// Returns true if we should parse the children of the box.
bool DoProcessChildren(const BoxTag& tag) {
  static const std::set<BoxTag> kTags = {NewTag("trak"), NewTag("moov"),
                                         NewTag("mdia"), NewTag("minf")};
  return kTags.find(tag) != kTags.end();
}

// Processes box and returns offset of the next box to process.
// A return value of 0 indicates an error.
//
// Outline of hierarchy and important boxes:
// ftyp
// moov
//   uuid (id is kUuidMoov)
//     ... boxes we skip ...
//     CMT1 (EXIF data)
//     CMT2 (EXIF data)
//     ... boxes we skip ...
//     THMB (160x120 JPEG thumbnail, embedded in this box)
//   trak
//     tkhd
//     mdia
//     ... boxes we skip ...
//     minf
//       ... boxes we skip ...
//       stbl
//         stsd
//           CRAW (Full image preview, type (JPEG or HEVC), width, height. The
//                 image data is found in mdat box, below.)
//       ... boxes we skip ...
//       stsz (Size of preview, in bytes)
//       ... boxes we skip ...
//       co64 (Location/offset of full preview data in mdat)
//   .. boxes we skip ...
// uuid (id is kUuidPrvw)
//   PRVW (1620x1080 JPEG preview, embedded in this box)
// mdat
//   Full image preview (JPEG or HEVC)
//   ... RAW image data ...
size_t ProcessBox(StreamInterface* stream, const Box& box, ProcessData* data) {
  // Parse child boxes.
  if (box.tag() == kUuidTag) {
    // Uuid box have extended box types.
    Uuid uuid;
    if (stream->GetData(box.header_offset(), uuid.size(), uuid.data()) != kOk) {
      return 0;
    }
    if (uuid == kUuidPrvw) {
      return box.header_offset() + uuid.size() + 8;
    } else if (uuid == kUuidMoov) {
      return box.header_offset() + uuid.size();
    }  // else skip the box, below.
  } else if (DoProcessChildren(box.tag())) {
    return box.header_offset();
  }

  // Potentially process the data contained in the box.
  bool success;
  if (box.tag() == kMdatTag) {
    // mdat_image.offset is relative to mdat's header, update it to be absolute
    // offset to the image data.
    data->mdat_image.offset += box.header_offset();
    success = true;
  } else if (box.tag() == kStblTag) {
    success = ProcessStblBox(stream, box, data);
  } else if (box.tag() == kPrvwTag) {
    // Preview jpeg. 1620x1080 for EOS R.
    success = ProcessJpegBox(stream, box, &data->prvw_image);
  } else if (box.tag() == kThmbTag) {
    // Thumbnail jpeg. 160x120 for EOS R.
    success = ProcessJpegBox(stream, box, &data->preview_image_data->thumbnail);
  } else if (box.tag() == kCmt1Tag) {
    success =
        ParseExifIfd0(stream, box.header_offset(), data->preview_image_data);
  } else if (box.tag() == kCmt2Tag) {
    success =
        ParseExifExifIfd(stream, box.header_offset(), data->preview_image_data);
  } else {
    // This box isn't interesting, skip it.
    success = true;
  }
  return success ? box.next_box_offset() : 0;
}

bool ProcessStream(StreamInterface* stream, const BoxTag& last_chunk,
                   ProcessData* data) {
  size_t offset = 0;
  while (true) {
    Box box = GetNextBox(stream, offset);
    if (!box.IsValid()) {
      return false;
    }
    size_t new_offset = ProcessBox(stream, box, data);
    if (new_offset <= offset) {
      return false;
    }
    if (box.tag() == last_chunk) {
      return true;
    }
    offset = new_offset;
  }
}

bool IsImage(StreamInterface* stream, const Image& image) {
  if (image.format != Image::kJpegCompressed) {
    // Pass responsibility to the caller.
    return true;
  }
  // Check for JPEG magic number at start. This could be HEVC data.
  constexpr std::array<uint8_t, 3> kJpegMagicNumber = {0xFF, 0xD8, 0xFF};
  std::array<uint8_t, 3> magic_number;
  if (stream->GetData(image.offset, magic_number.size(), magic_number.data()) !=
      kOk) {
    return false;
  }
  return magic_number == kJpegMagicNumber;
}

}  // namespace

Error Cr3GetPreviewData(StreamInterface* stream,
                        PreviewImageData* preview_image_data) {
  ProcessData data{.preview_image_data = preview_image_data};
  if (!ProcessStream(stream, kMdatTag, &data)) {
    return kFail;
  }
  // Prefer image in mdata box, as spec ensures it is the largest image.
  if (data.mdat_image.length > 0 && IsImage(stream, data.mdat_image)) {
    preview_image_data->preview = data.mdat_image;
  } else if (data.prvw_image.length > 0 && IsImage(stream, data.prvw_image)) {
    preview_image_data->preview = data.prvw_image;
  } else {
    return kFail;
  }
  return kOk;
}

bool Cr3GetOrientation(StreamInterface* stream, std::uint32_t* orientation) {
  PreviewImageData preview_image_data;
  ProcessData data{.preview_image_data = &preview_image_data};
  if (ProcessStream(stream, kCmt1Tag, &data)) {
    *orientation = preview_image_data.exif_orientation;
    return true;
  }
  return false;
}

}  // namespace piex