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

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#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/status/statusor.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/jpeg/app11_io.h"
#include "formats/jpeg/hasher.h"
#include "formats/jpeg/reader.h"
#include "generator/generator.h"
#include "proto/data_hash_assertion.pb.h"
#include "proto/generator.pb.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_writer.h"
#include "riegeli/bytes/writer.h"
#include "utils/byte_instruction.h"

namespace credentio {
namespace {

void AdjustDataHashExclusions(DataHashAssertion* absl_nonnull assertion,
                              int64_t c2pa_offset, int64_t c2pa_size) {
  for (auto& exclusion : *assertion->mutable_exclusions()) {
    if (exclusion.start() < c2pa_offset) {
      // Before the C2PA segment, so no adjustment needed.
      continue;
    }
    if (exclusion.start() == c2pa_offset) {
      // This is the C2PA segment, just need to fill in the length.
      exclusion.set_length(c2pa_size);
      continue;
    }
    // This is after the C2PA segment, so we need to add the manifest size to
    // the offset.
    exclusion.set_start(exclusion.start() + c2pa_size);
  }
}

}  // namespace

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

absl::Status JpegWriter::Write(
    GenerateManifestStoreRequest request, riegeli::Reader* absl_nonnull source,
    riegeli::Writer* absl_nonnull destination,
    OutputMetadata* absl_nullable output_metadata) const {
  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 JpegWriter::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) const {
  auto byte_instructions = std::make_shared<std::vector<ByteInstruction>>();
  WriteInternal(
      std::move(request), source, output_metadata, byte_instructions, context,
      [resolver = AsyncResolver<void>(context, std::move(cb)), source,
       destination, byte_instructions](absl::Status status) mutable {
        if (!status.ok()) {
          resolver.Return(std::move(status));
          return;
        }
        status = ApplyByteInstructions(source, *byte_instructions, destination);
        resolver.Return(std::move(status));
      });
}

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

void JpegWriter::Write(GenerateManifestStoreRequest request,
                       riegeli::Writer* absl_nonnull asset,
                       OutputMetadata* absl_nullable output_metadata,
                       AsyncContext context,
                       absl::AnyInvocable<void(absl::Status) &&> cb) const {
  AsyncResolver<void> resolver(context, std::move(cb));
  if (!asset->SupportsRandomAccess()) {
    resolver.Return(
        absl::InvalidArgumentError("asset does not support random access"));
    return;
  }
  if (!asset->SupportsReadMode()) {
    resolver.Return(
        absl::InvalidArgumentError("asset does not support read mode"));
    return;
  }

  riegeli::Reader* source = asset->ReadMode(0);
  if (source == nullptr) {
    resolver.Return(asset->ok()
                        ? absl::InternalError("failed to create source reader")
                        : asset->status());
    return;
  }
  if (!source->ok()) {
    resolver.Return(source->status());
    return;
  }

  auto byte_instructions = std::make_shared<std::vector<ByteInstruction>>();
  WriteInternal(std::move(request), source, output_metadata, byte_instructions,
                context,
                [resolver = std::move(resolver), asset,
                 byte_instructions](absl::Status status) mutable {
                  if (!status.ok()) {
                    resolver.Return(std::move(status));
                    return;
                  }
                  status = ApplyByteInstructions(*byte_instructions, asset);
                  resolver.Return(std::move(status));
                });
}

void JpegWriter::WriteInternal(
    GenerateManifestStoreRequest request, riegeli::Reader* absl_nonnull source,
    OutputMetadata* absl_nullable output_metadata,
    std::shared_ptr<std::vector<ByteInstruction>> byte_instructions,
    AsyncContext context, absl::AnyInvocable<void(absl::Status) &&> cb) const {
  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 "
        "JpegWriter)"));
    return;
  }

  auto exclusion_fn = [this, &request](riegeli::Reader& reader,
                                       const JpegSegment& segment) {
    return options_.exclusion_processor(reader, segment, request);
  };
  std::optional<JpegExclusionProcessor> exclusion_processor = std::nullopt;
  if (options_.exclusion_processor != nullptr) {
    exclusion_processor.emplace(std::move(exclusion_fn));
  }
  auto digest_result = DigestJpeg(*source, HashAlgorithm::kSha256,
                                  request.options().generate_multi_asset_hash(),
                                  exclusion_processor);
  if (!digest_result.ok()) {
    resolver.Return(digest_result.status());
    return;
  }

  request.mutable_hard_binding()->Swap(&digest_result->hard_binding);
  auto insertion_point = std::move(digest_result->insertion_point);

  generator_.GenerateManifestStoreTemplate(
      request, context,
      [this, request, insertion_point = std::move(insertion_point),
       output_metadata, byte_instructions, context,
       resolver = std::move(resolver)](
          absl::StatusOr<ManifestStoreTemplate> template_response) mutable {
        if (!template_response.ok()) {
          resolver.Return(template_response.status());
          return;
        }

        // Do things with Generate Template Response
        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;
        }

        uint64_t embedded_size = CalculateTotalJpegApp11SegmentsSize(
            template_response->manifest_store_size());
        if (embedded_size > UINT32_MAX) {
          resolver.Return(
              absl::InternalError("estimated manifest store size is too large "
                                  "to fit in a uint32_t"));
          return;
        }

        FinalizeManifestStoreRequest finalize_request;
        finalize_request.mutable_manifest_store_template()->Swap(
            &*template_response);
        finalize_request.mutable_hard_binding()->MergeFrom(
            request.hard_binding());

        AdjustDataHashExclusions(
            finalize_request.mutable_hard_binding()->mutable_data_hash(),
            insertion_point.offset, embedded_size);

        if (finalize_request.hard_binding().part_bindings_size() > 0) {
          AdjustDataHashExclusions(finalize_request.mutable_hard_binding()
                                       ->mutable_part_bindings(0)
                                       ->mutable_data_hash(),
                                   insertion_point.offset, embedded_size);

          for (int i = 0;
               i < finalize_request.hard_binding().part_bindings_size(); ++i) {
            auto* location = finalize_request.mutable_hard_binding()
                                 ->mutable_part_bindings(i)
                                 ->mutable_location();
            if (location->byte_offset() == 0) {
              location->set_length(location->length() + embedded_size);
            } else {
              location->set_byte_offset(location->byte_offset() +
                                        embedded_size);
            }
          }
        }

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

              std::string app11_payload;
              riegeli::StringWriter app11_writer(&app11_payload);
              auto write_status = WriteJpegApp11Payload(
                  insertion_point.segment_id,
                  finalize_response->manifest_store(), app11_writer);
              if (!write_status.ok()) {
                resolver.Return(std::move(write_status));
                return;
              }
              if (!app11_writer.Close()) {
                resolver.Return(app11_writer.status());
                return;
              }

              byte_instructions->push_back(ByteInstruction{
                  .operation = ByteInstruction::Operation::kInsert,
                  .offset = static_cast<uint64_t>(insertion_point.offset),
                  .bytes = std::move(app11_payload)});

              if (output_metadata != nullptr) {
                output_metadata->timestamp_included =
                    finalize_response->is_timestamped();
              }
              resolver.Return(absl::OkStatus());
            });
      });
}
}  // namespace credentio