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

#include <sys/types.h>

#include <algorithm>
#include <cstdint>
#include <optional>
#include <string>
#include <utility>
#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 "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "constants/labels.h"
#include "formats/asset_box.h"
#include "formats/byte_range.h"
#include "formats/png/constants.h"
#include "formats/png/crc.h"
#include "formats/png/reader.h"
#include "jumbf/utils.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/endian/endian_reading.h"

namespace credentio {

namespace {
constexpr uint64_t kMaxPayloadSize = 1024 * 1024 * 10;  // 10 MiB

absl::StatusOr<std::string> ValidateCrcAndReturnPayload(riegeli::Reader& input,
                                                        PngChunk chunk,
                                                        int64_t end_offset) {
  if (chunk.data_length > kMaxPayloadSize) {
    return absl::InvalidArgumentError(
        absl::Substitute("PNG C2PA chunk is too large to extract ($0 > $1)",
                         chunk.data_length, kMaxPayloadSize));
  }

  std::string payload;
  if (!input.Read(chunk.data_length, payload)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to read payload"));
  }

  uint32_t extracted_crc;
  if (!riegeli::ReadBigEndian<uint32_t>(input, extracted_crc)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read CRC"));
  }

  uint32_t computed_crc = PngChunkCrc(chunk.type, payload);
  if (extracted_crc != computed_crc) {
    return absl::InvalidArgumentError(
        absl::Substitute("chunk CRC 0x$0 does not match computed CRC 0x$1",
                         absl::Hex(extracted_crc), absl::Hex(computed_crc)));
  }

  return std::move(payload);
}

absl::Status ValidateAssetSize(riegeli::Reader& input, int64_t end_offset) {
  int64_t asset_size =
      (end_offset < 0 ? input.Size().value_or(input.pos()) : end_offset) -
      input.pos();
  if (asset_size < kPngMinimumAssetSize) {
    // The asset is too small, we can't possibly have a manifest store.
    return absl::NotFoundError("No manifest store found");
  }
  return absl::OkStatus();
}

absl::Status ValidateWithinAssetWindow(const PngChunk& chunk,
                                       int64_t end_offset) {
  // If we are iterating over the entire asset, return early.
  if (end_offset < 0) {
    return absl::OkStatus();
  }

  if (chunk.type == "c2pa.after") {
    return chunk.offset >= end_offset ? absl::AbortedError("asset-window-end")
                                      : absl::OkStatus();
  }

  if (chunk.offset >= end_offset) {
    // Chunk starts at or past the end of the asset window, early return.
    return absl::AbortedError("asset-window-end");
  }

  if (chunk.offset + chunk.length > end_offset) {
    // Chunk starts within the asset window, but extends beyond, return error.
    return absl::InvalidArgumentError(
        "PNG chunk extends beyond the end of the file");
  }

  // Chunk starts within the asset window and does not extend beyond it.
  return absl::OkStatus();
}

}  // namespace

absl::StatusOr<std::string> PngExtractor::ExtractManifestStore(
    riegeli::Reader& input) const {
  std::optional<std::string> result = std::nullopt;
  absl::Status iteration_status = IterateOverPngChunks(
      input, [&result, &input](const PngChunk& chunk) -> absl::Status {
        if (chunk.type != kPngChunkTypeC2pa) {
          return absl::OkStatus();
        }

        if (result.has_value()) {
          return absl::NotFoundError("Multiple manifest stores found");
        }

        ABSL_ASSIGN_OR_RETURN(
            result, ValidateCrcAndReturnPayload(input, chunk,
                                                input.Size().value_or(0)));
        return absl::OkStatus();
      });
  ABSL_RETURN_IF_ERROR(iteration_status);
  if (result.has_value()) {
    return *result;
  }
  return absl::NotFoundError("No manifest store found");
}

absl::StatusOr<std::optional<ByteRange>>
PngExtractor::ExtractManifestStoreLocation(riegeli::Reader& input,
                                           ExtractOptions options) const {
  std::optional<ByteRange> result = std::nullopt;

  if (!ValidateAssetSize(input, options.end_offset).ok()) {
    if (options.requires_c2pa) {
      return absl::NotFoundError("No manifest store found");
    }
    return result;
  }

  absl::Status iteration_status = IterateOverPngChunks(
      input,
      [&result, &options, &input](const PngChunk& chunk) -> absl::Status {
        ABSL_RETURN_IF_ERROR(
            ValidateWithinAssetWindow(chunk, options.end_offset));

        if (chunk.type != kPngChunkTypeC2pa) {
          return absl::OkStatus();
        }
        if (result.has_value()) {
          return absl::NotFoundError("Multiple manifest stores found");
        }

        // Ensure the C2PA chunk is valid before using it.
        ABSL_RETURN_IF_ERROR(
            ValidateCrcAndReturnPayload(input, chunk, input.Size().value_or(0))
                .status());

        result = {.offset = chunk.offset, .length = chunk.length};

        return absl::OkStatus();
      });
  if (iteration_status.code() != absl::StatusCode::kAborted) {
    // Ignore the aborted error, it means we've reached the end of the asset
    // window in a valid manner.
    ABSL_RETURN_IF_ERROR(iteration_status);
  }

  if (options.requires_c2pa && !result.has_value()) {
    return absl::NotFoundError("No manifest store found");
  }
  return result;
}

absl::StatusOr<std::vector<AssetBox>> PngExtractor::ExtractBoxes(
    riegeli::Reader& input, ExtractOptions options) const {
  std::vector<AssetBox> result;

  if (!ValidateAssetSize(input, options.end_offset).ok()) {
    if (options.requires_c2pa) {
      return absl::NotFoundError("No manifest store found");
    }
    return result;
  }

  int64_t c2pa_chunks_found = 0;
  absl::Status iteration_status = IterateOverPngChunks(
      input,
      [&result, &options, &c2pa_chunks_found,
       &input](const PngChunk& chunk) -> absl::Status {
        ABSL_RETURN_IF_ERROR(
            ValidateWithinAssetWindow(chunk, options.end_offset));

        std::string identifier = chunk.type;
        if (chunk.type == kPngChunkTypeC2pa) {
          c2pa_chunks_found++;
          identifier = "C2PA";

          if (c2pa_chunks_found > 1) {
            return absl::NotFoundError("Multiple manifest stores found");
          }
          // Ensure the C2PA chunk is valid.
          if (auto payload_or =
                  ValidateCrcAndReturnPayload(input, chunk, options.end_offset);
              !payload_or.ok()) {
            return payload_or.status();
          }
        }

        uint64_t length = chunk.length;
        if (identifier == "c2pa.after" && options.end_offset >= 0) {
          length = std::min(
              length, static_cast<uint64_t>(options.end_offset - chunk.offset));
        }
        result.push_back(AssetBox{
            .identifier = std::move(identifier),
            .byte_range = {.offset = chunk.offset, .length = length},
        });
        return absl::OkStatus();
      });
  if (iteration_status.code() != absl::StatusCode::kAborted) {
    // Ignore the aborted error, it means we've reached the end of the asset
    // window in a valid manner.
    ABSL_RETURN_IF_ERROR(iteration_status);
  }
  if (options.requires_c2pa && c2pa_chunks_found == 0) {
    return absl::NotFoundError("No manifest store found");
  }
  return result;
}

bool PngExtractor::MightBeC2paManifestStore(absl::string_view payload) const {
  return jumbf::HasDescriptionBoxMatching(payload, kManifestStoreUuid,
                                          kMinimumJumbfDescriptionToggles,
                                          kManifestStoreLabel)
      .value_or(false);
}

}  // namespace credentio