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

#include <algorithm>
#include <cstddef>
#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/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/synchronization/notification.h"
#include "async/async_context.h"
#include "crypto/algorithms.h"
#include "formats/bmff/box_header.h"
#include "formats/bmff/embedder.h"
#include "formats/bmff/hasher.h"
#include "formats/byte_range.h"
#include "proto/bmff_based_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 {

constexpr int64_t kCreateFreeBoxThreshold = 30;

// Scans through the source, modifying the source's values to adjust offsets to
// account for the manifest size as needed and hashing the applicable bytes.
// The method will return the offset where the manifest store should be
// inserted, the length of the manifest store, and an ordered vector of byte
// ranges of atoms which require offset adjustment before being written to the
// destination.
absl::StatusOr<std::vector<ByteInstruction>> InspectAndHash(
    riegeli::Reader& source, uint64_t manifest_store_size,
    const ByteRange& insertion_point, bool can_replace_manifest,
    BmffBasedHashAssertion& assertion) {
  if (!source.Seek(0) || source.pos() != 0) {
    return absl::InvalidArgumentError("cannot seek to start of input");
  }
  BmffEmbedder embedder;
  ABSL_ASSIGN_OR_RETURN(
      std::vector<ByteInstruction> instructions,
      embedder.InspectAndHash(source, manifest_store_size, insertion_point,
                              can_replace_manifest, assertion));
  return std::move(instructions);
}

// Wraps the manifest store in an atom.
absl::StatusOr<std::string> WrapManifestStoreInAtom(
    absl::string_view manifest_store) {
  std::string atom;
  riegeli::StringWriter writer(&atom);
  ABSL_RETURN_IF_ERROR(WriteC2paBmff(manifest_store, writer));
  if (!writer.Close()) {
    return absl::InternalError("failed to close string writer");
  }
  return std::move(atom);
}

absl::Status GenerateManifestStoreInstructions(
    const ByteRange& insertion_point, absl::string_view manifest_store,
    std::vector<ByteInstruction>& byte_instructions,
    bool client_supplied_hard_binding = false) {
  // Wrap the manifest store in an atom.
  ABSL_ASSIGN_OR_RETURN(std::string manifest_atom,
                        WrapManifestStoreInAtom(manifest_store));

  if (client_supplied_hard_binding &&
      insertion_point.length < manifest_atom.size()) {
    return absl::ResourceExhaustedError(absl::StrCat(
        "not enough space to embed C2PA manifest with client-supplied hard "
        "binding (insertion_point.length=",
        insertion_point.length, ", manifest_atom.size=", manifest_atom.size(),
        "); skipping embedding to avoid byte shifting"));
  }

  if (insertion_point.length == 0) {
    // No overwritable bytes, only insert.
    byte_instructions.push_back(
        {.operation = ByteInstruction::Operation::kInsert,
         .offset = insertion_point.offset,
         .bytes = std::move(manifest_atom)});
    return absl::OkStatus();
  }

  if (insertion_point.length == manifest_atom.size()) {
    // The atom fits exactly, just overwrite.
    byte_instructions.push_back(
        {.operation = ByteInstruction::Operation::kReplace,
         .offset = insertion_point.offset,
         .bytes = std::move(manifest_atom)});
    return absl::OkStatus();
  }

  if (insertion_point.length < manifest_atom.size()) {
    // Not enough space, overwrite the initial data, insert the
    // remaining.
    byte_instructions.push_back(
        {.operation = ByteInstruction::Operation::kReplace,
         .offset = insertion_point.offset,
         .bytes = manifest_atom.substr(0, insertion_point.length)});
    byte_instructions.push_back(
        {.operation = ByteInstruction::Operation::kInsert,
         .offset = insertion_point.offset + insertion_point.length,
         .bytes = manifest_atom.substr(insertion_point.length)});
    return absl::OkStatus();
  }

  // More than enough space
  int64_t manifest_size = manifest_atom.size();
  int64_t padding_size = insertion_point.length - manifest_size;

  if (padding_size < kCreateFreeBoxThreshold) {
    // Minor amount of remaining space, just pad with \0.
    ABSL_ASSIGN_OR_RETURN(
        manifest_atom, WrapManifestStoreInAtom(absl::StrCat(
                           manifest_store, std::string(padding_size, '\0'))));
    byte_instructions.push_back(
        {.operation = ByteInstruction::Operation::kReplace,
         .offset = insertion_point.offset,
         .bytes = std::move(manifest_atom)});
    return absl::OkStatus();
  }

  // Significant amount of remaining space, create a free box.
  std::string free_box;
  riegeli::StringWriter<> free_box_writer(&free_box);
  ABSL_RETURN_IF_ERROR(CreateFreeBox(padding_size, free_box_writer));
  if (!free_box_writer.Close()) {
    return absl::InternalError("failed to close string writer");
  }

  byte_instructions.push_back(
      {.operation = ByteInstruction::Operation::kReplace,
       .offset = insertion_point.offset,
       .bytes = std::move(manifest_atom)});
  byte_instructions.push_back(
      {.operation = ByteInstruction::Operation::kReplace,
       .offset = insertion_point.offset + manifest_size,
       .bytes = std::move(free_box)});

  return absl::OkStatus();
}

