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

#include <cstdint>
#include <string>
#include <utility>
#include <variant>

#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/ascii.h"
#include "absl/strings/match.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
#include "formats/pdf/objects.h"
#include "re2/re2.h"
#include "riegeli/bytes/reader.h"

namespace credentio {
namespace {

constexpr uint64_t kMaxPayloadSize = 1024 * 1024 * 10;        // 10 MiB
constexpr uint64_t kMaxStringTokenLength = 1024 * 1024 * 10;  // 10 MiB
constexpr uint64_t kMaxNameTokenLength = 4096;
constexpr uint64_t kMaxNumberTokenLength = 128;
constexpr int kMaxObjectRecursionDepth = 100;

// PDF White-space characters. See ISO 32000-2 Section 7.2.3 Table 1.
constexpr bool IsPdfWhiteSpace(uint8_t c) {
  return c == 0x00 || c == 0x09 || c == 0x0a || c == 0x0c || c == 0x0d ||
         c == 0x20;
}
// PDF Delimiter characters. See ISO 32000-2 Section 7.2.3 Table 2.
constexpr bool IsPdfDelimiter(uint8_t c) {
  return c == 0x28 || c == 0x29 || c == 0x3c || c == 0x3e || c == 0x5b ||
         c == 0x5d || c == 0x7b || c == 0x7d || c == 0x2f || c == 0x25;
}

static constexpr LazyRE2 kIndirectReferencePattern = {
    R"(^[+-]?(\d+)\s+(\d+)\s+R[^a-zA-Z])"};

// Reads either an integer or a real number from the given input stream at the
// current position.
absl::Status ReadNumber(riegeli::Reader& input, Object& obj) {
  uint8_t token = 0x00;
  std::string buffer;

  while (true) {
    if (buffer.size() >= kMaxNumberTokenLength) {
      return absl::InvalidArgumentError(
          absl::StrCat("Number token exceeds maximum allowed length: ",
                       kMaxNumberTokenLength));
    }
    if (!input.ReadByte(token)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read beginning delimiter"));
    }
    // A number can be terminated by any white-space character or delimiter.
    if (IsPdfWhiteSpace(token) || IsPdfDelimiter(token)) {
      if (!input.Seek(input.pos() - 1)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to seek to offset"));
      }
      break;
    }
    buffer += token;
  }
  if (absl::StrContains(buffer, '.')) {
    obj.value = Object::RealNumber{.raw_value = std::move(buffer)};
  } else {
    int64_t value;
    if (!absl::SimpleAtoi(buffer, &value)) {
      return absl::InvalidArgumentError(
          absl::StrCat("Failed to parse integer. Invalid value: ", buffer));
    }
    obj.value = Object::Integer{.value = std::move(value)};
  }
  return absl::OkStatus();
}

absl::Status ReadBoolean(riegeli::Reader& input, Object& obj) {
  std::string buffer;
  if (!input.Read(5, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  if (buffer.substr(0, 4) == "true") {
    obj.value = Object::Boolean{.value = true};
    if (!input.Seek(input.pos() - 1)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to seek to offset"));
    }
  } else if (buffer == "false") {
    obj.value = Object::Boolean{.value = false};
  } else {
    return absl::InvalidArgumentError(
        absl::StrCat("Failed to parse boolean. Invalid value: ", buffer));
  }
  return absl::OkStatus();
}

absl::Status ReadLiteralString(riegeli::Reader& input, Object& obj) {
  uint8_t token = 0x00;
  std::string buffer;
  int32_t parentheses_count = 1;

  // Consume the beginning delimiter.
  if (!input.ReadByte(token)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to read beginning delimiter"));
  }
  if (token != '(') {
    return absl::InvalidArgumentError(
        absl::StrCat("Invalid delimiter. Expected '(' but got 0x",
                     absl::Hex(token), " at offset ", input.pos() - 1));
  }

  while (true) {
    if (buffer.size() >= kMaxStringTokenLength) {
      return absl::InvalidArgumentError(
          absl::StrCat("Literal string exceeds maximum allowed length: ",
                       kMaxStringTokenLength));
    }
    if (!input.ReadByte(token)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read beginning delimiter"));
    }
    if (token == '(') {
      parentheses_count++;
    }
    if (token == ')' && --parentheses_count == 0) {
      break;
    }
    buffer += token;
  }
  obj.value = Object::LiteralString{.raw_value = std::move(buffer)};
  return absl::OkStatus();
}

absl::Status ReadHexadecimalString(riegeli::Reader& input, Object& obj) {
  uint8_t token = 0x00;
  std::string buffer;

  // Consume the beginning delimiter.
  if (!input.ReadByte(token)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to read beginning delimiter"));
  }
  if (token != '<') {
    return absl::InvalidArgumentError(
        absl::StrCat("Invalid delimiter. Expected '<' but got 0x",
                     absl::Hex(token), " at offset ", input.pos() - 1));
  }

  while (true) {
    if (buffer.size() >= kMaxStringTokenLength) {
      return absl::InvalidArgumentError(
          absl::StrCat("Hexadecimal string exceeds maximum allowed length: ",
                       kMaxStringTokenLength));
    }
    if (!input.ReadByte(token)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read beginning delimiter"));
    }
    if (token == '>') {
      break;
    }
    buffer += token;
  }
  obj.value = Object::HexadecimalString{.raw_value = std::move(buffer)};
  return absl::OkStatus();
}

absl::Status ReadName(riegeli::Reader& input, Object& obj) {
  uint8_t token = 0x00;
  std::string buffer;

  // Consume the beginning delimiter.
  if (!input.ReadByte(token)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to read beginning delimiter"));
  }
  if (token != '/') {
    return absl::InvalidArgumentError(
        absl::StrCat("Invalid delimiter. Expected '/' but got 0x",
                     absl::Hex(token), " at offset ", input.pos() - 1));
  }

  while (true) {
    if (buffer.size() >= kMaxNameTokenLength) {
      return absl::InvalidArgumentError(absl::StrCat(
          "Name token exceeds maximum allowed length: ", kMaxNameTokenLength));
    }
    if (!input.ReadByte(token)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read beginning delimiter"));
    }
    // A name can be terminated by any white-space character or delimiter.
    if (IsPdfWhiteSpace(token) || IsPdfDelimiter(token)) {
      if (!input.Seek(input.pos() - 1)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to seek to offset"));
      }
      break;
    }
    if (token == '#') {
      std::string hex_str(2, '\0');
      if (!input.Read(2, hex_str)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to read hex string"));
      }
      if (!absl::ascii_isxdigit(hex_str[0]) ||
          !absl::ascii_isxdigit(hex_str[1])) {
        return absl::InvalidArgumentError("Invalid hex string");
      }
      int hex_val = 0;
      if (!absl::SimpleHexAtoi(hex_str, &hex_val)) {
        return absl::InvalidArgumentError("Invalid hex string");
      }
      token = hex_val;
    }
    buffer += token;
  }
  obj.value = Object::Name{.value = std::move(buffer)};
  return absl::OkStatus();
}

absl::Status ReadArray(riegeli::Reader& input, Object& obj, int depth) {
  uint8_t token = 0x00;
  Object::Array array;

  // Consume the beginning delimiter.
  if (!input.ReadByte(token)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to read beginning delimiter"));
  }
  if (token != '[') {
    return absl::InvalidArgumentError(
        absl::StrCat("Invalid delimiter. Expected '[' but got '",
                     absl::Hex(token), "' at offset ", input.pos() - 1));
  }

  while (true) {
    ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
    if (!input.ReadByte(token)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read beginning delimiter"));
    }
    if (token == ']') {
      break;
    }
    Object sub_obj;
    ABSL_RETURN_IF_ERROR(ReadObject(input, input.pos() - 1, sub_obj, depth));
    array.objects.push_back(std::move(sub_obj));
  }
  obj.value = std::move(array);

  return absl::OkStatus();
}

absl::Status ReadNull(riegeli::Reader& input, Object& obj) {
  std::string buffer;
  if (!input.Read(4, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  if (buffer == "null") {
    obj.value = Object::Null{};
  } else {
    return absl::InvalidArgumentError(absl::StrCat("Invalid value: ", buffer));
  }
  return absl::OkStatus();
}

absl::Status ReadDictionary(riegeli::Reader& input, Object& obj, int depth) {
  std::string buffer;
  Object::Dictionary dictionary;

  if (!input.Read(2, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  if (buffer != "<<") {
    return absl::InvalidArgumentError(absl::StrCat(
        "Invalid delimiter. Expected '<<' but got '", buffer, "'"));
  }

  while (true) {
    ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
    if (!input.Read(2, buffer)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read buffer"));
    }
    if (buffer == ">>") {
      break;
    }
    if (!input.Seek(input.pos() - 2)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to seek to offset"));
    }
    Object name_obj;
    ABSL_RETURN_IF_ERROR(ReadName(input, name_obj));
    ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
    Object value_obj;
    ABSL_RETURN_IF_ERROR(ReadObject(input, input.pos(), value_obj, depth));
    ABSL_ASSIGN_OR_RETURN(auto name, Object::As<Object::Name>(name_obj));
    dictionary.entries.insert({name.value, std::move(value_obj)});
  }
  obj.value = std::move(dictionary);

  return absl::OkStatus();
}

absl::Status ReadStream(riegeli::Reader& input, int64_t length, Object& obj) {
  std::string buffer;
  Object::Stream stream;

  if (!input.Read(6, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  if (buffer != "stream") {
    return absl::InvalidArgumentError(absl::StrCat(
        "Malformed keyword. Expected 'stream' but got '", buffer, "'"));
  }
  if (!input.Read(2, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  // The spec requires "stream" keyword to be followed by a CARRIAGE RETURN and
  // a LINE FEED, or just a LINE FEED.
  if (buffer == "\r\n") {  // CARRIAGE RETURN + LINE FEED
    // Do nothing.
  } else if (buffer[0] == '\n') {  // LINE FEED
    // Seek back to the end of the LINE FEED.
    if (!input.Seek(input.pos() - 1)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to seek to offset"));
    }
  } else {
    return absl::InvalidArgumentError(absl::StrCat(
        "Malformed stream. Expected '\r\n' or '\n' but got '", buffer, "'"));
  }
  if (length > kMaxPayloadSize) {
    return absl::InvalidArgumentError(
        absl::StrCat("Stream length exceeds maximum allowed size: ", length,
                     " > ", kMaxPayloadSize));
  }
  if (!input.Read(length, stream.raw_value)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  obj.value = std::move(stream);
  // Skip over the end-of-line marker.
  ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
  if (!input.Read(9, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  if (buffer != "endstream") {
    return absl::InvalidArgumentError(absl::StrCat(
        "Malformed keyword. Expected 'endstream' but got '", buffer, "'"));
  }
  return absl::OkStatus();
}
}  // namespace

absl::Status ReadObject(riegeli::Reader& input, int64_t obj_offset, Object& obj,
                        int depth) {
  if (depth > kMaxObjectRecursionDepth) {
    return absl::InvalidArgumentError("Max nesting depth exceeded.");
  }
  if (!input.Seek(obj_offset)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to seek to offset"));
  }

  uint8_t token = 0x00;

  ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
  if (!input.ReadByte(token)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to read beginning delimiter"));
  }
  if (!input.Seek(input.pos() - 1)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to seek to offset"));
  }
  switch (token) {
    case '+':
    case '-':
    case '.':
    case '0':
    case '1':
    case '2':
    case '3':
    case '4':
    case '5':
    case '6':
    case '7':
    case '8':
    case '9': {  // Numbers or indirect reference
      std::string buffer(20, '\0');
      // Read more bytes to check if this is an indirect reference.
      if (!input.Read(20, buffer)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to read buffer"));
      }
      if (!input.Seek(input.pos() - 20)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to seek to offset"));
      }
      if (uint32_t obj_num, gen_num; RE2::PartialMatch(
              buffer, *kIndirectReferencePattern, &obj_num, &gen_num)) {
        obj.value = Object::IndirectReference{.object_number = obj_num,
                                              .generation_number = gen_num};
        auto pos = buffer.find('R');
        // This should never happen.
        if (pos == std::string::npos) {
          return absl::InternalError(
              "Failed to locate keyword 'R' in indirect reference buffer.");
        }
        if (!input.Seek(input.pos() + pos + 1)) {
          return input.StatusOrAnnotate(
              absl::DataLossError("Failed to seek to offset"));
        }
        break;
      }
      ABSL_RETURN_IF_ERROR(ReadNumber(input, obj));
      break;
    }
    case 't':
    case 'f': {  // Boolean
      ABSL_RETURN_IF_ERROR(ReadBoolean(input, obj));
      break;
    }
    case '(': {  // Literal string
      ABSL_RETURN_IF_ERROR(ReadLiteralString(input, obj));
      break;
    }
    case '<': {  // Hexadecimal string or Dictionary
      // Read one more byte to determine the object type.
      std::string buffer(2, '\0');
      if (!input.Read(2, buffer)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to read buffer"));
      }
      if (!input.Seek(input.pos() - 2)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to seek to offset"));
      }
      if (buffer == "<<") {  // Dictionary
        ABSL_RETURN_IF_ERROR(ReadDictionary(input, obj, depth + 1));
      } else {
        // Hexadecimal string
        ABSL_RETURN_IF_ERROR(ReadHexadecimalString(input, obj));
      }
      break;
    }
    case '/': {  // Name
      ABSL_RETURN_IF_ERROR(ReadName(input, obj));
      break;
    }
    case '[': {  // Array
      ABSL_RETURN_IF_ERROR(ReadArray(input, obj, depth + 1));
      break;
    }
    case 'n': {  // Null
      ABSL_RETURN_IF_ERROR(ReadNull(input, obj));
      break;
    }
    case 's': {  // Stream
      std::string buffer(6, '\0');
      if (!input.Read(6, buffer)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to read buffer"));
      }
      if (!input.Seek(input.pos() - 6)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to seek to offset"));
      }
      if (buffer == "stream") {
        return absl::InvalidArgumentError(
            "Stream must reside in an indirect object. Call "
            "ReadIndirectObject() instead.");
      } else {
        return absl::InvalidArgumentError(absl::StrCat(
            "Malformed keyword. Expected 'stream' but got '", buffer, "'"));
      }
      break;
    }
    default:
      return absl::InvalidArgumentError(
          absl::StrCat("Unsupported delimiter: 0x", absl::Hex(token),
                       " at offset ", input.pos()));
  }
  return absl::OkStatus();
}

absl::Status ReadIndirectObject(
    riegeli::Reader& input, int64_t obj_offset,
    absl::AnyInvocable<
        absl::StatusOr<IndirectObject>(const Object::IndirectReference&)>
        indirect_reference_resolver,
    IndirectObject& obj) {
  if (!input.Seek(obj_offset)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to seek to offset"));
  }

  uint8_t token = 0x00;
  std::string buffer;

  // Read the object number.
  while (true) {
    if (buffer.size() >= kMaxNumberTokenLength) {
      return absl::InvalidArgumentError(
          absl::StrCat("Object number token exceeds maximum allowed length: ",
                       kMaxNumberTokenLength));
    }
    if (!input.ReadByte(token)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read beginning delimiter"));
    }
    if (IsPdfWhiteSpace(token)) {
      break;
    }
    absl::StrAppend(&buffer, std::string(1, token));
  }
  bool success = absl::SimpleAtoi(buffer, &obj.object_number);
  if (!success) {
    return absl::InvalidArgumentError(
        absl::StrCat("Failed to parse object number. Invalid value: ", buffer));
  }

  // Read the generation number.
  buffer.clear();
  while (true) {
    if (buffer.size() >= kMaxNumberTokenLength) {
      return absl::InvalidArgumentError(absl::StrCat(
          "Generation number token exceeds maximum allowed length: ",
          kMaxNumberTokenLength));
    }
    if (!input.ReadByte(token)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read beginning delimiter"));
    }
    if (IsPdfWhiteSpace(token)) {
      break;
    }
    absl::StrAppend(&buffer, std::string(1, token));
  }
  success = absl::SimpleAtoi(buffer, &obj.generation_number);
  if (!success) {
    return absl::InvalidArgumentError(absl::StrCat(
        "Failed to parse generation number. Invalid value: ", buffer));
  }

  // Read the enclosed object.
  ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
  if (!input.Read(3, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  if (buffer != "obj") {
    return absl::InvalidArgumentError(absl::StrCat(
        "Malformed keyword. Expected 'obj' but got '", buffer, "'"));
  }
  Object temp_obj;
  ABSL_RETURN_IF_ERROR(ReadObject(input, input.pos(), temp_obj));
  ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
  if (!input.Read(6, buffer)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read buffer"));
  }
  if (!input.Seek(input.pos() - 6)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to seek to offset"));
  }
  if (buffer == "stream") {
    // Maker indicates that this is a stream object, hence the object preceding
    // it must be the stream dictionary.
    ABSL_ASSIGN_OR_RETURN(obj.stream_dictionary,
                          Object::As<Object::Dictionary>(temp_obj));

    ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
    // Determine the length of the stream.
    int64_t length = -1;
    if (const auto it = obj.stream_dictionary->entries.find("Length");
        it != obj.stream_dictionary->entries.end()) {
      if (std::holds_alternative<Object::IndirectReference>(it->second.value)) {
        ABSL_ASSIGN_OR_RETURN(
            auto indirect_ref,
            Object::As<Object::IndirectReference>(it->second));
        // Save the current input stream position, as the
        // `indirect_reference_resolver` might change it.
        auto pos = input.pos();
        ABSL_ASSIGN_OR_RETURN(auto indirect_length_obj,
                              indirect_reference_resolver(indirect_ref));

        ABSL_ASSIGN_OR_RETURN(auto length_obj, Object::As<Object::Integer>(
                                                   indirect_length_obj.object));
        length = length_obj.value;
        // Seek back to the saved position to continue reading the stream.
        if (!input.Seek(pos)) {
          return input.StatusOrAnnotate(
              absl::DataLossError("Failed to seek to offset"));
        }
      } else {
        ABSL_ASSIGN_OR_RETURN(auto length_obj,
                              Object::As<Object::Integer>(it->second));
        length = length_obj.value;
      }
    } else {
      return absl::DataLossError(
          "Stream dictionary does not contain a 'Length' entry.");
    }
    ABSL_RETURN_IF_ERROR(ReadStream(input, length, obj.object));
    // Skip over the end-of-object marker.
    ABSL_RETURN_IF_ERROR(SkipOverWhiteSpaces(input));
    if (!input.Read(6, buffer)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read buffer"));
    }
  } else {
    obj.object = std::move(temp_obj);
  }
  if (buffer != "endobj") {
    return absl::InvalidArgumentError(absl::StrCat(
        "Malformed keyword. Expected 'endobj' but got '", buffer, "'"));
  }
  return absl::OkStatus();
}

absl::Status SkipOverWhiteSpaces(riegeli::Reader& input) {
  uint8_t token = 0x00;
  while (true) {
    if (!input.ReadByte(token)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read beginning delimiter"));
    }
    if (token == '%') {
      // It's a comment, skip until EOL or EOF.
      while (true) {
        if (!input.ReadByte(token)) {
          return input.StatusOrAnnotate(
              absl::DataLossError("Failed to read beginning delimiter"));
        }
        if (token == 0x0a || token == 0x0d) {  // LF or CR
          break;
        }
      }
      continue;
    }
    if (!IsPdfWhiteSpace(token)) {
      // Seek back to the last non-whitespace character.
      if (!input.Seek(input.pos() - 1)) {
        return input.StatusOrAnnotate(
            absl::DataLossError("Failed to seek to offset"));
      }
      break;
    }
  }
  return absl::OkStatus();
}
}  // namespace credentio