// Copyright 2026 Google LLC
//
// 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
//
//    https://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 "formats/bmff/box_header.h"

#include <cstdint>
#include <limits>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "constants/labels.h"
#include "formats/bmff/constants.h"
#include "formats/bmff/xpath.h"
#include "formats/byte_range.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/writer.h"
#include "riegeli/endian/endian_reading.h"
#include "riegeli/endian/endian_writing.h"

namespace credentio {

namespace {

// Returns true if the given type is "uuid".
inline bool IsUuid(absl::string_view type) { return type == "uuid"; }

// ISOBMFF and QTFF define "meta" boxes differently. ISOBMFF defines it
// as a full box with flags and version (see 4.2 'Object Structure' in ISO/IEC
// 14496-12:2005), while QuickTime does not have such a full box structure.
// Here, we check the 8 bytes and see if it is the "hdlr" box, as defined by
// https://developer.apple.com/documentation/quicktime-file-format/metadata_atom.
bool IsQuickTimeMetaBox(const BmffBoxHeader& header, riegeli::Reader& input) {
  if (header.type != "meta") {
    return false;
  }
  // The next 8 bytes can be either:
  // (iso) [1 byte version + 3 bytes flags][4 byte size of next atom]
  // (qt)  [4 byte size of next atom      ][4 byte hdlr atom type   ]
  auto pin = input.pos();
  bool success = input.Seek(pin + 4);
  std::string type;
  success &= input.Read(4, type);
  success &= input.Seek(pin);
  return success && type == "hdlr";
};

constexpr int kAtomTypeSize = 4;
constexpr int kVersionFlagsSize = 4;
constexpr int kSizeOf32BitSize = 4;
constexpr int kSizeOf64BitSize = 8;

bool ShouldWriteSize64Bits(uint64_t box_size) {
  return box_size > std::numeric_limits<uint32_t>::max();
}

bool ShouldWriteUserType(const BmffBox& box) { return IsUuid(box.type); }

bool ShouldWriteVersionFlags(const BmffBox& box) {
  return IsBmffFullBox(box.type) || IsUuid(box.type);
}

constexpr int kMaxBmffDepth = 64;

void SetInsertionPoint(std::optional<ByteRange> candidate, int64_t c2pa_offset,
                       std::optional<ByteRange>& result) {
  if (!candidate.has_value()) {
    // No candidate
    return;
  }

  if (!result.has_value()) {
    // No existing
    result = candidate;
    return;
  }

  if (c2pa_offset == -1) {
    // No C2PA range found yet, use the larger range
    if (result->length < candidate->length) {
      // Existing is smaller and no C2PA range to account for
      result = candidate;
    }
    return;
  }

  if (c2pa_offset >= candidate->offset &&
      c2pa_offset <= candidate->offset + candidate->length) {
    // C2PA is contained within the candidate, use it.
    result = candidate;
  }
}

absl::StatusOr<bool> CanOverwrite(const BmffBoxHeader& box,
                                  riegeli::Reader& reader) {
  if (box.type == "free") {
    // Free boxes are always overwritable.
    return true;
  }
  if (box.type != "uuid" || box.user_type != kC2paBmffBoxUuid) {
    // The only other overwritable box is a C2PA uuid box.
    return false;
  }

  std::string purpose;
  if (!reader.Read(8, purpose)) {
    return reader.StatusOrAnnotate(
        absl::DataLossError("kUnexpectedEof; purpose"));
  }

  if (absl::StartsWith(purpose, kBmffC2paBoxPurposeMerkle)) {
    // Merkle boxes are not overwritable, we replace them with a free box.
    return false;
  }
  if (absl::StartsWith(purpose, kBmffC2paBoxPurposeUpdate)) {
    // Update boxes are not overwritable, we remove them.
    return false;
  }

  // Original and Manifest boxes are overwritable.
  return true;
}

absl::StatusOr<bool> IterateOverBmffBoxesInternal(XPath* parent,
                                                  riegeli::Reader& reader,
                                                  uint64_t end_offset,
                                                  BmffBoxProcessor processor,
                                                  int depth = 0) {
  if (depth > kMaxBmffDepth) {
    return absl::InvalidArgumentError("Too many nested BMFF boxes");
  }
  absl::flat_hash_map<std::string, int> count_by_type;
  while (reader.pos() < (end_offset - kBmffMinSizeofBoxHeader)) {
    BmffBoxHeader box;
    if (auto box_or = ReadBmffBoxHeader(reader); box_or.ok()) {
      box = std::move(box_or.value());
    } else {
      return box_or.status();
    }

    if (box.start > (std::numeric_limits<uint64_t>::max() - box.box_size)) {
      return reader.StatusOrAnnotate(
          absl::DataLossError("Final box offset exceeds uint64_t::max"));
    }
    uint64_t box_end_offset = box.start + box.box_size;

    // If the box extends beyond the end of the input, it's invalid.
    if (box_end_offset > end_offset) {
      if (parent == nullptr) {
        return reader.StatusOrAnnotate(
            absl::DataLossError("truncated BMFF box"));
      }
      return reader.StatusOrAnnotate(
          absl::DataLossError("sub-box extends beyond the parent"));
    }

    XPath xpath(parent, box.type);
    xpath.SetPosition(++count_by_type[box.type]);
    box.xpath = xpath.ToString();

    // If it fails during processing, we can't continue.
    bool continue_processing;
    if (auto result = processor(box); result.ok()) {
      continue_processing = *result;
    } else {
      return result.status();
    }
    if (!continue_processing) {
      return false;
    }

    if (reader.pos() < box_end_offset && box.IsContainerBox()) {
      if (auto sub_continue_or = IterateOverBmffBoxesInternal(
              &xpath, reader, box_end_offset, processor, depth + 1);
          sub_continue_or.ok()) {
        continue_processing = *sub_continue_or;
      } else {
        return sub_continue_or.status();
      }
      if (!continue_processing) {
        return false;
      }
    }

    reader.Seek(box_end_offset);
  }
  return true;
}

}  // namespace

bool BmffBoxHeader::IsContainerBox() const {
  // This list mirrors the C2PA SDK's implementation:
  // https://github.com/contentauth/c2pa-rs/blob/main/sdk/src/asset_handlers/bmff_io.rs#L982.
  return type == "moov" || type == "trak" || type == "mdia" || type == "minf" ||
         type == "stbl" || type == "moof" || type == "traf" || type == "edts" ||
         type == "udta" || type == "dinf" || type == "tref" || type == "treg" ||
         type == "mvex" || type == "mfra" || type == "meta" || type == "schi";
}

absl::StatusOr<BmffBoxHeader> ReadBmffBoxHeader(riegeli::Reader& input) {
  BmffBoxHeader header;

  header.start = input.pos();

  uint32_t box_size32;
  if (!riegeli::ReadBigEndian<uint32_t>(input, box_size32)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("kUnexpectedEof; box_size32"));
  }
  if (!input.Read(4, header.type)) {
    return input.StatusOrAnnotate(absl::DataLossError("kUnexpectedEof; type"));
  }
  header.header_size = 8;

