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

#include <optional>
#include <string>
#include <utility>
#include <variant>

#include "absl/algorithm/container.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "constants/labels.h"
#include "formats/byte_range.h"
#include "formats/pdf/constants.h"
#include "formats/pdf/objects.h"
#include "formats/pdf/reader.h"
#include "jumbf/utils.h"
#if defined(__GNUC__) || defined(__clang__)
#include <cxxabi.h>  // IWYU pragma: keep

#include <cstdlib>
#include <memory>
#endif
#include "absl/status/status_macros.h"
#include "riegeli/bytes/reader.h"

namespace credentio {
namespace {

std::string Demangle(const char* mangled) {
#if defined(__GNUC__) || defined(__clang__)
  int status = 0;
  std::unique_ptr<char, void (*)(void*)> res{
      abi::__cxa_demangle(mangled, nullptr, nullptr, &status), std::free};
  return (status == 0) ? res.get() : mangled;
#else
  return mangled;
#endif
}

struct C2paManifestStore {
  ByteRange location;
  std::string content;
};

// Checks if `file_spec_dict` is a file spec dictionary for an embedded C2PA
// file stream.
absl::StatusOr<bool> IsC2paManifestFileSpec(
    const Object::Dictionary& file_spec_dict) {
  const auto af_relationship = FindObjectInDictionary<Object::Name>(
      file_spec_dict, kPdfAfRelationshipKey);
  if (!af_relationship.ok()) {
    if (absl::IsNotFound(af_relationship.status())) {
      return false;
    }
    return af_relationship.status();
  }
  return af_relationship->value == kPdfC2paManifestAfRelationship;
}

// Checks that the embedded file stream subtype is as expected.
absl::StatusOr<bool> IsC2paManifestSubtype(
    const IndirectObject& embedded_file_obj) {
  if (!embedded_file_obj.stream_dictionary.has_value()) {
    return false;
  }
  const auto subtype = FindObjectInDictionary<Object::Name>(
      *embedded_file_obj.stream_dictionary, kPdfSubtypeKey);
  if (!subtype.ok()) {
    // Optionally check the subtype if it exists as it is only required in C2PA
    // v2.2 onward.
    if (absl::IsNotFound(subtype.status())) {
      return true;
    }
    return subtype.status();
  }
  return subtype->value == kPdfC2paManifestMimeType;
}

// Gets the indirect reference to the embedded file stream from the file spec
// dictionary under the '/EF/F' key.
absl::StatusOr<Object::IndirectReference> GetEmbeddedFileRef(
    const Object::Dictionary& file_spec_dict) {
  ABSL_ASSIGN_OR_RETURN(
      const auto ef_dict,
      FindObjectInDictionary<Object::Dictionary>(file_spec_dict, kPdfEfKey));

  return FindObjectInDictionary<Object::IndirectReference>(ef_dict, kPdfFKey);
}

// Gets the array of embedded file name and indirect reference pairs from the
// catalog dictionary under the '/Names/EmbeddedFiles/Names' key.
absl::StatusOr<Object::Array> GetEmbeddedFilesNameTreeArray(
    const Object::Dictionary& catalog_dict, PdfReader& reader) {
  if (auto it = catalog_dict.entries.find(kPdfNamesKey);
      it != catalog_dict.entries.end()) {
    // /Catalog/Names
    Object::Dictionary names_dict;
    if (std::holds_alternative<Object::IndirectReference>(it->second.value)) {
      const auto& names_ref =
          std::get<Object::IndirectReference>(it->second.value);
      ABSL_ASSIGN_OR_RETURN(const auto names_obj,
                            reader.GetObject(names_ref.object_number,
                                             names_ref.generation_number));
      ABSL_ASSIGN_OR_RETURN(names_dict,
                            Object::As<Object::Dictionary>(names_obj.object));
    } else if (std::holds_alternative<Object::Dictionary>(it->second.value)) {
      names_dict = std::get<Object::Dictionary>(it->second.value);
    } else {
      return absl::InvalidArgumentError(
          absl::StrCat("Unexpected object type found. Expected "
                       "IndirectReference or Dictionary, but got ",
                       Demangle(typeid(it->second.value).name())));
    }
    // /Catalog/Names/EmbeddedFiles
    auto name_tree_it = names_dict.entries.find(kPdfEmbeddedFilesKey);
    if (name_tree_it == names_dict.entries.end()) {
      return absl::NotFoundError("No '/EmbeddedFiles' dictionary found.");
    }
    Object::Dictionary name_tree_dict;
    if (std::holds_alternative<Object::IndirectReference>(
            name_tree_it->second.value)) {
      const auto& name_tree_ref =
          std::get<Object::IndirectReference>(name_tree_it->second.value);
      ABSL_ASSIGN_OR_RETURN(const auto name_tree_obj,
                            reader.GetObject(name_tree_ref.object_number,
                                             name_tree_ref.generation_number));
      ABSL_ASSIGN_OR_RETURN(
          name_tree_dict, Object::As<Object::Dictionary>(name_tree_obj.object));
    } else if (std::holds_alternative<Object::Dictionary>(
                   name_tree_it->second.value)) {
      name_tree_dict = std::get<Object::Dictionary>(name_tree_it->second.value);
    } else {
      return absl::InvalidArgumentError(
          absl::StrCat("Unexpected object type found. Expected "
                       "IndirectReference or Dictionary, but got ",
                       Demangle(typeid(name_tree_it->second.value).name())));
    }
    // /Catalog/Names/EmbeddedFiles/Names
    return FindObjectInDictionary<Object::Array>(name_tree_dict, kPdfNamesKey);
  } else {
    return absl::NotFoundError(
        "No '/Names' dictionary found in the document catalog.");
  }
}

absl::StatusOr<C2paManifestStore> ExtractManifestStoreInternal(
    riegeli::Reader& input) {
  ABSL_ASSIGN_OR_RETURN(auto reader, PdfReader::Create(&input));

  // Locate the active manifest store in the file trailer dictionary.
  const Object::Dictionary& file_trailer_dict =
      reader->file_trailer_dictionary();

  ABSL_ASSIGN_OR_RETURN(const auto root_ref,
                        FindObjectInDictionary<Object::IndirectReference>(
                            file_trailer_dict, kPdfRootKey));

  ABSL_ASSIGN_OR_RETURN(
      auto root_obj,
      reader->GetObject(root_ref.object_number, root_ref.generation_number));
  ABSL_ASSIGN_OR_RETURN(auto catalog_dict,
                        Object::As<Object::Dictionary>(root_obj.object));

  const auto& file_spec_array =
      FindObjectInDictionary<Object::Array>(catalog_dict, kPdfAfKey);
  if (absl::IsNotFound(file_spec_array.status())) {
    // Early return if the document catalog dictionary does not contain the
    // required '/AF' entry for C2PA manifest extraction.
    return absl::NotFoundError("No manifest store found.");
  }
  ABSL_RETURN_IF_ERROR(file_spec_array.status());

  ABSL_ASSIGN_OR_RETURN(auto name_tree_array,
                        GetEmbeddedFilesNameTreeArray(catalog_dict, *reader));

  for (const auto& file_spec_obj : file_spec_array->objects) {
    Object::Dictionary file_spec_dict;
    Object::IndirectReference file_spec_ref;
    if (std::holds_alternative<Object::IndirectReference>(
            file_spec_obj.value)) {
      // C2PA manifest file spec is expected to be referenced indirectly from
      // the '/AF' array.
      file_spec_ref = std::get<Object::IndirectReference>(file_spec_obj.value);
      ABSL_ASSIGN_OR_RETURN(auto file_spec_dict_obj,
                            reader->GetObject(file_spec_ref.object_number,
                                              file_spec_ref.generation_number));
      ABSL_ASSIGN_OR_RETURN(file_spec_dict, Object::As<Object::Dictionary>(
                                                file_spec_dict_obj.object));
    } else {
      // Not an indirect reference, continue searching.
      continue;
    }

    ABSL_ASSIGN_OR_RETURN(auto is_c2pa_manifest_file_spec,
                          IsC2paManifestFileSpec(file_spec_dict));
    if (!is_c2pa_manifest_file_spec) {
      continue;
    }
    ABSL_ASSIGN_OR_RETURN(auto embedded_file_ref,
                          GetEmbeddedFileRef(file_spec_dict));

    // The C2PA manifest file spec object should also be referenced from the
    // '/EmbeddedFiles' NameTree.
    if (absl::c_any_of(name_tree_array.objects, [&](const auto& name_tree_obj) {
          return std::holds_alternative<Object::IndirectReference>(
                     name_tree_obj.value) &&
                 std::get<Object::IndirectReference>(name_tree_obj.value)
                         .object_number == file_spec_ref.object_number &&
                 std::get<Object::IndirectReference>(name_tree_obj.value)
                         .generation_number == file_spec_ref.generation_number;
        })) {
      // The embedded file reference was found in the NameTree.
      // Proceed with extraction.
    } else {
      // The embedded file reference was NOT found in the NameTree.
      continue;
    }

    ABSL_ASSIGN_OR_RETURN(
        auto embedded_file_obj,
        reader->GetObject(embedded_file_ref.object_number,
                          embedded_file_ref.generation_number));
    ABSL_ASSIGN_OR_RETURN(auto is_c2pa_manifest_subtype,
                          IsC2paManifestSubtype(embedded_file_obj));
    if (!is_c2pa_manifest_subtype) {
      continue;
    }
    ABSL_ASSIGN_OR_RETURN(auto embedded_file_stream,
                          Object::As<Object::Stream>(embedded_file_obj.object));
    ABSL_ASSIGN_OR_RETURN(
        auto embedded_file_offset,
        reader->GetObjectOffset(embedded_file_ref.object_number,
                                embedded_file_ref.generation_number));
    return C2paManifestStore{
        .location = {.offset = embedded_file_offset,
                     .length = input.pos() - embedded_file_offset},
        .content = embedded_file_stream.raw_value};
  }
  return absl::NotFoundError("No manifest store found.");
}
}  // namespace

absl::StatusOr<std::string> PdfExtractor::ExtractManifestStore(
    riegeli::Reader& input) const {
  ABSL_ASSIGN_OR_RETURN(auto c2pa_manifest_store,
                        ExtractManifestStoreInternal(input));
  return c2pa_manifest_store.content;
}

absl::StatusOr<std::optional<ByteRange>>
PdfExtractor::ExtractManifestStoreLocation(riegeli::Reader& input,
                                           ExtractOptions options) const {
  auto c2pa_manifest_store = ExtractManifestStoreInternal(input);

  // Extraction failed but C2PA is required, return an error.
  if (!c2pa_manifest_store.ok() && options.requires_c2pa) {
    return absl::NotFoundError("No manifest store found");
  }
  // Extraction failed but C2PA is not required, return nullopt.
  if (!c2pa_manifest_store.ok() && !options.requires_c2pa) {
    return std::nullopt;
  }
  // Extraction succeeded, return the location.
  return c2pa_manifest_store->location;
}

bool PdfExtractor::MightBeC2paManifestStore(absl::string_view payload) const {
  return jumbf::HasDescriptionBoxMatching(payload, kManifestStoreUuid,
                                          kMinimumJumbfDescriptionToggles,
                                          kManifestStoreLabel)
      .value_or(false);
}
}  // namespace credentio