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

#include <sys/types.h>

#include <cstdint>
#include <string>
#include <utility>

#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/ascii.h"
#include "absl/strings/str_cat.h"
#include "formats/png/constants.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/endian/endian_reading.h"

namespace credentio {

namespace {

absl::StatusOr<PngChunk> ReadHeader(riegeli::Reader& reader) {
  uint64_t start_offset = reader.pos();
  std::string header;
  if (!reader.Read(kPngHeader.size(), header)) {
    return reader.StatusOrAnnotate(
        absl::InvalidArgumentError("Failed to read header"));
  }
  if (header != kPngHeader) {
    return absl::InvalidArgumentError("input does not start with PNG marker");
  }
  return PngChunk{
      .offset = start_offset,
      .length = kPngHeader.size(),
      .type = kPngChunkTypeHeader.data(),
      .data_length = static_cast<uint32_t>(kPngHeader.size()),
  };
}

bool ChunkGoesPastEndOffset(uint64_t chunk_start, uint64_t chunk_length,
                            uint64_t end_offset) {
  if (end_offset < chunk_length) {
    // Chunk is larger than the file size
    return true;
  }
  if (chunk_start > (end_offset - chunk_length)) {
    // Chunk starts too far into the file to read the entire chunk.
    return true;
  }
  return false;
}

absl::StatusOr<PngChunk> ReadChunk(riegeli::Reader& reader,
                                   int64_t end_offset) {
  PngChunk chunk{.offset = reader.pos()};
  if (!riegeli::ReadBigEndian<uint32_t>(reader, chunk.data_length)) {
    return reader.StatusOrAnnotate(
        absl::DataLossError("Failed to read chunk data length"));
  }
  chunk.length = sizeof(uint32_t)                            // Chunk length
                 + sizeof(uint32_t)                          // Chunk type
                 + static_cast<uint64_t>(chunk.data_length)  // Chunk data
                 + sizeof(uint32_t);                         // CRC
  if (ChunkGoesPastEndOffset(chunk.offset, chunk.length, end_offset)) {
    return absl::InvalidArgumentError(
        "PNG chunk extends beyond the end of the file");
  }
  if (!reader.Read(4, chunk.type)) {
    return reader.StatusOrAnnotate(
        absl::DataLossError("Failed to read chunk type"));
  }
  for (char c : chunk.type) {
    if (!absl::ascii_isalpha(c)) {
      return absl::InvalidArgumentError(
          absl::StrCat("Invalid PNG chunk type name: ", chunk.type));
    }
  }
  return std::move(chunk);
}

}  // namespace

absl::Status IterateOverPngChunks(riegeli::Reader& reader,
                                  PngChunkProcessor processor) {
  if (!reader.SupportsSize() || !reader.Size().has_value()) {
    return absl::InvalidArgumentError(
        "PNG reader does not support size, cannot iterate");
  }
  uint64_t end_offset = *reader.Size();
  if (kPngMinimumAssetSize > end_offset - reader.pos()) {
    // Asset size is smaller than the minimum size needed, just return.
    return absl::OkStatus();
  }

  ABSL_ASSIGN_OR_RETURN(PngChunk header, ReadHeader(reader));

  ABSL_RETURN_IF_ERROR(processor(header));
  uint64_t last_chunk_offset = header.offset;
  while (reader.pos() < end_offset) {
    ABSL_ASSIGN_OR_RETURN(PngChunk chunk, ReadChunk(reader, end_offset));

    if (chunk.offset <= last_chunk_offset) {
      return reader.StatusOrAnnotate(
          absl::InvalidArgumentError("PNG chunk offset did not advance"));
    }
    last_chunk_offset = chunk.offset;

    ABSL_RETURN_IF_ERROR(processor(chunk));
    if (!reader.Seek(chunk.offset + chunk.length)) {
      return reader.StatusOrAnnotate(
          absl::InvalidArgumentError("Failed to seek to next chunk"));
    }

    if (chunk.type == kPngChunkTypeEnd) {
      break;
    }
  }

  if (reader.pos() < end_offset) {
    PngChunk chunk{
        .offset = reader.pos(),
        .length = end_offset - reader.pos(),
        .type = "c2pa.after",
        .data_length = 0,
    };
    ABSL_RETURN_IF_ERROR(processor(chunk));
  }
  return absl::OkStatus();
}

}  // namespace credentio