absl::StatusOr<bool> CanReplaceManifest(
    const GenerateManifestStoreRequest& request) {
  if (!request.has_actions()) {
    return absl::InvalidArgumentError("request must include `actions`");
  }
  if (request.actions().inception_action_case() ==
      Actions::INCEPTION_ACTION_NOT_SET) {
    return absl::InvalidArgumentError(
        "request must include an inception action (created or opened)");
  }
  return request.actions().has_opened();
}

}  // namespace

absl::Status BmffWriter::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 BmffWriter::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 {
  if (!source->SupportsRandomAccess()) {
    AsyncResolver<void>(std::move(context), std::move(cb))
        .Return(absl::InvalidArgumentError(
            "source does not support random access"));
    return;
  }

  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,
       byte_instructions, destination](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 BmffWriter::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 BmffWriter::Write(GenerateManifestStoreRequest request,
                       riegeli::Writer* absl_nonnull asset,
                       OutputMetadata* absl_nullable output_metadata,
                       AsyncContext context,
                       absl::AnyInvocable<void(absl::Status) &&> cb) const {
  if (!asset->SupportsRandomAccess()) {
    AsyncResolver<void>(std::move(context), std::move(cb))
        .Return(
            absl::InvalidArgumentError("asset does not support random access"));
    return;
  }
  if (!asset->SupportsReadMode()) {
    AsyncResolver<void>(std::move(context), std::move(cb))
        .Return(absl::InvalidArgumentError("asset does not support read mode"));
    return;
  }
  if (request.options().has_max_bytes_to_shift() &&
      request.options().max_bytes_to_shift() < 0) {
    AsyncResolver<void>(std::move(context), std::move(cb))
        .Return(absl::InvalidArgumentError(
            "max_bytes_to_shift must be non-negative"));
    return;
  }

  auto options = request.options();
  riegeli::Reader* source = asset->ReadMode(0);
  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)), asset,
       options = std::move(options),
       byte_instructions](absl::Status status) mutable {
        if (!status.ok()) {
          resolver.Return(std::move(status));
          return;
        }

        std::optional<size_t> insert_idx;
        auto first_insert_it = std::find_if(
            byte_instructions->begin(), byte_instructions->end(),
            [](const ByteInstruction& instr) {
              return instr.operation == ByteInstruction::Operation::kInsert;
            });
        if (first_insert_it != byte_instructions->end()) {
          insert_idx = first_insert_it - byte_instructions->begin();
        }

        auto file_size = asset->Size();
        uint64_t bytes_shifted = 0;
        if (file_size.has_value() && insert_idx.has_value() &&
            *file_size > (*byte_instructions)[*insert_idx].offset) {
          bytes_shifted = *file_size - (*byte_instructions)[*insert_idx].offset;
        }

        if (options.has_max_bytes_to_shift() &&
            bytes_shifted >
                static_cast<uint64_t>(options.max_bytes_to_shift())) {
          uint64_t overwritable_bytes = 0;
          if (*insert_idx > 0) {
            const auto& prev_instruction =
                (*byte_instructions)[*insert_idx - 1];
            if (prev_instruction.operation ==
                    ByteInstruction::Operation::kReplace &&
                prev_instruction.offset + prev_instruction.bytes.size() ==
                    (*byte_instructions)[*insert_idx].offset) {
              overwritable_bytes = prev_instruction.bytes.size();
            }
          }
          uint64_t manifest_atom_size =
              overwritable_bytes +
              (*byte_instructions)[*insert_idx].bytes.size();
          resolver.Return(absl::InvalidArgumentError(absl::StrCat(
              "embedding requires shifting ", bytes_shifted,
              " bytes, which exceeds the limit of ",
              options.max_bytes_to_shift(),
              " bytes (pre-allocated manifest size: ", overwritable_bytes,
              " bytes, final manifest size: ", manifest_atom_size, " bytes)")));
          return;
        }

        if (!asset->Seek(0)) {
          resolver.Return(
              absl::InternalError("failed to seek to start of asset"));
          return;
        }
        status = ApplyByteInstructions(*byte_instructions, asset);
        resolver.Return(std::move(status));
      });
}

