// 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/embedder.h"

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/algorithm/container.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/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "bindings/input_hasher.h"
#include "constants/labels.h"
#include "formats/bmff/box_header.h"
#include "formats/bmff/box_matcher.h"
#include "formats/bmff/offset_adjustment.h"
#include "formats/byte_range.h"
#include "proto/bmff_based_hash_assertion.pb.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_writer.h"
#include "riegeli/bytes/writer.h"
#include "riegeli/endian/endian_writing.h"
#include "utils/byte_instruction.h"
#include "utils/byte_writers.h"

namespace credentio {

namespace {

std::string Uint64ToBytes(uint64_t value) {
  std::vector<uint8_t> bytes;
  WriteUint64NetworkOrder(value, &bytes);
  return std::string(bytes.begin(), bytes.end());
}

constexpr absl::string_view kZeroFlags("\x00\x00\x00", 3);
constexpr absl::string_view kC2paBoxPurpose = "manifest";
constexpr uint64_t kMaxChunkBytes = 1024 * 1024;  // 1 MiB

BmffBox C2paBmffBoxHeader(uint64_t manifest_store_size) {
  return BmffBox{
      .type = "uuid",
      .user_type = std::string(kC2paBmffBoxUuid),
      .flags = std::string(kZeroFlags),
      // Nul-terminated purpose + 8-byte merkle offset + manifest_store.
      .data_size = kC2paBoxPurpose.size() + 1 + 8 + manifest_store_size,
  };
}

bool BoxHasOffsets(absl::string_view type) {
  return type == "stco" || type == "saio" || type == "co64" || type == "iloc";
}

// Adjusts the offsets in the given box by adding the size of the manifest
// store and writes it to the destination.
absl::Status FixOffsets(const BmffBoxHeader& box_header,
                        riegeli::Reader& source, uint64_t byte_adjustment,
                        riegeli::Writer& destination) {
  ABSL_ASSIGN_OR_RETURN(
      int64_t written_header_size,
      WriteBmffBoxHeader(
          BmffBox{
              .type = box_header.type,
              .user_type = box_header.user_type,
              .version = box_header.version,
              .flags = box_header.flags,
              // The box_size includes the header size, the write method expects
              // only the data size.
              .data_size = box_header.box_size - box_header.header_size -
                           box_header.version_and_flags_size,
          },
          destination));
  // Since we didn't read the header from the source, we need to skip over those
  // bytes.
  if (!source.Seek(box_header.start + written_header_size)) {
    return absl::DataLossError("failed to skip header size");
  }

  if (box_header.type == "iloc") {
    ABSL_ASSIGN_OR_RETURN(IlocBox box, ParseIloc(source, box_header));
    AdjustOffsets(box, byte_adjustment);
    return SerializeIloc(box, box_header, destination);
  }

  if (box_header.type == "stco") {
    ABSL_ASSIGN_OR_RETURN(StcoBox box, ParseStco(source));
    AdjustOffsets(box, byte_adjustment);
    return SerializeStco(box, destination);
  }

  if (box_header.type == "co64") {
    ABSL_ASSIGN_OR_RETURN(Co64Box box, ParseCo64(source));
    AdjustOffsets(box, byte_adjustment);
    return SerializeCo64(box, destination);
  }

  if (box_header.type == "saio") {
    ABSL_ASSIGN_OR_RETURN(SaioBox box, ParseSaio(source, box_header));
    AdjustOffsets(box, byte_adjustment);
    return SerializeSaio(box, box_header, destination);
  }

  return absl::InvalidArgumentError(
      absl::StrCat("Unexpected box type for offset fixing: ", box_header.type));
}

absl::Status ProcessC2paAtom(const BmffBoxHeader& box, riegeli::Reader& source,
                             const ByteRange& insertion_point,
                             std::vector<ByteInstruction>& byte_instructions) {
  if (box.start >= insertion_point.offset &&
      box.start + box.box_size <=
          insertion_point.offset + insertion_point.length) {
    // This box is completely within the insertion point, it will be handled by
    // the instruction which inserts the manifest store.
    return absl::OkStatus();
  }

  // Just convert any other C2PA atom outside of the insertion point to a free
  // box.
  std::string free_box;
  riegeli::StringWriter<> free_box_writer(&free_box);
  ABSL_RETURN_IF_ERROR(CreateFreeBox(box.box_size, free_box_writer));
  if (!free_box_writer.Close()) {
    return absl::InternalError("failed to close string writer");
  }

  byte_instructions.push_back(
      ByteInstruction{.operation = ByteInstruction::Operation::kReplace,
                      .offset = box.start,
                      .bytes = std::move(free_box)});

  return absl::OkStatus();
}

absl::Status ProcessBox(const BmffBoxHeader& box, riegeli::Reader& source,
                        bool can_replace_manifest,
                        absl::Span<const BMFFBoxMatcher> exclusion_matchers,
                        const ByteRange& insertion_point,
                        uint64_t bytes_needed_to_add,
                        std::vector<ByteInstruction>& byte_instructions,
                        std::vector<std::string>& excluded_atoms,
                        InputHasher& hasher) {
  if (box.type == "uuid" && box.user_type == kC2paBmffBoxUuid) {
    if (!can_replace_manifest) {
      return absl::AlreadyExistsError(
          "C2PA chunk already exists, but no `opened` action was provided");
    }

    ABSL_RETURN_IF_ERROR(
        ProcessC2paAtom(box, source, insertion_point, byte_instructions));

    // Skip this box
    source.Seek(box.start + box.box_size);
    return absl::OkStatus();
  }

  bool added_byte_instruction = false;
  if (BoxHasOffsets(box.type) && bytes_needed_to_add > 0) {
    // Move the current position to this box.
    if (!source.Seek(box.start) || source.pos() != box.start) {
      return absl::InternalError("failed to seek to box start");
    }

    auto byte_instruction = ByteInstruction{
        .operation = ByteInstruction::Operation::kReplace, .offset = box.start};

    riegeli::StringWriter<> writer(&byte_instruction.bytes);
    ABSL_RETURN_IF_ERROR(FixOffsets(box, source, bytes_needed_to_add, writer));
    if (!writer.Close()) {
      return absl::InternalError("failed to close string writer");
    }

    byte_instructions.push_back(std::move(byte_instruction));
    added_byte_instruction = true;
  }

  if (absl::c_any_of(excluded_atoms, [&box](absl::string_view excluded_atom) {
        return absl::StartsWith(box.xpath, excluded_atom);
      })) {
    // Parent atom is excluded, so skip this box.
    return absl::OkStatus();
  }
  for (const auto& matcher : exclusion_matchers) {
    if (matcher.Matches(box)) {
      excluded_atoms.push_back(box.xpath);

      if (matcher.IsMatchingEntireBox()) {
        // The entire box is excluded, so we can skip it.
        return absl::OkStatus();
      }
      // Note: Excluding parts of a box is not supported, as best practice would
      // see no additional exclusions aside from required exclusions.
      // See:
      // https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html#_exclusion_list_requirements
      return absl::InvalidArgumentError(
          "partial exclusions of atom are not supported");
    }
  }

  if (box.IsRootBox()) {
    int64_t adjusted_box_start = box.start >= insertion_point.offset
                                     ? box.start + bytes_needed_to_add
                                     : box.start;

    hasher.Update(Uint64ToBytes(adjusted_box_start));
  }

  if (added_byte_instruction) {
    hasher.Update(byte_instructions.back().bytes);
    return absl::OkStatus();
  }

  if (box.IsContainerBox()) {
    // The box is a container, so we need to hash the header then continue for
    // the data.
    int64_t data_start = source.pos();

    source.Seek(box.start);
    std::string header_data;
    source.Read(data_start - box.start, header_data);
    hasher.Update(header_data);

    return absl::OkStatus();
  }

  // not a container, not excluded, and no offsets, so just hash the whole box
  source.Seek(box.start);
  uint64_t remaining_bytes = box.box_size;
  while (remaining_bytes > 0) {
    size_t chunk_size = std::min(remaining_bytes, kMaxChunkBytes);
    std::string chunk_data;
    if (!source.Read(chunk_size, chunk_data)) {
      return absl::DataLossError("failed to read box data");
    }
    hasher.Update(chunk_data);
    remaining_bytes -= chunk_size;
  }

  return absl::OkStatus();
}

}  // namespace

absl::StatusOr<std::vector<ByteInstruction>> BmffEmbedder::InspectAndHash(
    riegeli::Reader& source, uint64_t manifest_store_size,
    const ByteRange& insertion_point, bool can_replace_manifest,
    BmffBasedHashAssertion& assertion) const {
  uint64_t c2pa_atom_size = CalculateC2paSizeBmff(manifest_store_size);
  uint64_t bytes_needed_to_add = insertion_point.length >= c2pa_atom_size
                                     ? 0
                                     : c2pa_atom_size - insertion_point.length;

  std::vector<BMFFBoxMatcher> exclusion_matchers;
  for (const auto& bmff_range : assertion.exclusions()) {
    exclusion_matchers.push_back(BMFFBoxMatcher(bmff_range, &source));
  }

  source.Seek(0);

  ABSL_ASSIGN_OR_RETURN(auto hasher, InputHasher::Create(assertion.alg()));
  if (!source.SupportsSize()) {
    return absl::InvalidArgumentError(
        "manifest store not embedded: source does not support size");
  }

  std::vector<ByteInstruction> byte_instructions;
  std::vector<std::string> excluded_atoms;
  ABSL_RETURN_IF_ERROR(IterateOverBmffBoxes(
      source,
      [&source, &insertion_point, &can_replace_manifest, &byte_instructions,
       &exclusion_matchers, &excluded_atoms, &hasher,
       &bytes_needed_to_add](const BmffBoxHeader& box) -> absl::StatusOr<bool> {
        ABSL_RETURN_IF_ERROR(ProcessBox(box, source, can_replace_manifest,
                                        exclusion_matchers, insertion_point,
                                        bytes_needed_to_add, byte_instructions,
                                        excluded_atoms, *hasher));
        return true;
      }));

  assertion.set_hash(hasher->Digest());

  return byte_instructions;
}

int64_t CalculateC2paSizeBmff(uint64_t manifest_store_size) {
  BmffBox box = C2paBmffBoxHeader(manifest_store_size);
  return BmffBoxHeaderSize(box) + box.data_size;
}

absl::Status WriteC2paBmff(absl::string_view manifest_store,
                           riegeli::Writer& output) {
  bool failed = false;
  ABSL_RETURN_IF_ERROR(
      WriteBmffBoxHeader(C2paBmffBoxHeader(manifest_store.size()), output)
          .status());
  failed |= !output.Write(kC2paBoxPurpose);
  failed |= !output.WriteByte(0);
  failed |= !riegeli::WriteBigEndian<uint64_t>(0, output);
  failed |= !output.Write(manifest_store);
  if (failed) {
    return absl::InternalError("manifest store not embedded: write failed");
  }
  return absl::OkStatus();
}

absl::Status CreateFreeBox(uint64_t size, riegeli::Writer& output) {
  if (size < 8) {
    return absl::InvalidArgumentError(
        absl::StrCat("Invalid free box size: ", size));
  }

  uint64_t data_size = 0;
  if (size > std::numeric_limits<uint32_t>::max()) {
    data_size = size - 16;
  } else {
    data_size = size - 8;
  }

  ABSL_RETURN_IF_ERROR(
      WriteBmffBoxHeader(BmffBox{.type = "free", .data_size = data_size},
                         output)
          .status());

  uint64_t remaining_padding = data_size;
  constexpr size_t kChunkSize = 64 * 1024;  // 64 KiB
  std::string zero_chunk(kChunkSize, '\0');
  while (remaining_padding > 0) {
    size_t write_size = std::min<uint64_t>(remaining_padding, kChunkSize);
    if (!output.Write(absl::string_view(zero_chunk.data(), write_size))) {
      return absl::InternalError("manifest store not embedded: write failed");
    }
    remaining_padding -= write_size;
  }

  return absl::OkStatus();
}

}  // namespace credentio