  if (box_size32 == 1) {
    if (!riegeli::ReadBigEndian<uint64_t>(input, header.box_size)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("kUnexpectedEof; box_size64"));
    }
    header.header_size = 16;
  } else {
    header.box_size = box_size32;
  }

  if (header.box_size < kBmffMinSizeofBoxHeader) {
    // Unsupported box size.
    return input.StatusOrAnnotate(
        absl::DataLossError("kInvalidData; unsupported box size"));
  }

  if (IsUuid(header.type)) {
    if (!input.Read(kBmffUserTypeSize, header.user_type)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("kUnexpectedEof; user_type"));
    }
    header.header_size += kBmffUserTypeSize;
  }
  if ((IsBmffFullBox(header.type) ||
       (IsUuid(header.type) && header.user_type == kC2paBmffBoxUuid)) &&
      !IsQuickTimeMetaBox(header, input)) {
    header.version_and_flags_size = 4;
    if (!input.ReadByte(header.version)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("kUnexpectedEof; version"));
    }
    if (!input.Read(3, header.flags)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("kUnexpectedEof; flags"));
    }
  }
  if (header.box_size < header.header_size + header.version_and_flags_size) {
    return input.StatusOrAnnotate(
        absl::DataLossError("kInvalidData; box size too small for header"));
  }
  return header;
}

int64_t BmffBoxHeaderSize(const BmffBox& box) {
  uint64_t header_size = kSizeOf32BitSize + kAtomTypeSize;
  if (ShouldWriteUserType(box)) {
    header_size += kBmffUserTypeSize;
  }
  if (ShouldWriteVersionFlags(box)) {
    header_size += kVersionFlagsSize;
  }
  if (ShouldWriteSize64Bits(header_size + box.data_size)) {
    header_size += kSizeOf64BitSize;
  }
  return header_size;
}

