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

#include <cstdint>
#include <iterator>
#include <optional>
#include <string>
#include <tuple>
#include <utility>
#include <variant>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/container/btree_set.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "absl/types/span.h"
#include "formats/pdf/constants.h"
#include "formats/pdf/extractor.h"
#include "formats/pdf/object_writer.h"
#include "formats/pdf/objects.h"
#include "formats/pdf/reader.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/writer.h"

namespace credentio {
namespace {

using Dictionary = Object::Dictionary;

constexpr absl::string_view kPdfHeader = "%PDF-1.7\n%\x90\x91\xC2\x9A\n";

struct XrefTableEntry {
  uint32_t object_number;
  uint32_t generation_number;

  int64_t offset;
  bool in_use;

  friend bool operator<(const XrefTableEntry& lhs, const XrefTableEntry& rhs) {
    return std::tie(lhs.object_number, lhs.generation_number, lhs.offset,
                    lhs.in_use) < std::tie(rhs.object_number,
                                           rhs.generation_number, rhs.offset,
                                           rhs.in_use);
  }
};

template <typename T>
absl::StatusOr<T> DereferenceObject(const Object::IndirectReference& ref,
                                    PdfReader& reader) {
  ABSL_ASSIGN_OR_RETURN(
      const auto& indirect_obj,
      reader.GetObject(ref.object_number, ref.generation_number));
  return Object::As<T>(indirect_obj.object);
}

// Using the basic mechanism to form a linked list from the free entries.
// See ISO 32000-2 Section 7.5.4 for more details.
std::vector<XrefTableEntry> ReconstructXrefTableFreeEntries(
    const absl::btree_set<Object::IndirectReference>& free_entries) {
  if (free_entries.empty()) {
    return {};
  }
  std::vector<XrefTableEntry> xref_table_entries;
  xref_table_entries.reserve(free_entries.size());

  for (const auto& free_entry : free_entries) {
    xref_table_entries.push_back(
        XrefTableEntry{.object_number = free_entry.object_number,
                       .generation_number = free_entry.generation_number,
                       .offset = 0,
                       .in_use = false});
  }

  // Set the offset of each free entry to the object number of the next free
  // entry.
  for (int i = 0; i < xref_table_entries.size() - 1; ++i) {
    xref_table_entries[i].offset = xref_table_entries[i + 1].object_number;
  }

  // Signify the end of the list by having an offset value of 0.
  xref_table_entries.back().offset = 0;

  return xref_table_entries;
}

// Returns true if the object numbers in the given vector are consecutive,
// starting from 0. Assumes the vector is already sorted by object number.
bool AreObjectNumbersConsecutive(
    absl::Span<const XrefTableEntry> xref_table_entries) {
  for (auto i = 0; i < xref_table_entries.size(); i++) {
    if (xref_table_entries[i].object_number != i) {
      return false;
    }
  }
  return true;
}

absl::Status WritePdfHeader(riegeli::Writer& output) {
  if (!output.Write(kPdfHeader)) {
    return output.status();
  }
  return absl::OkStatus();
}

// See ISO 32000-2 Section 7.5.4 for the cross-reference table structure.
absl::Status WriteXrefTable(std::vector<XrefTableEntry>& xref_table_entries,
                            riegeli::Writer& output) {
  // Write the xref header.
  if (!output.Write("xref\n")) {
    return output.status();
  }

  absl::c_sort(xref_table_entries);
  if (!AreObjectNumbersConsecutive(xref_table_entries)) {
    return absl::UnimplementedError(
        "Non-consecutive object numbers in the xref table are not supported.");
  }

  // Write the object count section.
  if (!output.Write(absl::StrCat(0, " ", xref_table_entries.size(), "\n"))) {
    return output.status();
  }

  // Write the xref table entries.
  for (const auto& entry : xref_table_entries) {
    if (!output.Write(absl::StrFormat("%010d %05d %c \n", entry.offset,
                                      entry.generation_number,
                                      entry.in_use ? 'n' : 'f'))) {
      return output.status();
    }
  }
  return absl::OkStatus();
}

// See ISO 32000-2 Section 7.5.5 for the trailer structure.
absl::Status WriteTrailer(const Dictionary& trailer_dict,
                          int64_t xref_table_offset, riegeli::Writer& output) {
  // Write the trailer header.
  if (!output.Write("trailer\n")) {
    return output.status();
  }

  // Write the trailer dictionary object.
  Object obj = Object{.value = trailer_dict};
  ABSL_RETURN_IF_ERROR(WriteObject(obj, output));

  // Write the trailer footer and the offset to the last cross-reference table.
  if (!output.Write(
          absl::StrCat("\nstartxref\n", xref_table_offset, "\n%%EOF"))) {
    return output.status();
  }

  return absl::OkStatus();
}

absl::Status CopyObjects(
    PdfReader& reader,
    absl::Span<const std::optional<Object::IndirectReference>> objects_to_skip,
    int64_t start_offset, std::vector<XrefTableEntry>& xref_table_entries,
    riegeli::Writer& output) {
  while (reader.HasNext()) {
    ABSL_ASSIGN_OR_RETURN(auto obj, reader.Next());

    bool should_skip = absl::c_any_of(
        objects_to_skip,
        [&](const std::optional<Object::IndirectReference>& skip_ref) {
          return skip_ref.has_value() &&
                 obj.object_number == skip_ref->object_number &&
                 obj.generation_number == skip_ref->generation_number;
        });

    if (should_skip) continue;
    int64_t obj_offset = output.pos();
    ABSL_RETURN_IF_ERROR(WriteIndirectObject(obj, output));
    xref_table_entries.push_back(
        XrefTableEntry{.object_number = obj.object_number,
                       .generation_number = obj.generation_number,
                       .offset = obj_offset - start_offset,
                       .in_use = true});
  }
  return absl::OkStatus();
}

// Creates an indirect stream Object with a zero-filled manifest store.
IndirectObject BuildC2paManifestTemplate(uint32_t object_number,
                                         uint32_t generation_number,
                                         int64_t manifest_store_size) {
  IndirectObject store_template;
  store_template.object_number = object_number;
  store_template.generation_number = generation_number;

  std::string formatted_date = absl::FormatTime(
      "D:%Y%m%d%H%M%S+00'00'", absl::Now(), absl::UTCTimeZone());
  Dictionary params_dict{
      .entries = {
          {std::string(kPdfSizeKey),
           Object{.value = Object::Integer{.value = manifest_store_size}}},
          {std::string(kPdfModDateKey),
           Object{.value = Object::LiteralString{.raw_value = formatted_date}}},
      }};

  store_template.stream_dictionary = Dictionary{
      .entries = {
          {"Length",
           Object{.value = Object::Integer{.value = manifest_store_size}}},
          {std::string(kPdfSubtypeKey),
           Object{.value = Object::Name{.value = std::string(
                                            kPdfC2paManifestMimeType)}}},
          {std::string(kPdfParamsKey), Object{.value = std::move(params_dict)}},
      }};
  store_template.object =
      Object{.value = Object::Stream{
                 .raw_value = std::string(manifest_store_size, '\0')}};
  return store_template;
}

IndirectObject BuildC2paManifestFileSpec(
    uint32_t object_number, uint32_t generation_number,
    const IndirectObject& c2pa_manifest_store_obj) {
  IndirectObject file_spec;
  file_spec.object_number = object_number;
  file_spec.generation_number = generation_number;
  file_spec.object = Object{
      .value = Dictionary{
          .entries = {
              {std::string(kPdfTypeKey),
               Object{.value = Object::Name{.value = std::string(
                                                kPdfFileSpecType)}}},
              {std::string(kPdfAfRelationshipKey),
               Object{.value =
                          Object::Name{.value = std::string(
                                           kPdfC2paManifestAfRelationship)}}},
              {std::string(kPdfFKey),
               Object{.value =
                          Object::LiteralString{
                              .raw_value =
                                  std::string(kPdfC2paEmbeddedFileName)}}},
              {std::string(kPdfUfKey),
               Object{.value =
                          Object::LiteralString{
                              .raw_value =
                                  std::string(kPdfC2paEmbeddedFileName)}}},
              {std::string(kPdfEfKey),
               Object{
                   .value = Dictionary{
                       .entries = {
                           {std::string(kPdfFKey),
                            Object{.value = Object::IndirectReference{
                                       .object_number = c2pa_manifest_store_obj
                                                            .object_number,
                                       .generation_number =
                                           c2pa_manifest_store_obj
                                               .generation_number}}}}}}}}}};
  return file_spec;
}

absl::Status InsertNameTreePair(absl::string_view key,
                                const Object::IndirectReference& value,
                                Object::Array& name_tree_array) {
  if (name_tree_array.objects.size() % 2 != 0) {
    return absl::InternalError(
        "Name tree array must contain an even number of objects.");
  }

  // Find the correct position to insert the new key-value pair to maintain
  // lexicographical order of keys. The `name_tree_array` stores keys and
  // values interleaved, e.g., [key1, value1, key2, value2, ...].
  auto insertion_it = name_tree_array.objects.begin();
  while (insertion_it != name_tree_array.objects.end()) {
    // We only need to check the keys, which are at even indices.
    ABSL_ASSIGN_OR_RETURN(const auto& current_key,
                          Object::As<Object::LiteralString>(*insertion_it));
    if (key < current_key.raw_value) {
      // Found the position. The new key should be inserted before this one.
      break;
    }
    // Advance the iterator to the next key.
    insertion_it += 2;
  }

  name_tree_array.objects.insert(
      insertion_it,
      {Object{.value = Object::LiteralString{.raw_value = std::string(key)}},
       Object{.value = value}});

  return absl::OkStatus();
}

}  // namespace

absl::StatusOr<int64_t> PdfEmbedder::RewriteForEmbedding(
    riegeli::Reader& input, int64_t estimated_manifest_store_size,
    riegeli::Writer& output) const {
  ABSL_ASSIGN_OR_RETURN(auto reader, PdfReader::Create(&input));

  if (auto manifest_store = PdfExtractor().ExtractManifestStore(input);
      manifest_store.ok()) {
    return absl::AlreadyExistsError(
        "Rewriting failed because a C2PA manifest store already exists ");
  }

  // The start offset of the output stream.
  int64_t start_offset = output.pos();

  // Step 1:
  // Locate and identify objects that need to be updated/created for C2PA
  // embedding.
  auto file_trailer_dict = reader->file_trailer_dictionary();
  ABSL_ASSIGN_OR_RETURN(
      auto object_count_obj,
      FindObjectInDictionary<Object::Integer>(file_trailer_dict, kPdfSizeKey));
  uint32_t next_available_object_number = object_count_obj.value;

  // /Root
  ABSL_ASSIGN_OR_RETURN(auto doc_catalog_ref,
                        FindObjectInDictionary<Object::IndirectReference>(
                            file_trailer_dict, kPdfRootKey));
  ABSL_ASSIGN_OR_RETURN(auto doc_catalog_dict, DereferenceObject<Dictionary>(
                                                   doc_catalog_ref, *reader));

  // /Root/Names
  std::optional<Dictionary> catalog_names = std::nullopt;
  std::optional<Object::IndirectReference> catalog_names_ref = std::nullopt;
  if (auto it = doc_catalog_dict.entries.find(kPdfNamesKey);
      it != doc_catalog_dict.entries.end()) {
    if (std::holds_alternative<Dictionary>(it->second.value)) {
      // If we find the catalog names dictionary to be a direct object, move it
      // into an indirect object then update the document catalog dictionary to
      // point to it.
      catalog_names = std::move(std::get<Dictionary>(it->second.value));
      catalog_names_ref = Object::IndirectReference{
          .object_number = next_available_object_number++,
          .generation_number = 0};
      doc_catalog_dict.entries[kPdfNamesKey] =
          Object{.value = *catalog_names_ref};
    } else {
      ABSL_ASSIGN_OR_RETURN(catalog_names_ref,
                            Object::As<Object::IndirectReference>(it->second));
      ABSL_ASSIGN_OR_RETURN(
          catalog_names,
          DereferenceObject<Dictionary>(catalog_names_ref.value(), *reader));
    }
  }

  // /Root/Names/EmbeddedFiles
  std::optional<Dictionary> name_tree = std::nullopt;
  std::optional<Object::IndirectReference> name_tree_ref = std::nullopt;
  if (catalog_names.has_value()) {
    if (auto it = catalog_names->entries.find(kPdfEmbeddedFilesKey);
        it != catalog_names->entries.end()) {
      ABSL_ASSIGN_OR_RETURN(name_tree_ref,
                            Object::As<Object::IndirectReference>(it->second));
      ABSL_ASSIGN_OR_RETURN(name_tree, DereferenceObject<Dictionary>(
                                           name_tree_ref.value(), *reader));
    }
  }

  // Step 2: Write the header.
  ABSL_RETURN_IF_ERROR(WritePdfHeader(output));

  // Step 3: Copy the objects that can be written as-is from the input stream to
  // the output stream.

  // A list of xref table entries that will be written to the output stream,
  // including both in-use and free entries. In-use entry's offset is relative
  // to `start_offset`.
  std::vector<XrefTableEntry> xref_table_entries;
  ABSL_RETURN_IF_ERROR(CopyObjects(
      *reader,
      /*objects_to_skip=*/{doc_catalog_ref, catalog_names_ref, name_tree_ref},
      start_offset, xref_table_entries, output));

  // Step 4.1: Create the embedded file stream object that holds the C2PA
  // manifest store.
  IndirectObject c2pa_manifest_store_obj = BuildC2paManifestTemplate(
      /*object_number=*/next_available_object_number++,
      /*generation_number=*/0, estimated_manifest_store_size);

  int64_t c2pa_manifest_store_offset = output.pos();
  ABSL_RETURN_IF_ERROR(WriteIndirectObject(c2pa_manifest_store_obj, output));
  xref_table_entries.push_back(XrefTableEntry{
      .object_number = c2pa_manifest_store_obj.object_number,
      .generation_number = c2pa_manifest_store_obj.generation_number,
      .offset = c2pa_manifest_store_offset - start_offset,
      .in_use = true});

  // Step 4.2: Create a file specification object for the C2PA manifest
  // store.
  IndirectObject c2pa_manifest_store_file_spec = BuildC2paManifestFileSpec(
      /*object_number=*/next_available_object_number++,
      /*generation_number=*/0, c2pa_manifest_store_obj);

  int64_t c2pa_manifest_store_file_spec_offset = output.pos();
  ABSL_RETURN_IF_ERROR(
      WriteIndirectObject(c2pa_manifest_store_file_spec, output));
  xref_table_entries.push_back(XrefTableEntry{
      .object_number = c2pa_manifest_store_file_spec.object_number,
      .generation_number = c2pa_manifest_store_file_spec.generation_number,
      .offset = c2pa_manifest_store_file_spec_offset - start_offset,
      .in_use = true});

  // Step 4.3: Add/Update the embedded files name tree.
  if (!name_tree.has_value()) {
    // Since we don't find an existing name tree in the input PDF, we can safely
    // create a new one.
    name_tree = Dictionary{};
    name_tree_ref = Object::IndirectReference{
        .object_number = next_available_object_number++,
        .generation_number = 0};
  }
  auto name_tree_it = name_tree->entries.find(kPdfNamesKey);
  if (name_tree_it == name_tree->entries.end() ||
      !std::holds_alternative<Object::Array>(name_tree_it->second.value)) {
    name_tree_it =
        name_tree->entries
            .insert_or_assign(kPdfNamesKey, Object{.value = Object::Array{}})
            .first;
  }
  Object::Array& name_tree_array =
      std::get<Object::Array>(name_tree_it->second.value);
  ABSL_RETURN_IF_ERROR(InsertNameTreePair(
      kPdfC2paEmbeddedFileName,
      Object::IndirectReference{
          .object_number = c2pa_manifest_store_file_spec.object_number,
          .generation_number = c2pa_manifest_store_file_spec.generation_number},
      name_tree_array));

  int64_t name_tree_offset = output.pos();
  ABSL_RETURN_IF_ERROR(WriteIndirectObject(
      IndirectObject{.object_number = name_tree_ref->object_number,
                     .generation_number = name_tree_ref->generation_number,
                     .object = Object{.value = *name_tree}},
      output));
  xref_table_entries.push_back(
      XrefTableEntry{.object_number = name_tree_ref->object_number,
                     .generation_number = name_tree_ref->generation_number,
                     .offset = name_tree_offset - start_offset,
                     .in_use = true});

  // Step 4.4: Add/Update the catalog names dictionary.
  if (!catalog_names.has_value()) {
    // Since we don't find an existing catalog names dictionary in the input
    // PDF, we can safely create a new one.
    catalog_names = Dictionary{};
    catalog_names_ref = Object::IndirectReference{
        .object_number = next_available_object_number++,
        .generation_number = 0};
    doc_catalog_dict.entries[kPdfNamesKey] =
        Object{.value = *catalog_names_ref};
  }
  if (!catalog_names->entries.contains(kPdfEmbeddedFilesKey)) {
    catalog_names->entries[kPdfEmbeddedFilesKey] =
        Object{.value = *name_tree_ref};
  }

  int64_t catalog_names_offset = output.pos();
  ABSL_RETURN_IF_ERROR(WriteIndirectObject(
      IndirectObject{.object_number = catalog_names_ref->object_number,
                     .generation_number = catalog_names_ref->generation_number,
                     .object = Object{.value = *catalog_names}},
      output));
  xref_table_entries.push_back(
      XrefTableEntry{.object_number = catalog_names_ref->object_number,
                     .generation_number = catalog_names_ref->generation_number,
                     .offset = catalog_names_offset - start_offset,
                     .in_use = true});

  // Step 4.5: Update the document catalog dictionary.
  auto af_it = doc_catalog_dict.entries.find(kPdfAfKey);
  if (af_it == doc_catalog_dict.entries.end() ||
      !std::holds_alternative<Object::Array>(af_it->second.value)) {
    af_it = doc_catalog_dict.entries
                .insert_or_assign(kPdfAfKey, Object{.value = Object::Array{}})
                .first;
  }
  std::get<Object::Array>(af_it->second.value)
      .objects.push_back(Object{
          .value = Object::IndirectReference{
              .object_number = c2pa_manifest_store_file_spec.object_number,
              .generation_number =
                  c2pa_manifest_store_file_spec.generation_number}});
  int64_t doc_catalog_dict_offset = output.pos();
  ABSL_RETURN_IF_ERROR(WriteIndirectObject(
      IndirectObject{.object_number = doc_catalog_ref.object_number,
                     .generation_number = doc_catalog_ref.generation_number,
                     .object = Object{.value = doc_catalog_dict}},
      output));
  xref_table_entries.push_back(
      XrefTableEntry{.object_number = doc_catalog_ref.object_number,
                     .generation_number = doc_catalog_ref.generation_number,
                     .offset = doc_catalog_dict_offset - start_offset,
                     .in_use = true});

  // Step 5: Compute and write the xref table.
  int64_t xref_table_offset = output.pos() - start_offset;

  // Add all existing free entries to the xref table.
  auto xref_table_free_entries =
      ReconstructXrefTableFreeEntries(reader->free_entries());
  xref_table_entries.insert(
      xref_table_entries.end(),
      std::make_move_iterator(xref_table_free_entries.begin()),
      std::make_move_iterator(xref_table_free_entries.end()));

  ABSL_RETURN_IF_ERROR(WriteXrefTable(xref_table_entries, output));

  // Step 6: Write the trailer section, including the trailer dictionary and
  // startxref pointer.

  // Update the object count in the trailer dictionary.
  file_trailer_dict.entries[kPdfSizeKey] =
      Object{.value = Object::Integer{
                 .value = static_cast<int64_t>(xref_table_entries.size())}};
  ABSL_RETURN_IF_ERROR(
      WriteTrailer(file_trailer_dict, xref_table_offset, output));

  ABSL_ASSIGN_OR_RETURN(int64_t header_len, ComputeIndirectStreamHeaderByteSize(
                                                c2pa_manifest_store_obj));
  return c2pa_manifest_store_offset + header_len;
}

absl::Status PdfEmbedder::Embed(absl::string_view manifest_store,
                                int64_t insert_at_offset,
                                absl::Cord& embeddable_asset) const {
  absl::Cord result = embeddable_asset.Subcord(0, insert_at_offset);
  result.Append(manifest_store);
  result.Append(embeddable_asset.Subcord(
      insert_at_offset + manifest_store.size(), embeddable_asset.size()));
  embeddable_asset = std::move(result);
  return absl::OkStatus();
}
}  // namespace credentio