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

#include <algorithm>
#include <cstdint>
#include <optional>
#include <string>
#include <vector>

#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "formats/bmff/box_header.h"
#include "riegeli/base/types.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 {

// Caps upfront `.reserve()` at 2,000,000 entries (~8 MB for `stco`, ~16 MB for
// `co64`/`saio`, and up to ~80 MB for v2 `iloc` items; `iloc` `extent_count` is
// uint16_t, so at most ~1.6 MB per item).
constexpr uint64_t kMaxEntryCountToReserve = 2000000;

absl::StatusOr<uint64_t> ReadVariableSize(riegeli::Reader& reader,
                                          uint8_t size) {
  if (size == 0) return 0;
  if (size == 4) {
    uint32_t val;
    if (!riegeli::ReadBigEndian<uint32_t>(reader, val)) {
      return absl::DataLossError("Failed to read 4 byte integer");
    }
    return val;
  }
  if (size == 8) {
    uint64_t val;
    if (!riegeli::ReadBigEndian<uint64_t>(reader, val)) {
      return absl::DataLossError("Failed to read 8 byte integer");
    }
    return val;
  }
  return absl::DataLossError(absl::StrCat("invalid byte_size: ", size));
}

absl::Status WriteVariableSize(uint64_t value, uint8_t size,
                               riegeli::Writer& writer) {
  if (size == 0) return absl::OkStatus();
  if (size == 4) {
    if (value > 0xffffffff) {
      return absl::InvalidArgumentError("Value too large for 4 bytes");
    }
    if (!riegeli::WriteBigEndian<uint32_t>(static_cast<uint32_t>(value),
                                           writer)) {
      return writer.status();
    }
    return absl::OkStatus();
  }
  if (size == 8) {
    if (!riegeli::WriteBigEndian<uint64_t>(value, writer)) {
      return writer.status();
    }
    return absl::OkStatus();
  }
  return absl::InvalidArgumentError(absl::StrCat("invalid byte_size: ", size));
}

// Returns how many of `count` entries can be reserved up front (capped at
// `kMaxEntryCountToReserve`), or an error if `reader` does not have enough data
// left for `count` entries of at least `min_entry_size` bytes each. `count` is
// read from the input, so reserving it unchecked lets a few bytes of input
// allocate gigabytes of memory.
//
// Returns 0 if the size of `reader` is unknown, in which case entries should be
// added as they are read.
absl::StatusOr<uint64_t> EntriesToReserve(riegeli::Reader& reader,
                                          uint64_t count,
                                          uint64_t min_entry_size) {
  if (!reader.SupportsSize()) {
    return 0;
  }
  const std::optional<riegeli::Position> size = reader.Size();
  if (!size.has_value()) {
    return 0;
  }
  const riegeli::Position remaining =
      *size > reader.pos() ? *size - reader.pos() : 0;
  if (count > remaining / min_entry_size) {
    return absl::DataLossError(absl::StrCat(
        "entry count ", count, " exceeds the remaining ", remaining, " bytes"));
  }
  return std::min(count, kMaxEntryCountToReserve);
}

}  // namespace

