// 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 <string>

#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "crypto/algorithms.h"
#include "formats/pdf/constants.h"
#include "formats/pdf/extractor.h"
#include "formats/pdf/objects.h"
#include "formats/pdf/reader.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "proto/data_hash_assertion.pb.h"
#include "proto/generator.pb.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"
#include "testing/nonprod_generator.h"

namespace credentio {
namespace {

using ::absl_testing::IsOk;
using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::testing::ElementsAre;
using ::testing::HasSubstr;

constexpr absl::string_view kTestPdfContent = R"(%PDF-1.7
%%%%%
1 0 obj
/test_name
endobj
2 0 obj
123.456
endobj
3 0 obj
<</Type /Catalog>>
endobj
xref
0 4
0000000000 65535 f 
0000000015 00000 n 
0000000041 00000 n 
0000000064 00000 n 
trailer
<</Root 3 0 R /Size 4>>
startxref
98
%%EOF
)";

constexpr absl::string_view kTestPdfContentWithDirectNamesDict = R"(%PDF-1.7
%%%%%
1 0 obj
/test_name
endobj
2 0 obj
123.456
endobj
3 0 obj
<</Names <<>> /Type /Catalog>>
endobj
xref
0 4
0000000000 65535 f 
0000000015 00000 n 
0000000041 00000 n 
0000000064 00000 n 
trailer
<</Root 3 0 R /Size 4>>
startxref
110
%%EOF
)";

constexpr absl::string_view kTestPdfContentWithFreeObjectEntries = R"(%PDF-1.7
%%%%%
1 0 obj
/test_name
endobj
3 0 obj
123.456
endobj
5 0 obj
<</Type /Catalog>>
endobj
xref
0 6
0000000002 65535 f 
0000000015 00000 n 
0000000004 00000 f 
0000000041 00000 n 
0000000000 00000 f 
0000000064 00000 n 
trailer
<</Root 5 0 R /Size 6>>
startxref
98
%%EOF
)";

absl::StatusOr<std::string> CreateManifestStore() {
  ABSL_ASSIGN_OR_RETURN(auto generator,
                        CreateNonProdManifestStoreGenerator({}));
  GenerateManifestStoreRequest request;
  DataHashAssertion* data_hash =
      request.mutable_hard_binding()->mutable_data_hash();
  data_hash->set_alg(absl::StrCat(HashAlgorithm::kSha256));
  data_hash->set_hash("data_hash");
  request.mutable_actions()->mutable_created();
  ABSL_ASSIGN_OR_RETURN(auto response,
                        generator->GenerateManifestStore(request));
  return std::string(response.manifest_store());
}

TEST(PdfRewriterTest, RewriteForEmbedding) {
  riegeli::StringReader<> input(kTestPdfContent);
  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  EXPECT_THAT(PdfEmbedder().RewriteForEmbedding(
                  input, /*estimated_manifest_store_size=*/100, output),
              IsOk());
  ASSERT_TRUE(output.Close());

  riegeli::StringReader<> output_reader(output_s);
  EXPECT_THAT(PdfExtractor().ExtractManifestStore(output_reader),
              IsOkAndHolds(std::string(100, '\0')));
}