void BmffWriter::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 (source->pos() != 0) {
    resolver.Return(absl::InvalidArgumentError(
        "source must be positioned at the start of the file"));
    return;
  }
  auto can_replace_manifest = CanReplaceManifest(request);
  if (!can_replace_manifest.ok()) {
    resolver.Return(
        absl::Status(can_replace_manifest.status().code(),
                     absl::StrCat(can_replace_manifest.status().message(),
                                  "; verifying request has valid actions")));
    return;
  }

  if (request.has_hard_binding() &&
      !request.hard_binding().has_bmff_based_hash()) {
    resolver.Return(absl::InvalidArgumentError(
        "request `hard_binding` must be a `bmff_based_hash` assertion"));
    return;
  }

  auto insertion_point = LocateManifestInsertionPoint(*source);
  if (!insertion_point.ok()) {
    resolver.Return(
        absl::Status(insertion_point.status().code(),
                     absl::StrCat(insertion_point.status().message(),
                                  "; locating manifest insertion point")));
    return;
  }

  if (request.has_hard_binding()) {
    generator_.GenerateManifestStore(
        std::move(request), context,
        [output_metadata, byte_instructions, insertion_point = *insertion_point,
         resolver = std::move(resolver)](
            absl::StatusOr<GenerateManifestStoreResponse> response) mutable {
          if (!response.ok()) {
            resolver.Return(response.status());
            return;
          }

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

          // Add the byte instructions for the manifest store atom to the
          // result.
          auto gen_status = GenerateManifestStoreInstructions(
              insertion_point, response->manifest_store(), *byte_instructions,
              /*client_supplied_hard_binding=*/true);
          if (!gen_status.ok()) {
            resolver.Return(std::move(gen_status));
            return;
          }

          // Sort all byte instructions by offset.
          std::sort(byte_instructions->begin(), byte_instructions->end(),
                    [](const ByteInstruction& a, const ByteInstruction& b) {
                      return a.offset < b.offset;
                    });

          if (output_metadata != nullptr) {
            output_metadata->timestamp_included = response->is_timestamped();
          }
          resolver.Return(absl::OkStatus());
        });
    return;
  }

  // Step 1: Estimate the manifest size
  *request.mutable_hard_binding()->mutable_bmff_based_hash() =
      GetBmffHashTemplate(HashAlgorithm::kSha256);

  if (options_.exclusion_processor != nullptr) {
    auto status = options_.exclusion_processor(*source, request);
    if (!status.ok()) {
      resolver.Return(std::move(status));
      return;
    }
  }

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

        // Step 2: Construct the final generated manifest store request.
        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_bmff_based_hash() =
            request.hard_binding().bmff_based_hash();

        // Step 3: Determine where the manifest should be inserted, where atoms
        // which need adjustment are located, and compute the hash of the
        // applicable bytes.
        auto inspect_res = InspectAndHash(
            *source,
            finalize_request.manifest_store_template().manifest_store_size(),
            insertion_point, can_replace_manifest,
            *finalize_request.mutable_hard_binding()
                 ->mutable_bmff_based_hash());
        if (!inspect_res.ok()) {
          resolver.Return(absl::Status(
              inspect_res.status().code(),
              absl::StrCat(inspect_res.status().message(),
                           "; embedding manifest store template")));
          return;
        }
        *byte_instructions = std::move(*inspect_res);

        // Step 4: Generate the finalized manifest store.
        const size_t expected_manifest_store_size =
            finalize_request.manifest_store_template().manifest_store_size();

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

              if (response->manifest_store().size() !=
                  expected_manifest_store_size) {
                resolver.Return(absl::InternalError(
                    "manifest store size does not match template"));
                return;
              }

              // Add the byte instructions for the manifest store atom to the
              // result.
              auto gen_status = GenerateManifestStoreInstructions(
                  insertion_point, response->manifest_store(),
                  *byte_instructions);
              if (!gen_status.ok()) {
                resolver.Return(std::move(gen_status));
                return;
              }

              // Sort all byte instructions by offset.
              std::sort(byte_instructions->begin(), byte_instructions->end(),
                        [](const ByteInstruction& a, const ByteInstruction& b) {
                          return a.offset < b.offset;
                        });

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

}  // namespace credentio