absl::StatusOr<IlocBox> ParseIloc(riegeli::Reader& reader,
                                  const BmffBoxHeader& header) {
  IlocBox box;
  uint8_t offset_and_length_size;
  if (!riegeli::ReadBigEndian<uint8_t>(reader, offset_and_length_size)) {
    return absl::DataLossError("Failed to read offset and length size");
  }
  box.offset_size = offset_and_length_size >> 4;
  box.length_size = offset_and_length_size & 0xf;

  uint8_t base_offset_and_index_size;
  if (!riegeli::ReadBigEndian<uint8_t>(reader, base_offset_and_index_size)) {
    return absl::DataLossError("Failed to read base offset and index size");
  }
  box.base_offset_size = base_offset_and_index_size >> 4;
  box.index_size = base_offset_and_index_size & 0xf;

  uint32_t item_count = 0;
  if (header.version < 2) {
    uint16_t count;
    if (!riegeli::ReadBigEndian<uint16_t>(reader, count)) {
      return absl::DataLossError("Failed to read item count");
    }
    item_count = count;
  } else if (header.version == 2) {
    if (!riegeli::ReadBigEndian<uint32_t>(reader, item_count)) {
      return absl::DataLossError("Failed to read item count");
    }
  }

  // item_ID, [reserved and construction_method], data_reference_index,
  // base_offset, and extent_count.
  const uint64_t min_item_size =
      (header.version == 2 ? 4 : 2) +
      (header.version == 1 || header.version == 2 ? 2 : 0) + 2 +
      box.base_offset_size + 2;
  ABSL_ASSIGN_OR_RETURN(uint64_t items_to_reserve,
                        EntriesToReserve(reader, item_count, min_item_size));
  box.items.reserve(items_to_reserve);
  const uint64_t extent_size =
      (header.version == 1 || header.version == 2 ? box.index_size : 0) +
      box.offset_size + box.length_size;
  for (uint32_t i = 0; i < item_count; ++i) {
    IlocItem item;
    if (header.version < 2) {
      uint16_t item_id;
      if (!riegeli::ReadBigEndian<uint16_t>(reader, item_id)) {
        return absl::DataLossError("Failed to read item ID");
      }
      item.item_id = item_id;
    } else if (header.version == 2) {
      if (!riegeli::ReadBigEndian<uint32_t>(reader, item.item_id)) {
        return absl::DataLossError("Failed to read item ID");
      }
    }

    if (header.version == 1 || header.version == 2) {
      if (!riegeli::ReadBigEndian<uint16_t>(reader, item.construction_method)) {
        return absl::DataLossError(
            "Failed to read reserved and construction method");
      }
    }

    if (!riegeli::ReadBigEndian<uint16_t>(reader, item.data_reference_index)) {
      return absl::DataLossError("Failed to read data reference index");
    }

    ABSL_ASSIGN_OR_RETURN(item.base_offset,
                          ReadVariableSize(reader, box.base_offset_size));

    uint16_t extent_count;
    if (!riegeli::ReadBigEndian<uint16_t>(reader, extent_count)) {
      return absl::DataLossError("Failed to read extent count");
    }

    // With all extent fields absent, an extent occupies no bytes and refers to
    // the entire referenced container, so an item cannot be fragmented into
    // multiple extents.
    uint64_t extents_to_reserve = extent_count;
    if (extent_size > 0) {
      ABSL_ASSIGN_OR_RETURN(
          extents_to_reserve,
          EntriesToReserve(reader, extent_count, extent_size));
    } else if (extent_count > 1) {
      return absl::DataLossError(absl::StrCat("iloc item ", item.item_id,
                                              " has ", extent_count,
                                              " extents of size 0"));
    }
    item.extents.reserve(extents_to_reserve);
    for (int j = 0; j < extent_count; ++j) {
      IlocExtent extent;
      if ((header.version == 1 || header.version == 2) && box.index_size > 0) {
        ABSL_ASSIGN_OR_RETURN(extent.index,
                              ReadVariableSize(reader, box.index_size));
      }
      ABSL_ASSIGN_OR_RETURN(extent.offset,
                            ReadVariableSize(reader, box.offset_size));
      ABSL_ASSIGN_OR_RETURN(extent.length,
                            ReadVariableSize(reader, box.length_size));
      item.extents.push_back(extent);
    }
    box.items.push_back(item);
  }
  return box;
}