TEST(PdfRewriterTest, RewriteForEmbeddingMovesDirectNamesDict) {
  riegeli::StringReader<> input(kTestPdfContentWithDirectNamesDict);
  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  EXPECT_THAT(PdfEmbedder().RewriteForEmbedding(
                  input, /*estimated_manifest_store_size=*/100, output),
              IsOk());
  ASSERT_TRUE(output.Close());

  riegeli::StringReader<> output_reader(output_s);
  auto reader = PdfReader::Create(&output_reader);
  ABSL_ASSERT_OK(reader);
  auto trailer_dict = (*reader)->file_trailer_dictionary();
  auto doc_catalog_dict_obj = (*reader)->GetObject(3, 0);
  ABSL_ASSERT_OK(doc_catalog_dict_obj);
  auto doc_catalog_dict =
      Object::As<Object::Dictionary>(doc_catalog_dict_obj->object);
  ABSL_ASSERT_OK(doc_catalog_dict);

  // Verify that the catalog names dictionary is moved into an indirect object
  // and referenced from the document catalog dictionary.
  EXPECT_THAT(
      (*reader)->GetObject(4, 0),
      IsOkAndHolds(IndirectObject{
          .object_number = 4,
          .generation_number = 0,
          .object = Object{
              .value = Object::Dictionary{
                  .entries = {{"EmbeddedFiles",
                               Object{.value = Object::IndirectReference{
                                          .object_number = 7,
                                          .generation_number = 0}}}}}}}));
  EXPECT_THAT(doc_catalog_dict->entries.at("Names"),
              (Object{.value = Object::IndirectReference{
                          .object_number = 4, .generation_number = 0}}));
}

TEST(PdfRewriterTest, RewriteForEmbeddingRespectsFreeObjectEntries) {
  riegeli::StringReader<> input(kTestPdfContentWithFreeObjectEntries);
  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  EXPECT_THAT(PdfEmbedder().RewriteForEmbedding(
                  input, /*estimated_manifest_store_size=*/100, output),
              IsOk());
  ASSERT_TRUE(output.Close());

  riegeli::StringReader<> output_reader(output_s);
  auto reader = PdfReader::Create(&output_reader);
  ABSL_ASSERT_OK(reader);
  EXPECT_THAT((*reader)->free_entries(),
              ElementsAre(Object::IndirectReference{.object_number = 0,
                                                    .generation_number = 65535},
                          Object::IndirectReference{.object_number = 2,
                                                    .generation_number = 0},
                          Object::IndirectReference{.object_number = 4,
                                                    .generation_number = 0}));
}

TEST(PdfRewriterTest, RewriteForEmbeddingFailsIfC2paManifestAlreadyExists) {
  // Prepare a C2PA-embeddable PDF.
  riegeli::StringReader<> input(kTestPdfContent);
  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  EXPECT_THAT(PdfEmbedder().RewriteForEmbedding(
                  input, /*estimated_manifest_store_size=*/100, output),
              IsOk());
  ASSERT_TRUE(output.Close());

  // Rewrite the PDF again, it should fail.
  riegeli::StringReader<> output_reader(output_s);
  std::string output_s2;
  riegeli::StringWriter<> output2(&output_s2);
  EXPECT_THAT(
      PdfEmbedder().RewriteForEmbedding(
          output_reader, /*estimated_manifest_store_size=*/100, output2),
      StatusIs(absl::StatusCode::kAlreadyExists,
               HasSubstr("C2PA manifest store already exists")));
}

TEST(PdfRewriterTest, Embed) {
  auto manifest_store = CreateManifestStore();
  ABSL_ASSERT_OK(manifest_store);

  // Prepare a C2PA-embeddable PDF.
  riegeli::StringReader<> input(kTestPdfContent);
  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  absl::StatusOr<int64_t> manifest_store_offset =
      PdfEmbedder().RewriteForEmbedding(input, manifest_store->size(), output);
  ABSL_ASSERT_OK(manifest_store_offset);
  ASSERT_TRUE(output.Close());

  // Embed the manifest store.
  absl::Cord embeddable_asset(output_s);
  EXPECT_THAT(PdfEmbedder().Embed(*manifest_store, *manifest_store_offset,
                                  embeddable_asset),
              IsOk());

  riegeli::StringReader<> output_reader(embeddable_asset.Flatten());
  EXPECT_THAT(PdfExtractor().ExtractManifestStore(output_reader),
              IsOkAndHolds(*manifest_store));
}

