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

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <memory>
#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 "bindings/input_hasher.h"
#include "crypto/algorithms.h"
#include "formats/byte_range.h"
#include "formats/png/constants.h"
#include "formats/png/reader.h"
#include "proto/data_hash_assertion.pb.h"
#include "riegeli/bytes/reader.h"

namespace credentio {

namespace {

constexpr uint64_t kMaxChunkBytes = 1024 * 1024;  // 1 MiB

}  // namespace

absl::StatusOr<DigestPngResult> DigestPng(
    riegeli::Reader& input, HashAlgorithm alg,
    std::optional<PngExclusionProcessor> exclusion_processor) {
  DigestPngResult result;
  result.assertion.set_alg(absl::StrCat(alg));

  std::unique_ptr<InputHasher> hasher;
  ABSL_ASSIGN_OR_RETURN(hasher, InputHasher::Create(alg));

  std::optional<int64_t> insertion_offset = std::nullopt;

  ABSL_RETURN_IF_ERROR(IterateOverPngChunks(
      input,
      [&result, &input, &hasher, &insertion_offset,
       &exclusion_processor](const PngChunk& chunk) -> absl::Status {
        if (chunk.type == kPngChunkTypeC2pa) {
          if (insertion_offset.has_value()) {
            // Multiple C2PA chunks, or C2PA chunk after IDAT chunk.
            return absl::InvalidArgumentError("Unexpected C2PA chunk found");
          }
          // Since there is an existing C2PA chunk, overwrite it with the new
          // one.
          insertion_offset = chunk.offset;
          result.existing_c2pa_length = chunk.length;
          DataRange* exclusion = result.assertion.add_exclusions();
          exclusion->set_start(chunk.offset);
          exclusion->set_length(0);
          return absl::OkStatus();
        }

        if (!insertion_offset.has_value() && chunk.type == kPngChunkTypeIdat) {
          // Insert the manifest store immediately before the first IDAT chunk
          // unless there is an existing C2PA chunk.
          insertion_offset = chunk.offset;
          DataRange* exclusion = result.assertion.add_exclusions();
          exclusion->set_start(chunk.offset);
          exclusion->set_length(0);
        }

        if (exclusion_processor.has_value()) {
          ABSL_ASSIGN_OR_RETURN(std::vector<ByteRange> exclusions,
                                (*exclusion_processor)(input, chunk));
          if (!exclusions.empty()) {
            int64_t next_idx_to_hash = chunk.offset;
            for (const auto& ex : exclusions) {
              if (ex.offset < next_idx_to_hash ||
                  ex.offset + ex.length > chunk.offset + chunk.length) {
                return absl::InvalidArgumentError(
                    absl::StrCat("Invalid or out-of-order exclusion range [",
                                 ex.offset, ", ", ex.offset + ex.length,
                                 ") for chunk at ", chunk.offset));
              }

              if (next_idx_to_hash < ex.offset) {
                ABSL_RETURN_IF_ERROR(hasher->Update(
                    input, next_idx_to_hash, ex.offset - next_idx_to_hash));
              }

              DataRange* exclusion = result.assertion.add_exclusions();
              exclusion->set_start(ex.offset);
              exclusion->set_length(ex.length);

              next_idx_to_hash = ex.offset + ex.length;
            }
            // Hash remaining bytes (includes CRC)
            if (next_idx_to_hash < (chunk.offset + chunk.length)) {
              ABSL_RETURN_IF_ERROR(hasher->Update(
                  input, next_idx_to_hash,
                  chunk.length - (next_idx_to_hash - chunk.offset)));
            }
            return absl::OkStatus();
          }
        }

        if (!input.Seek(chunk.offset)) {
          return absl::InternalError("failed to seek to chunk offset");
        }
        uint64_t remaining_bytes = chunk.length;
        while (remaining_bytes > 0) {
          size_t chunk_size = std::min(remaining_bytes, kMaxChunkBytes);
          std::string chunk_data;
          if (!input.Read(chunk_size, chunk_data)) {
            return absl::DataLossError("failed to read box data");
          }
          hasher->Update(chunk_data);
          remaining_bytes -= chunk_size;
        }
        return absl::OkStatus();
      }));

  if (!insertion_offset.has_value()) {
    return absl::InvalidArgumentError("No IDAT or C2PA chunk found");
  }
  result.insertion_offset = *insertion_offset;
  // Complete the entry for the post-C2PA boxes.
  result.assertion.set_hash(hasher->Digest());
  return result;
}

}  // namespace credentio