absl::Status SerializeIloc(const IlocBox& box, const BmffBoxHeader& header,
                           riegeli::Writer& writer) {
  uint8_t offset_and_length_size = (box.offset_size << 4) | box.length_size;
  if (!riegeli::WriteBigEndian<uint8_t>(offset_and_length_size, writer)) {
    return writer.status();
  }

  uint8_t base_offset_and_index_size =
      (box.base_offset_size << 4) | box.index_size;
  if (!riegeli::WriteBigEndian<uint8_t>(base_offset_and_index_size, writer)) {
    return writer.status();
  }

  if (header.version < 2) {
    if (box.items.size() > 0xffff) {
      return absl::InvalidArgumentError("Too many items for version < 2");
    }
    if (!riegeli::WriteBigEndian<uint16_t>(
            static_cast<uint16_t>(box.items.size()), writer)) {
      return writer.status();
    }
  } else if (header.version == 2) {
    if (!riegeli::WriteBigEndian<uint32_t>(
            static_cast<uint32_t>(box.items.size()), writer)) {
      return writer.status();
    }
  }

  for (const auto& item : box.items) {
    if (header.version < 2) {
      if (item.item_id > 0xffff) {
        return absl::InvalidArgumentError("Item ID too large for version < 2");
      }
      if (!riegeli::WriteBigEndian<uint16_t>(
              static_cast<uint16_t>(item.item_id), writer)) {
        return writer.status();
      }
    } else if (header.version == 2) {
      if (!riegeli::WriteBigEndian<uint32_t>(item.item_id, writer)) {
        return writer.status();
      }
    }

    if (header.version == 1 || header.version == 2) {
      if (!riegeli::WriteBigEndian<uint16_t>(item.construction_method,
                                             writer)) {
        return writer.status();
      }
    }

    if (!riegeli::WriteBigEndian<uint16_t>(item.data_reference_index, writer)) {
      return writer.status();
    }

    ABSL_RETURN_IF_ERROR(
        WriteVariableSize(item.base_offset, box.base_offset_size, writer));

    if (item.extents.size() > 0xffff) {
      return absl::InvalidArgumentError("Too many extents");
    }
    if (!riegeli::WriteBigEndian<uint16_t>(
            static_cast<uint16_t>(item.extents.size()), writer)) {
      return writer.status();
    }

    for (const auto& extent : item.extents) {
      if ((header.version == 1 || header.version == 2) && box.index_size > 0) {
        ABSL_RETURN_IF_ERROR(
            WriteVariableSize(extent.index, box.index_size, writer));
      }
      ABSL_RETURN_IF_ERROR(
          WriteVariableSize(extent.offset, box.offset_size, writer));
      ABSL_RETURN_IF_ERROR(
          WriteVariableSize(extent.length, box.length_size, writer));
    }
  }
  return absl::OkStatus();
}

void AdjustOffsets(IlocBox& box, int64_t adjustment) {
  for (auto& item : box.items) {
    if (item.data_reference_index == 0) {
      uint16_t method = item.construction_method & 0xf;
      if (method == 0) {
        if (box.base_offset_size > 0) {
          // The `extent.offset` is a relative offset. We must add it to the
          // base_offset to find the absolute position of the data.
          item.base_offset += adjustment;
        } else if (box.offset_size > 0) {
          // The `extent.offset` is an absolute offset. We must add the
          // adjustment to it directly.
          for (auto& extent : item.extents) {
            extent.offset += adjustment;
          }
        }
      }
    }
  }
}

absl::StatusOr<StcoBox> ParseStco(riegeli::Reader& reader) {
  StcoBox box;
  uint32_t entry_count;
  if (!riegeli::ReadBigEndian<uint32_t>(reader, entry_count)) {
    return absl::DataLossError("Failed to read entry count");
  }
  ABSL_ASSIGN_OR_RETURN(
      uint64_t entries_to_reserve,
      EntriesToReserve(reader, entry_count, sizeof(uint32_t)));
  box.offsets.reserve(entries_to_reserve);
  for (uint32_t i = 0; i < entry_count; ++i) {
    uint32_t offset;
    if (!riegeli::ReadBigEndian<uint32_t>(reader, offset)) {
      return absl::DataLossError("Failed to read offset");
    }
    box.offsets.push_back(offset);
  }
  return box;
}

absl::Status SerializeStco(const StcoBox& box, riegeli::Writer& writer) {
  if (box.offsets.size() > 0xffffffff) {
    return absl::InvalidArgumentError("Too many offsets");
  }
  if (!riegeli::WriteBigEndian<uint32_t>(
          static_cast<uint32_t>(box.offsets.size()), writer)) {
    return writer.status();
  }
  for (uint32_t offset : box.offsets) {
    if (!riegeli::WriteBigEndian<uint32_t>(offset, writer)) {
      return writer.status();
    }
  }
  return absl::OkStatus();
}

void AdjustOffsets(StcoBox& box, int64_t adjustment) {
  for (auto& offset : box.offsets) {
    offset += adjustment;
  }
}

