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

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

#include "absl/base/nullability.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/types/span.h"
#include "bindings/input_hasher.h"
#include "crypto/algorithms.h"
#include "formats/hasher_data.h"
#include "formats/riff/constants.h"
#include "formats/riff/reader.h"
#include "proto/generator.pb.h"
#include "riegeli/base/byte_fill.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_writer.h"
#include "riegeli/bytes/writer.h"
#include "riegeli/endian/endian_writing.h"
#include "utils/byte_instruction.h"

namespace credentio {
namespace {

absl::StatusOr<uint64_t> FindC2paInsertionPoint(riegeli::Reader& input) {
  uint64_t c2pa_insertion_point = 0;

  bool first = true;
  ABSL_RETURN_IF_ERROR(IterateOverRiffChunks(
      input,
      [&c2pa_insertion_point,
       &first](const RiffChunk& chunk) -> absl::StatusOr<bool> {
        if (first) {
          // Verify that the input starts with RIFF chunk.
          if (chunk.id != kRiffChunkIdRiff) {
            return absl::InvalidArgumentError(
                "Input does not start with RIFF chunk");
          }
          first = false;
          c2pa_insertion_point = chunk.offset + chunk.length;
          return true;
        }
        if (chunk.id == kRiffChunkIdC2pa) {
          return absl::AlreadyExistsError("C2PA chunk already exists");
        }
        if (chunk.offset >= c2pa_insertion_point) {
          return false;  // Stop iterating, we finished the first RIFF chunk
        }
        return true;
      }));
  if (first) {
    return absl::InvalidArgumentError("Input is empty");
  }
  return c2pa_insertion_point;
}

absl::Status AddRiffSizeInstruction(
    uint32_t riff_size, std::vector<ByteInstruction>& byte_instructions) {
  byte_instructions.push_back({
      .operation = ByteInstruction::Operation::kReplace,
      .offset = 4,
      .bytes = "",
  });
  riegeli::StringWriter size_writer(&byte_instructions.back().bytes);
  if (!riegeli::WriteLittleEndian<uint32_t>(riff_size, size_writer)) {
    return absl::InternalError("failed to create size writer");
  }
  if (!size_writer.Close()) {
    return absl::InternalError("failed to write updated RIFF size");
  }
  return absl::OkStatus();
}

absl::Status AddC2paManifestInstructionPlaceholder(
    uint64_t c2pa_insertion_point, uint64_t manifest_store_size,
    uint64_t pad_size, std::vector<ByteInstruction>& byte_instructions) {
  byte_instructions.push_back({
      .operation = ByteInstruction::Operation::kInsert,
      .offset = c2pa_insertion_point,
      .bytes = "",
  });
  riegeli::StringWriter c2pa_writer(&byte_instructions.back().bytes);
  if (!c2pa_writer.Write(kRiffChunkIdC2pa)) {
    return absl::InternalError("failed to write C2PA chunk ID");
  }
  if (!riegeli::WriteLittleEndian<uint32_t>(manifest_store_size + pad_size,
                                            c2pa_writer)) {
    return absl::InternalError("failed to write C2PA chunk size");
  }
  if (pad_size > 0) {
    if (!c2pa_writer.Write(riegeli::ByteFill(pad_size, '\0'))) {
      return absl::InternalError("failed to write padding");
    }
  }
  if (!c2pa_writer.Close()) {
    return absl::InternalError("failed to close c2pa writer");
  }
  return absl::OkStatus();
}

absl::StatusOr<std::string> ComputeHash(
    riegeli::Reader& source, uint64_t file_size,
    absl::Span<const ByteInstruction> instructions) {
  if (!source.Seek(0) || source.pos() != 0) {
    return absl::InternalError("failed to seek to start of source");
  }

  ABSL_ASSIGN_OR_RETURN(std::unique_ptr<InputHasher> hasher,
                        InputHasher::Create(HashAlgorithm::kSha256));
  // Hash the RIFF header.
  ABSL_RETURN_IF_ERROR(hasher->Update(source, 0, instructions[0].offset));
  // Hash the updated RIFF size.
  hasher->Update(instructions[0].bytes);
  // Hash up to the C2PA chunk insertion point.
  ABSL_RETURN_IF_ERROR(hasher->Update(source, 8, instructions[1].offset - 8));

  // Hash the C2PA chunk placeholder.
  hasher->Update(instructions[1].bytes);

  // Hash the remainder of the file.
  ABSL_RETURN_IF_ERROR(hasher->Update(source, instructions[1].offset,
                                      file_size - instructions[1].offset));
  return hasher->Digest();
}

}  // namespace

absl::StatusOr<std::vector<ByteInstruction>> RiffWriter::WriteInternal(
    GenerateManifestStoreRequest request, riegeli::Reader* absl_nonnull source,
    OutputMetadata* absl_nullable output_metadata) const {
  if (request.has_hard_binding()) {
    return absl::InvalidArgumentError(
        "request must not include `hard_binding` (binding is computed by "
        "RiffWriter)");
  }
  uint64_t file_size = source->Size().value_or(0);
  if (file_size == 0) {
    return absl::InvalidArgumentError("source is empty");
  }

  // The two-phase approach is required for RIFF: The box hash can't be computed
  // correctly until the manifest store size is known, because the manifest size
  // affects the top-level RIFF chunk size.

  // Phase 1: Compute preliminary hard binding, produce a manifest store
  // template, and embed a zero-filled placeholder in the asset.
  ABSL_ASSIGN_OR_RETURN(*request.mutable_hard_binding()->mutable_data_hash(),
                        DataHashTemplate(HashAlgorithm::kSha256));
  ABSL_ASSIGN_OR_RETURN(ManifestStoreTemplate manifest_store_template,
                        generator_.GenerateManifestStoreTemplate(request));
  constexpr int64_t kMaxManifestStoreSizeBytes = 100 * 1024 * 1024;  // 100 MiB
  if (manifest_store_template.manifest_store_size() <= 0 ||
      manifest_store_template.manifest_store_size() >
          kMaxManifestStoreSizeBytes) {
    return absl::InvalidArgumentError(
        absl::StrCat("invalid manifest store template size: ",
                     manifest_store_template.manifest_store_size()));
  }

  // The insertion point is the end of the first RIFF chunk, or an error if the
  // input is not a RIFF file or if a C2PA chunk already exists.
  ABSL_ASSIGN_OR_RETURN(uint64_t c2pa_insertion_point,
                        FindC2paInsertionPoint(*source));

  uint64_t manifest_store_size = manifest_store_template.manifest_store_size();
  uint64_t pad_size = (manifest_store_size % 2 == 1) ? 1 : 0;

  uint64_t updated_riff_size =
      c2pa_insertion_point + manifest_store_size + pad_size;
  if (updated_riff_size > std::numeric_limits<uint32_t>::max()) {
    return absl::InvalidArgumentError("chunk size overflow");
  }

  std::vector<ByteInstruction> byte_instructions;
  ABSL_RETURN_IF_ERROR(
      AddRiffSizeInstruction(updated_riff_size, byte_instructions));
  ABSL_RETURN_IF_ERROR(AddC2paManifestInstructionPlaceholder(
      c2pa_insertion_point, manifest_store_size, pad_size, byte_instructions));

  ABSL_ASSIGN_OR_RETURN(std::string hash,
                        ComputeHash(*source, file_size, byte_instructions));

  FinalizeManifestStoreRequest finalize_request;
  *finalize_request.mutable_manifest_store_template() =
      std::move(manifest_store_template);
  *finalize_request.mutable_hard_binding() = std::move(request.hard_binding());
  finalize_request.mutable_hard_binding()
      ->mutable_data_hash()
      ->mutable_exclusions(0)
      ->set_start(c2pa_insertion_point + 8);
  finalize_request.mutable_hard_binding()
      ->mutable_data_hash()
      ->mutable_exclusions(0)
      ->set_length(manifest_store_size);
  finalize_request.mutable_hard_binding()->mutable_data_hash()->set_hash(
      std::move(hash));

  ABSL_ASSIGN_OR_RETURN(auto response,
                        generator_.FinalizeManifestStore(finalize_request));
  if (output_metadata != nullptr) {
    output_metadata->timestamp_included = response.is_timestamped();
  }
  if (response.manifest_store().size() !=
      finalize_request.manifest_store_template().manifest_store_size()) {
    return absl::InternalError("manifest store size does not match template");
  }

  // Update the C2PA chunk with the manifest store and padding.
  byte_instructions.back().bytes = absl::StrCat(
      byte_instructions.back().bytes.substr(0, 8), response.manifest_store(),
      pad_size > 0 ? std::string(pad_size, '\0') : "");

  return std::move(byte_instructions);
}

absl::Status RiffWriter::Write(
    GenerateManifestStoreRequest request, riegeli::Reader* absl_nonnull source,
    riegeli::Writer* absl_nonnull destination,
    OutputMetadata* absl_nullable output_metadata) const {
  ABSL_ASSIGN_OR_RETURN(
      std::vector<ByteInstruction> byte_instructions,
      WriteInternal(std::move(request), source, output_metadata));
  return ApplyByteInstructions(source, byte_instructions, destination);
}

absl::Status RiffWriter::Write(
    GenerateManifestStoreRequest request, riegeli::Writer* absl_nonnull asset,
    OutputMetadata* absl_nullable output_metadata) const {
  if (!asset->SupportsReadMode()) {
    return absl::InvalidArgumentError("asset does not support read mode");
  }

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

  ABSL_ASSIGN_OR_RETURN(
      std::vector<ByteInstruction> byte_instructions,
      WriteInternal(std::move(request), source, output_metadata));

  return ApplyByteInstructions(byte_instructions, asset);
}

}  // namespace credentio