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

#include <cstdint>
#include <optional>
#include <utility>

#include "absl/base/nullability.h"
#include "absl/functional/any_invocable.h"
#include "absl/log/die_if_null.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/notification.h"
#include "async/async_context.h"
#include "crypto/algorithms.h"
#include "formats/png/constants.h"
#include "formats/png/embedder.h"
#include "formats/png/hasher.h"
#include "formats/png/reader.h"
#include "generator/generator.h"
#include "proto/generator.pb.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/writer.h"

namespace credentio {

PngWriter::PngWriter(const ManifestStoreGenerator* absl_nonnull generator,
                     Options options)
    : generator_(*ABSL_DIE_IF_NULL(generator)), options_(std::move(options)) {}

absl::Status PngWriter::Write(GenerateManifestStoreRequest request,
                              riegeli::Reader* absl_nonnull source,
                              riegeli::Writer* absl_nonnull destination,
                              OutputMetadata* absl_nullable output_metadata) {
  absl::Notification done;
  absl::Status status;
  Write(std::move(request), source, destination, output_metadata,
        AsyncContext(), [&](absl::Status s) {
          status = std::move(s);
          done.Notify();
        });
  done.WaitForNotification();
  return status;
}

void PngWriter::Write(GenerateManifestStoreRequest request,
                      riegeli::Reader* absl_nonnull source,
                      riegeli::Writer* absl_nonnull destination,
                      OutputMetadata* absl_nullable output_metadata,
                      AsyncContext context,
                      absl::AnyInvocable<void(absl::Status) &&> cb) {
  AsyncResolver<void> resolver(context, std::move(cb));
  if (request.has_hard_binding()) {
    resolver.Return(absl::InvalidArgumentError(
        "request must not include `hard_binding` (binding is computed by "
        "PngWriter)"));
    return;
  }

  // DigestPng() takes an optional processor over (reader, chunk), so bind the
  // request, which stays alive for the duration of the digest below.
  auto exclusion_fn = [this, &request](riegeli::Reader& reader,
                                       const PngChunk& chunk) {
    return options_.exclusion_processor(reader, chunk, request);
  };
  std::optional<PngExclusionProcessor> exclusion_processor = std::nullopt;
  if (options_.exclusion_processor != nullptr) {
    exclusion_processor.emplace(std::move(exclusion_fn));
  }
  // Hashes the source and reports where the manifest store chunk goes. The
  // hash covers the asset as it will look once that chunk is in place: the
  // chunk itself is recorded as a zero-length exclusion for now, and is sized
  // below once the manifest store template tells us how big it will be.
  auto digest_result =
      DigestPng(*source, HashAlgorithm::kSha256, exclusion_processor);
  if (!digest_result.ok()) {
    resolver.Return(digest_result.status());
    return;
  }

  if (digest_result->existing_c2pa_length.has_value() &&
      (!request.has_actions() || !request.actions().has_opened())) {
    resolver.Return(absl::AlreadyExistsError(
        "C2PA chunk already exists, but no `opened` action was provided"));
    return;
  }

  *request.mutable_hard_binding()->mutable_data_hash() =
      digest_result->assertion;

  // The manifest store covers its own size, so it is generated in two phases:
  // a template reserves the space, then the exclusions are fixed up to account
  // for it, and only then is the manifest store finalized and embedded.
  generator_.GenerateManifestStoreTemplate(
      request, context,
      [this, source, request, digest_result = std::move(*digest_result),
       destination, output_metadata, context, resolver = std::move(resolver)](
          absl::StatusOr<ManifestStoreTemplate> template_response) mutable {
        if (!template_response.ok()) {
          resolver.Return(template_response.status());
          return;
        }

        constexpr int64_t kMaxManifestStoreSizeBytes =
            100 * 1024 * 1024;  // 100 MiB
        if (template_response->manifest_store_size() <= 0 ||
            template_response->manifest_store_size() >
                kMaxManifestStoreSizeBytes) {
          resolver.Return(absl::InvalidArgumentError(
              absl::StrCat("invalid manifest store template size: ",
                           template_response->manifest_store_size())));
          return;
        }

        FinalizeManifestStoreRequest finalize_request;
        *finalize_request.mutable_manifest_store_template() =
            std::move(*template_response);
        *finalize_request.mutable_hard_binding()->mutable_data_hash() =
            request.hard_binding().data_hash();

        // Size of the whole `caBX` chunk record that will be inserted.
        int64_t embedded_size =
            finalize_request.manifest_store_template().manifest_store_size() +
            kPngChunkLengthSize + kPngChunkTypeC2pa.size() + kPngCrcSize;
        // The exclusions were recorded against the source asset, but they must
        // describe the output asset. Inserting the chunk shifts everything
        // after the insertion point by the size it adds (the size of the new
        // chunk, less any existing `caBX` chunk it replaces), and the
        // placeholder exclusion at the insertion point becomes the chunk.
        for (auto& exclusion : *finalize_request.mutable_hard_binding()
                                    ->mutable_data_hash()
                                    ->mutable_exclusions()) {
          if (exclusion.start() > digest_result.insertion_offset) {
            exclusion.set_start(exclusion.start() + embedded_size -
                                digest_result.existing_c2pa_length.value_or(0));
          } else if (exclusion.start() == digest_result.insertion_offset) {
            exclusion.set_length(embedded_size);
          }
        }

        generator_.FinalizeManifestStore(
            finalize_request, context,
            [source, digest_result = std::move(digest_result), destination,
             output_metadata, resolver = std::move(resolver)](
                absl::StatusOr<FinalizeManifestStoreResponse>
                    finalize_response) mutable {
              if (!finalize_response.ok()) {
                resolver.Return(finalize_response.status());
                return;
              }

              if (output_metadata != nullptr) {
                output_metadata->timestamp_included =
                    finalize_response->is_timestamped();
              }

              PngEmbedder embedder;
              auto status = embedder.CopyAndEmbed(
                  *source, finalize_response->manifest_store(), digest_result,
                  *destination);
              resolver.Return(status);
            });
      });
}

}  // namespace credentio