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

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

#include "absl/base/nullability.h"
#include "absl/container/btree_map.h"
#include "absl/container/btree_set.h"
#include "absl/container/flat_hash_set.h"
#include "absl/memory/memory.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/substitute.h"
#include "formats/pdf/object_reader.h"
#include "formats/pdf/objects.h"
#include "re2/re2.h"
#include "riegeli/bytes/reader.h"

namespace credentio {
namespace {

// Matches the cross-reference table offset in the trailer. See ISO 32000-2
// Section 7.5.5
static constexpr LazyRE2 kXrefOffsetPattern = {R"(startxref\s(\d+)\s%%EOF)"};

// Matches a single cross-reference entry in the cross-reference table. See ISO
// 32000-2 Section 7.5.4
static constexpr LazyRE2 kXrefEntryPattern = {
    R"(^(\d{10}) (\d{5}) ([nf])(?: \r| \n|\r\n))"};

// Reads the first cross-reference table offset from the file trailer.
absl::StatusOr<uint64_t> GetFirstXrefTableOffset(riegeli::Reader& input) {
  // Scan the last 30 bytes for the cross-reference table offset. This allows
  // for the "startxref" keyword and an offset value large enough for a
  // multi-terabyte file.
  if (!input.Seek(input.Size().value_or(30) - 30)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to seek to offset"));
  }

  std::string buffer;
  if (!input.Read(30, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  uint64_t offset;
  if (!RE2::PartialMatch(buffer, *kXrefOffsetPattern, &offset)) {
    return absl::InvalidArgumentError(
        absl::StrCat("Failed to find cross-reference table offset in the "
                     "following payload: ",
                     buffer));
  }
  return offset;
}

// Reads the offset to the next cross-reference table from the file trailer.
// Returns `std::nullopt` if the file trailer dictionary does not contain a
// "Prev" entry.
absl::StatusOr<std::optional<uint64_t>> GetPrevXrefTableOffset(
    const Object::Dictionary& file_trailer_dict) {
  if (const auto it = file_trailer_dict.entries.find("Prev");
      it != file_trailer_dict.entries.end()) {
    ABSL_ASSIGN_OR_RETURN(const auto prev_xref_offset,
                          Object::As<Object::Integer>(it->second));
    return prev_xref_offset.value;
  }
  return std::nullopt;
}

// Reads the standard cross-reference table and the file trailer from the input
// stream at the given offset.
absl::Status ReadXrefTableAndTrailer(
    riegeli::Reader& input, uint64_t xref_offset,
    absl::btree_map<Object::IndirectReference, uint64_t>& xref_table,
    absl::btree_set<Object::IndirectReference>& free_entries,
    Object::Dictionary& file_trailer_dict) {
  if (!input.Seek(xref_offset)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to seek to offset"));
  }

  std::string buffer;
  if (!input.Read(4, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  if (buffer != "xref") {
    return absl::InvalidArgumentError(
        absl::StrCat("Malformed cross-reference table. Expected 'xref' but got "
                     "'",
                     buffer, "'"));
  }
  while (true) {
    ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
    if (!input.Read(7, buffer)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read buffer"));
    }
    if (buffer == "trailer") {
      break;
    }
    if (!input.Seek(input.pos() - 7)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to seek to offset"));
    }

    Object int_obj;
    ABSL_RETURN_IF_ERROR(ReadObject(input, input.pos(), int_obj));
    ABSL_ASSIGN_OR_RETURN(const auto obj_num,
                          Object::As<Object::Integer>(int_obj));

    ABSL_RETURN_IF_ERROR(ReadObject(input, input.pos(), int_obj));
    ABSL_ASSIGN_OR_RETURN(const auto obj_count,
                          Object::As<Object::Integer>(int_obj));

    ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));

    for (uint32_t i = 0; i < obj_count.value; ++i) {
      // Each cross-reference entry is expected to be exactly 20 bytes long.
      if (!input.Read(20, buffer)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to read buffer"));
      }
      uint64_t obj_offset;
      uint32_t gen_num;
      std::string type;
      if (RE2::PartialMatch(buffer, *kXrefEntryPattern, &obj_offset, &gen_num,
                            &type)) {
        auto xref_entry_key = Object::IndirectReference{
            .object_number = static_cast<uint32_t>(obj_num.value + i),
            .generation_number = gen_num};
        if (type == "n") {  // 'n' represents an in-use entry.
          xref_table.insert({xref_entry_key, obj_offset});
        } else {  // 'f' represents a free entry.
          free_entries.insert(xref_entry_key);
        }
      } else {
        return absl::InvalidArgumentError(
            absl::StrCat("Failed to parse cross-reference table entry. Invalid "
                         "value: ",
                         buffer));
      }
    }
  }
  Object dict_obj;
  ABSL_RETURN_IF_ERROR(ReadObject(input, input.pos(), dict_obj));
  ABSL_ASSIGN_OR_RETURN(file_trailer_dict,
                        Object::As<Object::Dictionary>(dict_obj));

  return absl::OkStatus();
}
}  // namespace

absl::StatusOr<std::unique_ptr<PdfReader>> PdfReader::Create(
    riegeli::Reader* absl_nonnull input) {
  if (input == nullptr) {
    return absl::InvalidArgumentError("input cannot be null");
  }
  absl::flat_hash_set<uint64_t> xref_offsets_visited;
  ABSL_ASSIGN_OR_RETURN(uint64_t xref_offset, GetFirstXrefTableOffset(*input));
  xref_offsets_visited.insert(xref_offset);

  absl::btree_map<Object::IndirectReference, uint64_t> xref_table;
  absl::btree_set<Object::IndirectReference> free_entries;
  Object::Dictionary file_trailer_dict;
  ABSL_RETURN_IF_ERROR(ReadXrefTableAndTrailer(
      *input, xref_offset, xref_table, free_entries, file_trailer_dict));

  ABSL_ASSIGN_OR_RETURN(std::optional<uint64_t> prev_xref_offset,
                        GetPrevXrefTableOffset(file_trailer_dict));
  while (prev_xref_offset.has_value()) {
    if (xref_offsets_visited.contains(*prev_xref_offset)) {
      return absl::InvalidArgumentError(
          "Initialization failed: Circular cross-reference table detected.");
    }
    xref_offsets_visited.insert(*prev_xref_offset);

    absl::btree_map<Object::IndirectReference, uint64_t> prev_xref_table;
    absl::btree_set<Object::IndirectReference> prev_free_entries;
    Object::Dictionary prev_file_trailer_dict;
    ABSL_RETURN_IF_ERROR(
        ReadXrefTableAndTrailer(*input, *prev_xref_offset, prev_xref_table,
                                prev_free_entries, prev_file_trailer_dict));
    ABSL_ASSIGN_OR_RETURN(prev_xref_offset,
                          GetPrevXrefTableOffset(prev_file_trailer_dict));

    xref_table.merge(prev_xref_table);
    if (!prev_xref_table.empty()) {
      return absl::InvalidArgumentError(
          "Initialization failed: Collision detected in cross-reference "
          "table for in-use entries.");
    }
    free_entries.merge(prev_free_entries);
    if (!prev_free_entries.empty()) {
      return absl::InvalidArgumentError(
          "Initialization failed: Collision detected in cross-reference "
          "table for free entries.");
    }
  }
  auto reader = absl::WrapUnique(new PdfReader(input, std::move(xref_table),
                                               std::move(free_entries),
                                               std::move(file_trailer_dict)));
  return reader;
}

bool PdfReader::HasNext() { return current_it_ != xref_table_.end(); }

absl::StatusOr<IndirectObject> PdfReader::Next() {
  if (!HasNext()) {
    return absl::OutOfRangeError("No more objects to read.");
  }
  IndirectObject obj;
  if (absl::Status status = ReadIndirectObject(
          input_, current_it_->second,
          [this](const Object::IndirectReference& ref) {
            return this->GetObject(ref.object_number, ref.generation_number);
          },
          obj);
      !status.ok()) {
    return status;
  }
  ++current_it_;
  return obj;
}

absl::StatusOr<uint64_t> PdfReader::GetObjectOffset(
    uint32_t object_number, uint32_t generation_number) {
  auto it = xref_table_.find(Object::IndirectReference{
      .object_number = object_number, .generation_number = generation_number});
  if (it == xref_table_.end()) {
    return absl::NotFoundError(absl::Substitute(
        "No offset found for object number: $0, generation number: $1",
        object_number, generation_number));
  }
  return it->second;
}

absl::StatusOr<IndirectObject> PdfReader::GetObject(
    uint32_t object_number, uint32_t generation_number) {
  Object::IndirectReference ref{.object_number = object_number,
                                .generation_number = generation_number};
  if (active_objects_.contains(ref)) {
    return absl::InvalidArgumentError(
        absl::StrCat("Circular reference detected for object: ", object_number,
                     " ", generation_number));
  }

  struct ActiveObjectGuard {
    absl::btree_set<Object::IndirectReference>& active_objects;
    Object::IndirectReference ref;
    ~ActiveObjectGuard() { active_objects.erase(ref); }
  };
  active_objects_.insert(ref);
  ActiveObjectGuard guard{active_objects_, ref};

  ABSL_ASSIGN_OR_RETURN(uint64_t offset,
                        GetObjectOffset(object_number, generation_number));
  IndirectObject obj;
  if (absl::Status status = ReadIndirectObject(
          input_, offset,
          [this](const Object::IndirectReference& ref) {
            return this->GetObject(ref.object_number, ref.generation_number);
          },
          obj);
      !status.ok()) {
    return status;
  }
  return obj;
}

}  // namespace credentio