TEST(PdfRewriterTest, RewriteForEmbeddingCorrectLayout) {
  riegeli::StringReader<> input(kTestPdfContent);
  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  EXPECT_THAT(PdfEmbedder().RewriteForEmbedding(
                  input, /*estimated_manifest_store_size=*/100, output),
              IsOk());
  ASSERT_TRUE(output.Close());

  riegeli::StringReader<> output_reader(output_s);
  auto reader = PdfReader::Create(&output_reader);
  ABSL_ASSERT_OK(reader);

  // The catalog should have /AF pointing to the file spec.
  auto doc_catalog_dict_obj = (*reader)->GetObject(3, 0);
  ABSL_ASSERT_OK(doc_catalog_dict_obj);
  auto doc_catalog_dict =
      Object::As<Object::Dictionary>(doc_catalog_dict_obj->object);
  ABSL_ASSERT_OK(doc_catalog_dict);
  auto af_array =
      FindObjectInDictionary<Object::Array>(*doc_catalog_dict, kPdfAfKey);
  ABSL_ASSERT_OK(af_array);
  ASSERT_EQ(af_array->objects.size(), 1);
  auto file_spec_ref =
      Object::As<Object::IndirectReference>(af_array->objects[0]);
  ABSL_ASSERT_OK(file_spec_ref);

  // Inspect the file spec object.
  auto file_spec_obj = (*reader)->GetObject(file_spec_ref->object_number,
                                            file_spec_ref->generation_number);
  ABSL_ASSERT_OK(file_spec_obj);
  auto file_spec_dict = Object::As<Object::Dictionary>(file_spec_obj->object);
  ABSL_ASSERT_OK(file_spec_dict);

  // Subtype should NOT be in file spec dictionary.
  EXPECT_THAT(
      FindObjectInDictionary<Object::Name>(*file_spec_dict, kPdfSubtypeKey),
      StatusIs(absl::StatusCode::kNotFound));

  // AFRelationship should be C2PA_Manifest.
  EXPECT_THAT(FindObjectInDictionary<Object::Name>(*file_spec_dict,
                                                   kPdfAfRelationshipKey),
              IsOkAndHolds(Object::Name{
                  .value = std::string(kPdfC2paManifestAfRelationship)}));

  // Get the embedded file stream ref.
  auto ef_dict =
      FindObjectInDictionary<Object::Dictionary>(*file_spec_dict, kPdfEfKey);
  ABSL_ASSERT_OK(ef_dict);
  auto stream_ref =
      FindObjectInDictionary<Object::IndirectReference>(*ef_dict, kPdfFKey);
  ABSL_ASSERT_OK(stream_ref);

  // Inspect the stream object.
  auto stream_obj = (*reader)->GetObject(stream_ref->object_number,
                                         stream_ref->generation_number);
  ABSL_ASSERT_OK(stream_obj);
  ASSERT_TRUE(stream_obj->stream_dictionary.has_value());
  const auto& stream_dict = *stream_obj->stream_dictionary;

  // Subtype should be in stream dictionary.
  EXPECT_THAT(FindObjectInDictionary<Object::Name>(stream_dict, kPdfSubtypeKey),
              IsOkAndHolds(Object::Name{
                  .value = std::string(kPdfC2paManifestMimeType)}));

  // Params should be in stream dictionary.
  auto params_dict =
      FindObjectInDictionary<Object::Dictionary>(stream_dict, kPdfParamsKey);
  ABSL_ASSERT_OK(params_dict);
  // Params should contain Size.
  EXPECT_THAT(
      FindObjectInDictionary<Object::Integer>(*params_dict, kPdfSizeKey),
      IsOkAndHolds(Object::Integer{.value = 100}));

  // Params should contain ModDate.
  auto mod_date_obj = FindObjectInDictionary<Object::LiteralString>(
      *params_dict, kPdfModDateKey);
  ABSL_ASSERT_OK(mod_date_obj);
  EXPECT_THAT(mod_date_obj->raw_value,
              testing::MatchesRegex("D:[0-9]{14}[+]00'00'"));
}

}  // namespace
}  // namespace credentio