absl::StatusOr<Co64Box> ParseCo64(riegeli::Reader& reader) {
  Co64Box box;
  uint32_t entry_count;
  if (!riegeli::ReadBigEndian<uint32_t>(reader, entry_count)) {
    return absl::DataLossError("Failed to read entry count");
  }
  ABSL_ASSIGN_OR_RETURN(
      uint64_t entries_to_reserve,
      EntriesToReserve(reader, entry_count, sizeof(uint64_t)));
  box.offsets.reserve(entries_to_reserve);
  for (uint32_t i = 0; i < entry_count; ++i) {
    uint64_t offset;
    if (!riegeli::ReadBigEndian<uint64_t>(reader, offset)) {
      return absl::DataLossError("Failed to read offset");
    }
    box.offsets.push_back(offset);
  }
  return box;
}

absl::Status SerializeCo64(const Co64Box& box, riegeli::Writer& writer) {
  if (box.offsets.size() > 0xffffffff) {
    return absl::InvalidArgumentError("Too many offsets");
  }
  if (!riegeli::WriteBigEndian<uint32_t>(
          static_cast<uint32_t>(box.offsets.size()), writer)) {
    return writer.status();
  }
  for (uint64_t offset : box.offsets) {
    if (!riegeli::WriteBigEndian<uint64_t>(offset, writer)) {
      return writer.status();
    }
  }
  return absl::OkStatus();
}

void AdjustOffsets(Co64Box& box, int64_t adjustment) {
  for (auto& offset : box.offsets) {
    offset += adjustment;
  }
}

absl::StatusOr<SaioBox> ParseSaio(riegeli::Reader& reader,
                                  const BmffBoxHeader& header) {
  SaioBox box;
  if ((header.flags[2] & 1) == 1) {
    if (!riegeli::ReadBigEndian<uint32_t>(reader, box.aux_info_type)) {
      return absl::DataLossError("Failed to read aux_info_type");
    }
    if (!riegeli::ReadBigEndian<uint32_t>(reader,
                                          box.aux_info_type_parameter)) {
      return absl::DataLossError("Failed to read aux_info_type_parameter");
    }
  }

  uint32_t entry_count;
  if (!riegeli::ReadBigEndian<uint32_t>(reader, entry_count)) {
    return absl::DataLossError("Failed to read entry count");
  }

  ABSL_ASSIGN_OR_RETURN(
      uint64_t entries_to_reserve,
      EntriesToReserve(
          reader, entry_count,
          header.version == 0 ? sizeof(uint32_t) : sizeof(uint64_t)));
  box.offsets.reserve(entries_to_reserve);
  for (uint32_t i = 0; i < entry_count; ++i) {
    if (header.version == 0) {
      uint32_t offset;
      if (!riegeli::ReadBigEndian<uint32_t>(reader, offset)) {
        return absl::DataLossError("Failed to read offset");
      }
      box.offsets.push_back(offset);
    } else {
      uint64_t offset;
      if (!riegeli::ReadBigEndian<uint64_t>(reader, offset)) {
        return absl::DataLossError("Failed to read offset");
      }
      box.offsets.push_back(offset);
    }
  }
  return box;
}

absl::Status SerializeSaio(const SaioBox& box, const BmffBoxHeader& header,
                           riegeli::Writer& writer) {
  if ((header.flags[2] & 1) == 1) {
    if (!riegeli::WriteBigEndian<uint32_t>(box.aux_info_type, writer)) {
      return writer.status();
    }
    if (!riegeli::WriteBigEndian<uint32_t>(box.aux_info_type_parameter,
                                           writer)) {
      return writer.status();
    }
  }

  if (box.offsets.size() > 0xffffffff) {
    return absl::InvalidArgumentError("Too many offsets");
  }
  if (!riegeli::WriteBigEndian<uint32_t>(
          static_cast<uint32_t>(box.offsets.size()), writer)) {
    return writer.status();
  }

  for (uint64_t offset : box.offsets) {
    if (header.version == 0) {
      if (offset > 0xffffffff) {
        return absl::InvalidArgumentError("Offset too large for version 0");
      }
      if (!riegeli::WriteBigEndian<uint32_t>(static_cast<uint32_t>(offset),
                                             writer)) {
        return writer.status();
      }
    } else {
      if (!riegeli::WriteBigEndian<uint64_t>(offset, writer)) {
        return writer.status();
      }
    }
  }
  return absl::OkStatus();
}

void AdjustOffsets(SaioBox& box, int64_t adjustment) {
  for (auto& offset : box.offsets) {
    offset += adjustment;
  }
}

}  // namespace credentio