absl::StatusOr<int64_t> WriteBmffBoxHeader(const BmffBox& box,
                                           riegeli::Writer& destination) {
  bool success = true;

  uint64_t header_size = BmffBoxHeaderSize(box);
  uint64_t box_size = header_size + box.data_size;
  bool should_write_size_64bits = ShouldWriteSize64Bits(box_size);
  uint32_t box_size32 = should_write_size_64bits ? 1 : box_size;

  success &= riegeli::WriteBigEndian<uint32_t>(box_size32, destination);
  success &= box.type.size() == kAtomTypeSize && destination.Write(box.type);
  if (should_write_size_64bits) {
    success &= riegeli::WriteBigEndian<uint64_t>(box_size, destination);
  }
  if (ShouldWriteUserType(box)) {
    success &= box.user_type.size() == kBmffUserTypeSize &&
               destination.Write(box.user_type);
  }
  if (ShouldWriteVersionFlags(box)) {
    success &= destination.WriteByte(box.version);
    success &= box.flags.size() == (kVersionFlagsSize - 1) &&
               destination.Write(box.flags);
  }
  if (!success) {
    return absl::InternalError("failed to write box header");
  }
  return header_size;
}

absl::StatusOr<std::vector<BmffBoxHeader>> ReadBmffBoxHeaders(
    riegeli::Reader& input) {
  std::vector<BmffBoxHeader> result;
  if (!input.Seek(0) || input.pos() != 0) {
    return input.StatusOrAnnotate(
        absl::DataLossError("failed to seek to start of input"));
  }
  auto result_status =
      IterateOverBmffBoxes(input, [&result](const BmffBoxHeader& box) {
        result.push_back(std::move(box));
        return true;
      });
  if (result_status.ok()) {
    return std::move(result);
  }
  return result_status;
}

absl::Status IterateOverBmffBoxes(riegeli::Reader& reader,
                                  BmffBoxProcessor processor) {
  if (!reader.SupportsSize() || !reader.Size().has_value()) {
    return absl::InvalidArgumentError(
        "manifest store not embedded: reader size cannot be determined");
  }
  if ((*reader.Size() - reader.pos()) == 0) {
    return absl::OkStatus();
  }

  return IterateOverBmffBoxesInternal(nullptr, reader, *reader.Size(),
                                      processor)
      .status();
}

absl::StatusOr<ByteRange> LocateManifestInsertionPoint(
    riegeli::Reader& reader) {
  std::optional<ByteRange> result = std::nullopt;
  std::optional<ByteRange> range = std::nullopt;

  int64_t c2pa_offset = -1;
  std::string last_atom_type_processed = "";
  auto iterate_status = IterateOverBmffBoxes(
      reader,
      [&reader, &range, &c2pa_offset, &last_atom_type_processed,
       &result](const BmffBoxHeader& box) -> absl::StatusOr<bool> {
        ABSL_ASSIGN_OR_RETURN(bool overwritable, CanOverwrite(box, reader));

        // We only care about the top level boxes, so skip all sub-boxes.
        reader.Seek(box.start + box.box_size);
        last_atom_type_processed = box.type;

        if (!overwritable) {
          SetInsertionPoint(range, c2pa_offset, result);
          range = std::nullopt;

          if (box.type == "mdat" || box.type == "moov") {
            SetInsertionPoint(ByteRange{.offset = box.start, .length = 0},
                              c2pa_offset, result);
            return false;  // Terminate loop
          }
          return true;  // Continue
        }

        if (box.type == "uuid") {
          c2pa_offset = box.start;
        }

        if (!range.has_value()) {
          range = ByteRange{.offset = box.start, .length = 0};
        }
        range->length += box.box_size;

        return true;  // Continue
      });
  ABSL_RETURN_IF_ERROR(iterate_status);
  if (last_atom_type_processed != "mdat" &&
      last_atom_type_processed != "moov") {
    return absl::InvalidArgumentError("no mdat/moov box found");
  }
  SetInsertionPoint(range, c2pa_offset, result);
  return result.value_or(ByteRange{.offset = 0, .length = 0});
}

}  // namespace credentio