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

#include <cstdint>
#include <optional>
#include <string>
#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/string_view.h"
#include "absl/types/span.h"
#include "formats/zip/constants.h"
#include "formats/zip/content_type_stream_util.h"
#include "formats/zip/crc.h"
#include "formats/zip/reader.h"
#include "riegeli/base/byte_fill.h"
#include "riegeli/base/types.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/writer.h"
#include "riegeli/endian/endian_writing.h"

namespace credentio {
namespace {

// A file entry that has been copied from the source to the destination.
struct FileCopied {
  ZipReader::FileEntry entry;
  riegeli::Position offset_in_destination;
};

// Copies a file entry from the source to the destination.
absl::Status CopyFileEntry(riegeli::Reader* absl_nonnull source,
                           const ZipReader::FileEntry& entry,
                           riegeli::Writer* absl_nonnull destination) {
  int64_t file_end = entry.file_range.offset + entry.file_range.length;
  int64_t data_descriptor_length = 0;
  if (entry.has_data_descriptor) {
    source->Seek(file_end);
    ABSL_RETURN_IF_ERROR(source->status());
    std::string buffer;
    source->Read(kZipDataDescriptorSignature.size(), buffer);
    ABSL_RETURN_IF_ERROR(source->status());
    if (buffer == kZipDataDescriptorSignature) {
      data_descriptor_length = kDataDescriptorWithSignatureSize;
    } else {
      data_descriptor_length = kDataDescriptorSize;
    }
  }
  int64_t length =
      file_end + data_descriptor_length - entry.local_file_header_offset;
  source->Seek(entry.local_file_header_offset);
  ABSL_RETURN_IF_ERROR(source->status());
  source->Copy(length, *destination);
  ABSL_RETURN_IF_ERROR(source->status());
  ABSL_RETURN_IF_ERROR(destination->status());

  return absl::OkStatus();
}

// Reads the content types stream from the source, writes the updated stream to
// the destination, including the local file header and the data descriptor (if
// used). Returns a `ContentTypesStreamMetadata` containing the updated metadata
// of the content types stream.
absl::StatusOr<ContentTypesStreamMetadata> WriteUpdatedContentTypesStream(
    riegeli::Reader* absl_nonnull source,
    const ZipReader::FileEntry& content_types_file_entry,
    riegeli::Writer* absl_nonnull destination) {
  if (content_types_file_entry.has_data_descriptor) {
    // Copy the original local file header since they don't need to be updated.
    source->Seek(content_types_file_entry.local_file_header_offset);
    ABSL_RETURN_IF_ERROR(source->status());
    source->Copy(kLocalFileHeaderFixedSize + kOoxmlContentTypesFileName.size() +
                     content_types_file_entry.extra_field_length,
                 *destination);
    ABSL_RETURN_IF_ERROR(source->status());
    ABSL_RETURN_IF_ERROR(destination->status());

    // Read and update the original content types stream from the source.
    ABSL_ASSIGN_OR_RETURN(
        ContentTypesStream content_types_stream,
        ReadContentTypesStream(source, content_types_file_entry));
    ABSL_RETURN_IF_ERROR(UpdateContentTypesStream(
        content_types_stream, kZipC2paManifestMimeType, kZipManifestFileName));
    return WriteContentTypesStream(content_types_stream,
                                   content_types_file_entry.has_data_descriptor,
                                   destination);
  } else {
    return absl::UnimplementedError(
        "Content types stream without a data descriptor is not supported.");
  }
}

// Writes the central directory header for the updated content types stream to
// the destination.
absl::Status WriteUpdatedContentTypesStreamCentralDirectoryHeader(
    riegeli::Reader* absl_nonnull source,
    const ZipReader::FileEntry& original_file_entry,
    const ContentTypesStreamMetadata& content_types_stream_metadata,
    int64_t updated_local_file_header_offset,
    riegeli::Writer* absl_nonnull destination) {
  // Copy the following fields from the original central directory header:
  //   - Central directory file header signature (4 bytes)
  //   - Version made by (2 bytes)
  //   - Version needed to extract (2 bytes)
  //   - General purpose bit flag (2 bytes)
  //   - Compression method (2 bytes)
  //   - File modification time (2 bytes)
  //   - File modification date (2 bytes)
  source->Seek(original_file_entry.central_directory_header_offset);
  ABSL_RETURN_IF_ERROR(source->status());
  source->Copy(16, *destination);
  ABSL_RETURN_IF_ERROR(source->status());
  ABSL_RETURN_IF_ERROR(destination->status());

  // CRC-32 (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(content_types_stream_metadata.crc32,
                                       *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Compressed size (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(
      content_types_stream_metadata.compressed_size, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Uncompressed size (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(
      content_types_stream_metadata.uncompressed_size, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Copy the remaining fields from the original central directory header:
  //   - File name length (2 bytes)
  //   - Extra field length (2 bytes)
  //   - File comment length (2 bytes)
  //   - Disk number start (2 bytes)
  //   - Internal file attributes (2 bytes)
  //   - External file attributes (4 bytes)
  source->Seek(original_file_entry.central_directory_header_offset + 28);
  ABSL_RETURN_IF_ERROR(source->status());
  source->Copy(14, *destination);
  ABSL_RETURN_IF_ERROR(source->status());
  ABSL_RETURN_IF_ERROR(destination->status());

  // Local file header offset (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(updated_local_file_header_offset,
                                       *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // File name
  destination->Write(kOoxmlContentTypesFileName);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Extra field and File comment data
  int64_t extra_and_comment_length = original_file_entry.extra_field_length +
                                     original_file_entry.file_comment_length;
  if (extra_and_comment_length > 0) {
    source->Seek(original_file_entry.central_directory_header_offset +
                 kCentralDirectoryHeaderFixedSize +
                 original_file_entry.file_name.size());
    ABSL_RETURN_IF_ERROR(source->status());
    source->Copy(extra_and_comment_length, *destination);
    ABSL_RETURN_IF_ERROR(source->status());
    ABSL_RETURN_IF_ERROR(destination->status());
  }

  return absl::OkStatus();
}

// Writes the local file header for the manifest store to the destination.
absl::Status WriteManifestLocalFileHeader(
    int64_t manifest_store_size, riegeli::Writer* absl_nonnull destination) {
  // Local file header signature (4 bytes)
  destination->Write(kZipLocalFileHeaderSignature);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Version needed to extract (2 bytes)
  riegeli::WriteLittleEndian<uint16_t>(10, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Zero-filled fields:
  //   - General purpose bit flag (2 bytes)
  //   - Compression method (2 bytes, 0 means no compression)
  //   - File modification time (2 bytes)
  //   - File modification date (2 bytes)
  //   - CRC-32 (4 bytes, placeholder)
  destination->Write(riegeli::ByteFill(12));
  ABSL_RETURN_IF_ERROR(destination->status());

  // Compressed size (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(manifest_store_size, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Uncompressed size (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(manifest_store_size, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // File name length (2 bytes, little-endian)
  riegeli::WriteLittleEndian<uint16_t>(kZipManifestFileName.size(),
                                       *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Zero-filled extra field length (2 bytes)
  destination->Write(riegeli::ByteFill(2));
  ABSL_RETURN_IF_ERROR(destination->status());

  // File name
  destination->Write(kZipManifestFileName);
  ABSL_RETURN_IF_ERROR(destination->status());

  return absl::OkStatus();
}

// Writes the central directory header for the manifest store to the
// destination.
absl::Status WriteManifestCentralDirectoryHeader(
    uint32_t manifest_store_size,
    uint32_t manifest_store_local_file_header_offset,
    riegeli::Writer* absl_nonnull destination) {
  // Central directory file header signature (4 bytes)
  destination->Write(kZipCentralDirectorySignature);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Version made by (2 bytes)
  riegeli::WriteLittleEndian<uint16_t>(20, *destination);  // Version 2.0
  ABSL_RETURN_IF_ERROR(destination->status());

  // Version needed to extract (2 bytes)
  riegeli::WriteLittleEndian<uint16_t>(10, *destination);  // Version 1.0
  ABSL_RETURN_IF_ERROR(destination->status());

  // Zero-filled fields:
  //   - General purpose bit flag (2 bytes)
  //   - Compression method (2 bytes, 0 means no compression)
  //   - File modification time (2 bytes)
  //   - File modification date (2 bytes)
  //   - CRC-32 (4 bytes, placeholder)
  destination->Write(riegeli::ByteFill(12));
  ABSL_RETURN_IF_ERROR(destination->status());

  // Compressed size (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(manifest_store_size, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Uncompressed size (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(manifest_store_size, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // File name length (2 bytes, little-endian)
  riegeli::WriteLittleEndian<uint16_t>(kZipManifestFileName.size(),
                                       *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Zero-filled fields:
  //   - Extra field length (2 bytes)
  //   - File comment length (2 bytes)
  //   - Disk number start (2 bytes)
  //   - Internal file attributes (2 bytes)
  //   - External file attributes (4 bytes)
  destination->Write(riegeli::ByteFill(12));
  ABSL_RETURN_IF_ERROR(destination->status());

  // Local file header offset (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(manifest_store_local_file_header_offset,
                                       *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // File name
  destination->Write(kZipManifestFileName);
  ABSL_RETURN_IF_ERROR(destination->status());

  return absl::OkStatus();
}

// Writes the central directory header for a copied file entry to the
// destination.
absl::Status WriteCentralDirectoryHeader(
    riegeli::Reader* absl_nonnull source, const FileCopied& file_copied,
    riegeli::Writer* absl_nonnull destination) {
  source->Seek(file_copied.entry.central_directory_header_offset);
  ABSL_RETURN_IF_ERROR(source->status());

  int64_t total_cd_length = kCentralDirectoryHeaderFixedSize +
                            file_copied.entry.file_name.size() +
                            file_copied.entry.extra_field_length +
                            file_copied.entry.file_comment_length;

  if (file_copied.offset_in_destination ==
      file_copied.entry.local_file_header_offset) {
    source->Copy(total_cd_length, *destination);
    ABSL_RETURN_IF_ERROR(source->status());
    ABSL_RETURN_IF_ERROR(destination->status());
  } else {
    source->Copy(kCentralDirectoryHeaderFixedSize - 4, *destination);
    ABSL_RETURN_IF_ERROR(source->status());
    ABSL_RETURN_IF_ERROR(destination->status());

    riegeli::WriteLittleEndian<uint32_t>(file_copied.offset_in_destination,
                                         *destination);
    ABSL_RETURN_IF_ERROR(destination->status());

    source->Skip(4);
    ABSL_RETURN_IF_ERROR(source->status());

    int64_t remaining_length = file_copied.entry.file_name.size() +
                               file_copied.entry.extra_field_length +
                               file_copied.entry.file_comment_length;
    source->Copy(remaining_length, *destination);
    ABSL_RETURN_IF_ERROR(source->status());
    ABSL_RETURN_IF_ERROR(destination->status());
  }
  return absl::OkStatus();
}

// Updates the End of Central Directory (EOCD) record with the new central
// directory offset, length and entry count.
absl::Status UpdateEndOfCentralDirectory(
    riegeli::Reader* absl_nonnull source, uint32_t central_directory_offset,
    uint32_t central_directory_length, uint16_t entry_count,
    riegeli::Writer* absl_nonnull destination) {
  // The first 8 bytes of the EOCD don't need to be updated.
  auto source_size = source->Size();
  if (!source_size.has_value()) return source->status();
  source->Seek(*source_size - kEocdRecordSize);
  ABSL_RETURN_IF_ERROR(source->status());
  source->Copy(8, *destination);
  ABSL_RETURN_IF_ERROR(source->status());
  ABSL_RETURN_IF_ERROR(destination->status());

  // Total number of entries in the central directory on this disk (2 bytes,
  // little-endian)
  riegeli::WriteLittleEndian<uint16_t>(entry_count, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Total number of entries in the central directory (2 bytes, little-endian)
  riegeli::WriteLittleEndian<uint16_t>(entry_count, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Size of the central directory (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(central_directory_length, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Offset of the start of the central directory (4 bytes, little-endian)
  riegeli::WriteLittleEndian<uint32_t>(central_directory_offset, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  // Zero-filled comment length (2 bytes)
  destination->Write(riegeli::ByteFill(2));
  ABSL_RETURN_IF_ERROR(destination->status());

  return absl::OkStatus();
}

}  // namespace

absl::StatusOr<ZipEmbedder::CopyAndAllocateResult> ZipEmbedder::CopyAndAllocate(
    riegeli::Reader* absl_nonnull source,
    absl::Span<const ZipReader::FileEntry> file_entries,
    int64_t manifest_store_size,
    riegeli::Writer* absl_nonnull destination) const {
  ABSL_ASSIGN_OR_RETURN(std::vector<ZipReader::FileEntry> sorted_entries,
                        ZipReader::SortFileEntries(file_entries));

  // Step 1: Copy all the files in the ZIP archive from the source to the
  // destination, except the content types stream.
  std::vector<FileCopied> files_copied;
  std::optional<ZipReader::FileEntry> content_types_file_entry;
  for (const auto& entry : sorted_entries) {
    if (entry.file_name == kOoxmlContentTypesFileName) {
      content_types_file_entry = entry;
      continue;
    }
    if (entry.file_name == kZipManifestFileName) {
      return absl::AlreadyExistsError("Manifest store already exists.");
    }
    files_copied.push_back(
        {.entry = entry, .offset_in_destination = destination->pos()});
    ABSL_RETURN_IF_ERROR(CopyFileEntry(source, entry, destination));
  }

  // Step 2: Write the updated content types stream to the destination.
  if (!content_types_file_entry.has_value()) {
    return absl::NotFoundError("Content types stream not found.");
  }
  auto content_types_stream_output_offset = destination->pos();
  ABSL_ASSIGN_OR_RETURN(ContentTypesStreamMetadata content_types_metadata,
                        WriteUpdatedContentTypesStream(
                            source, *content_types_file_entry, destination));

  // Step 3: Write the zero-filled manifest store to the destination.
  auto manifest_store_local_file_header_offset = destination->pos();
  ABSL_RETURN_IF_ERROR(
      WriteManifestLocalFileHeader(manifest_store_size, destination));
  auto manifest_store_offset = destination->pos();
  destination->Write(riegeli::ByteFill(manifest_store_size));
  ABSL_RETURN_IF_ERROR(destination->status());

  // Step 4: Write the central directory headers.
  auto central_directory_offset = destination->pos();
  for (const auto& file : files_copied) {
    ABSL_RETURN_IF_ERROR(
        WriteCentralDirectoryHeader(source, file, destination));
  }
  ABSL_RETURN_IF_ERROR(WriteUpdatedContentTypesStreamCentralDirectoryHeader(
      source, *content_types_file_entry, content_types_metadata,
      content_types_stream_output_offset, destination));
  auto manifest_store_central_directory_header_offset = destination->pos();
  ABSL_RETURN_IF_ERROR(WriteManifestCentralDirectoryHeader(
      manifest_store_size, manifest_store_local_file_header_offset,
      destination));

  // Step 5: Update the End of Central Directory (EOCD).
  ABSL_RETURN_IF_ERROR(UpdateEndOfCentralDirectory(
      source, central_directory_offset,
      /*central_directory_length=*/destination->pos() -
          central_directory_offset,
      files_copied.size() + 2, destination));

  destination->Flush();
  ABSL_RETURN_IF_ERROR(destination->status());

  return CopyAndAllocateResult{
      .manifest_store_offset = manifest_store_offset,
      .manifest_store_central_directory_header_offset =
          manifest_store_central_directory_header_offset,
      .central_directory_offset = central_directory_offset,
  };
}

absl::Status ZipEmbedder::Embed(
    absl::string_view manifest_store,
    const CopyAndAllocateResult& copy_and_allocate_result,
    riegeli::Writer* absl_nonnull destination) const {
  if (!destination->SupportsRandomAccess()) {
    return absl::InvalidArgumentError(
        "ZipEmbedder::Embed requires a riegeli::Writer that supports random "
        "access.");
  }
  uint32_t crc32 = ZipFileCrc(manifest_store);

  int64_t local_file_header_crc_offset =
      copy_and_allocate_result.manifest_store_offset -
      kZipManifestFileName.size() - 16;
  int64_t central_directory_header_crc_offset =
      copy_and_allocate_result.manifest_store_central_directory_header_offset +
      16;

  // Replace the zero-filled manifest store placeholder with the actual manifest
  // store. Also update the CRC-32 fields in both the local file header and the
  // central directory header with the actual CRC-32 of the manifest store.
  destination->Seek(local_file_header_crc_offset);
  ABSL_RETURN_IF_ERROR(destination->status());
  riegeli::WriteLittleEndian<uint32_t>(crc32, *destination);
  ABSL_RETURN_IF_ERROR(destination->status());

  destination->Seek(copy_and_allocate_result.manifest_store_offset);
  ABSL_RETURN_IF_ERROR(destination->status());
  destination->Write(manifest_store);
  ABSL_RETURN_IF_ERROR(destination->status());

  destination->Seek(central_directory_header_crc_offset);
  ABSL_RETURN_IF_ERROR(destination->status());
  riegeli::WriteLittleEndian<uint32_t>(crc32, *destination);
  return destination->status();
}

}  // namespace credentio