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

#include <algorithm>
#include <cstddef>
#include <ostream>
#include <string>

#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "formats/pdf/objects.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "riegeli/bytes/string_writer.h"
#include "riegeli/bytes/writer.h"

namespace credentio {
namespace {

using ::absl_testing::StatusIs;
using ::testing::HasSubstr;

class FailingWriter : public riegeli::Writer {
 public:
  explicit FailingWriter(absl::Status status) { Fail(status); }

  bool PushSlow(size_t min_length, size_t recommended_length) override {
    return false;
  }
};

class StepFailingWriter : public riegeli::Writer {
 public:
  explicit StepFailingWriter(int fail_after_bytes, absl::Status status)
      : fail_after_bytes_(fail_after_bytes), status_(status) {
    set_buffer(buffer_, 1);
  }

  bool PushSlow(size_t min_length, size_t recommended_length) override {
    if (fail_after_bytes_ <= 0) {
      Fail(status_);
      return false;
    }
    size_t write_len =
        std::min(min_length, static_cast<size_t>(fail_after_bytes_));
    fail_after_bytes_ -= write_len;
    if (fail_after_bytes_ <= 0) {
      Fail(status_);
      return false;
    }
    set_buffer(buffer_, 1);
    return true;
  }

 private:
  int fail_after_bytes_;
  absl::Status status_;
  char buffer_[1];
};

TEST(WriteObjectTest, WriteBoolean) {
  Object obj = {.value = Object::Boolean{.value = true}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "true");
}

TEST(WriteObjectTest, WriteInteger) {
  Object obj = {.value = Object::Integer{.value = 123}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "123");
}

TEST(WriteObjectTest, WriteIntegerPreservesNegativeSign) {
  Object obj = {.value = Object::Integer{.value = -123}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "-123");
}

TEST(WriteObjectTest, WriteRealNumber) {
  Object obj = {.value = Object::RealNumber{.raw_value = "1.23"}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "1.23");
}

TEST(WriteObjectTest, WriteLiteralString) {
  Object obj = {.value = Object::LiteralString{.raw_value = "abc"}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "(abc)");
}

TEST(WriteObjectTest, WriteHexadecimalString) {
  Object obj = {.value = Object::HexadecimalString{.raw_value = "4E6F76"}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "<4E6F76>");
}

TEST(WriteObjectTest, WriteName) {
  Object obj = {.value = Object::Name{.value = "Catalog"}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "/Catalog");
}

TEST(WriteObjectTest, WriteNameWithEscapedCharacters) {
  Object obj = {.value = Object::Name{.value = "/\x20\x80"}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "/#2f#20#80");
}

TEST(WriteObjectTest, WriteArray) {
  Object obj = {
      .value = Object::Array{
          .objects = {
              Object{.value = Object::LiteralString{.raw_value = "abc"}},
              Object{.value = Object::LiteralString{.raw_value = "def"}}}}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "[(abc) (def)]");
}

TEST(WriteObjectTest, WriteDictionary) {
  Object obj = {
      .value = Object::Dictionary{
          .entries = {
              {"key_1", Object{.value = Object::Name{.value = "name"}}},
              {"key_2", Object{.value = Object::Integer{.value = 123}}}}}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s.substr(0, 2), "<<");
  // The order of the entries is non-deterministic, so we check the substrings
  // instead of the entire string.
  EXPECT_THAT(output_s, HasSubstr("/key_1 /name"));
  EXPECT_THAT(output_s, HasSubstr("/key_2 123"));
  EXPECT_EQ(output_s.substr(output_s.size() - 2, 2), ">>");
}

TEST(WriteObjectTest, WriteStreamFailure) {
  Object obj = {.value = Object::Stream{.raw_value = "1234567890"}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);

  EXPECT_THAT(
      WriteObject(obj, output),
      StatusIs(absl::StatusCode::kInvalidArgument,
               HasSubstr("Stream must be written as an indirect "
                         "object. Call WriteIndirectObject() instead.")));
}

TEST(WriteObjectTest, WriteNull) {
  Object obj = {.value = Object::Null{}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteObject(obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, "null");
}

TEST(WriteIndirectObjectTest, WriteStream) {
  IndirectObject indirect_obj = {
      .object_number = 1,
      .generation_number = 2,
      .stream_dictionary =
          Object::Dictionary{
              .entries = {{"Length",
                           Object{.value = Object::Integer{.value = 10}}}}},
      .object = {.value = Object::Stream{.raw_value = "1234567890"}}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteIndirectObject(indirect_obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, R"(1 2 obj
<</Length 10>>
stream
1234567890
endstream
endobj
)");
}

TEST(WriteIndirectObjectTest, WriteName) {
  IndirectObject indirect_obj = {
      .object_number = 1,
      .generation_number = 2,
      .object = {.value = Object::Name{.value = "Catalog"}}};

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(WriteIndirectObject(indirect_obj, output));
  ASSERT_TRUE(output.Close());

  EXPECT_EQ(output_s, R"(1 2 obj
/Catalog
endobj
)");
}

TEST(ComputeIndirectStreamHeaderLengthTest, ComputeHeaderLengthCorrectly) {
  IndirectObject indirect_obj = {
      .object_number = 1,
      .generation_number = 2,
      .stream_dictionary =
          Object::Dictionary{
              .entries = {{"Length",
                           Object{.value = Object::Integer{.value = 10}}}}},
      .object = {.value = Object::Stream{.raw_value = "1234567890"}}};

  auto header_len = ComputeIndirectStreamHeaderByteSize(indirect_obj);
  ABSL_ASSERT_OK(header_len);
  EXPECT_EQ(*header_len, 30);
}

TEST(ComputeIndirectStreamHeaderLengthTest, FailsForNonStreamObject) {
  IndirectObject indirect_obj = {
      .object_number = 1,
      .generation_number = 2,
      .object = {.value = Object::LiteralString{.raw_value = "abc"}}};

  EXPECT_THAT(
      ComputeIndirectStreamHeaderByteSize(indirect_obj),
      StatusIs(absl::StatusCode::kInvalidArgument,
               HasSubstr("The object is not an indirect stream object")));
}

TEST(ComputeIndirectStreamHeaderLengthTest,
     HandlesStreamSubstringInDictionaryField) {
  IndirectObject indirect_obj = {
      .object_number = 1,
      .generation_number = 2,
      .stream_dictionary =
          Object::Dictionary{
              .entries = {{"Title",
                           Object{.value =
                                      Object::LiteralString{
                                          .raw_value = "my stream\n title"}}}}},
      .object = {.value = Object::Stream{.raw_value = "1234567890"}}};

  auto header_len = ComputeIndirectStreamHeaderByteSize(indirect_obj);
  ABSL_ASSERT_OK(header_len);
  EXPECT_EQ(*header_len, 45);
}

TEST(WriteObjectTest, InjectionBlocked) {
  // 1. LiteralString Injection is escaped safely
  {
    Object obj = {.value = Object::LiteralString{
                      .raw_value = "payload) /Injected /Dictionary ("}};
    std::string output_s;
    riegeli::StringWriter<> output(&output_s);
    ABSL_ASSERT_OK(WriteObject(obj, output));
    ASSERT_TRUE(output.Close());
    EXPECT_EQ(output_s, "(payload\\) /Injected /Dictionary \\()");
  }

  // 2. HexadecimalString Injection is rejected as invalid hex
  {
    Object obj = {.value = Object::HexadecimalString{
                      .raw_value = "4E6F76> /Injected /Dictionary <"}};
    std::string output_s;
    riegeli::StringWriter<> output(&output_s);
    EXPECT_THAT(WriteObject(obj, output),
                StatusIs(absl::StatusCode::kInvalidArgument,
                         HasSubstr("Invalid hexadecimal string")));
  }

  // 3. RealNumber Injection is rejected as invalid real number
  {
    Object obj = {.value = Object::RealNumber{.raw_value = "1.23 /Injected"}};
    std::string output_s;
    riegeli::StringWriter<> output(&output_s);
    EXPECT_THAT(WriteObject(obj, output),
                StatusIs(absl::StatusCode::kInvalidArgument,
                         HasSubstr("Invalid real number")));
  }

  // 4. Name Injection is escaped safely
  {
    Object obj = {.value = Object::Name{.value = "MyName)(InjectedString)"}};
    std::string output_s;
    riegeli::StringWriter<> output(&output_s);
    ABSL_ASSERT_OK(WriteObject(obj, output));
    ASSERT_TRUE(output.Close());
    EXPECT_EQ(output_s, "/MyName#29#28InjectedString#29");
  }
}

TEST(WriteIndirectObjectTest, EmptyStreamLengthSynchronization) {
  // 1. Genuinely empty stream (not C2PA template) with non-zero /Length should
  // have its /Length forced to 0.
  {
    IndirectObject indirect_obj = {
        .object_number = 1,
        .generation_number = 2,
        .stream_dictionary =
            Object::Dictionary{
                .entries = {{"Length",
                             Object{.value = Object::Integer{.value = 10}}}}},
        .object = {.value = Object::Stream{.raw_value = ""}}};

    std::string output_s;
    riegeli::StringWriter<> output(&output_s);
    ABSL_ASSERT_OK(WriteIndirectObject(indirect_obj, output));
    ASSERT_TRUE(output.Close());

    EXPECT_EQ(output_s, R"(1 2 obj
<</Length 0>>
stream

endstream
endobj
)");
  }

  // 2. C2PA manifest store template stream during measurement (empty raw_value)
  // with is_measurement = true should preserve its original /Length.
  {
    IndirectObject indirect_obj = {
        .object_number = 1,
        .generation_number = 2,
        .stream_dictionary =
            Object::Dictionary{
                .entries =
                    {{"Length", Object{.value = Object::Integer{.value = 100}}},
                     {"Subtype",
                      Object{.value =
                                 Object::Name{.value = "application/c2pa"}}}}},
        .object = {.value = Object::Stream{.raw_value = ""}}};

    std::string output_s;
    riegeli::StringWriter<> output(&output_s);
    ABSL_ASSERT_OK(
        WriteIndirectObject(indirect_obj, output, /*is_measurement=*/true));
    ASSERT_TRUE(output.Close());

    // Substrings checked as dict serialization order is non-deterministic.
    EXPECT_THAT(output_s, HasSubstr("/Length 100"));
    EXPECT_THAT(output_s, HasSubstr("/Subtype /application#2fc2pa"));
    EXPECT_THAT(output_s, HasSubstr("stream\n\nendstream"));
  }

  // 3. Attack vector: Malicious stream dictionary containing
  // '_IsMeasurementCopy' key
  //    with is_measurement = false should still have its /Length forced to 0.
  {
    IndirectObject indirect_obj = {
        .object_number = 1,
        .generation_number = 2,
        .stream_dictionary =
            Object::Dictionary{
                .entries = {{"Length",
                             Object{.value = Object::Integer{.value = 100}}},
                            {"_IsMeasurementCopy",
                             Object{.value = Object::Boolean{.value = true}}}}},
        .object = {.value = Object::Stream{.raw_value = ""}}};

    std::string output_s;
    riegeli::StringWriter<> output(&output_s);
    ABSL_ASSERT_OK(
        WriteIndirectObject(indirect_obj, output, /*is_measurement=*/false));
    ASSERT_TRUE(output.Close());

    EXPECT_THAT(output_s, HasSubstr("/Length 0"));
    EXPECT_THAT(output_s, HasSubstr("/_IsMeasurementCopy true"));
  }
}

struct FailureTestCase {
  Object obj;
  std::string name;

  friend std::ostream& operator<<(std::ostream& os, const FailureTestCase& tc) {
    return os << tc.name;
  }
};

class WriteObjectFailureTest
    : public ::testing::TestWithParam<FailureTestCase> {};

TEST_P(WriteObjectFailureTest, WriteFailsAndReturnsWriterStatus) {
  const FailureTestCase& param = GetParam();
  absl::Status expected_status = absl::InternalError("Simulated write failure");
  FailingWriter writer(expected_status);

  absl::Status status = WriteObject(param.obj, writer);

  EXPECT_THAT(status, StatusIs(absl::StatusCode::kInternal,
                               HasSubstr("Simulated write failure")));
}

INSTANTIATE_TEST_SUITE_P(
    WriteObjectFailureTests, WriteObjectFailureTest,
    ::testing::Values(
        FailureTestCase{.obj = Object{.value = Object::Boolean{.value = true}},
                        .name = "Boolean"},
        FailureTestCase{.obj = Object{.value = Object::Integer{.value = 42}},
                        .name = "Integer"},
        FailureTestCase{
            .obj = Object{.value = Object::RealNumber{.raw_value = "3.14"}},
            .name = "RealNumber"},
        FailureTestCase{
            .obj = Object{.value = Object::LiteralString{.raw_value = "test"}},
            .name = "LiteralString"},
        FailureTestCase{
            .obj =
                Object{.value = Object::HexadecimalString{.raw_value = "AB"}},
            .name = "HexadecimalString"},
        FailureTestCase{.obj = Object{.value = Object::Name{.value = "test"}},
                        .name = "Name"},
        FailureTestCase{.obj = Object{.value = Object::Array{.objects = {}}},
                        .name = "Array"},
        FailureTestCase{
            .obj = Object{.value = Object::Dictionary{.entries = {}}},
            .name = "Dictionary"},
        FailureTestCase{.obj = Object{.value = Object::Null{}}, .name = "Null"},
        FailureTestCase{
            .obj = Object{.value =
                              Object::IndirectReference{
                                  .object_number = 1, .generation_number = 0}},
            .name = "IndirectReference"}),
    [](const ::testing::TestParamInfo<WriteObjectFailureTest::ParamType>&
           info) { return info.param.name; });

TEST(WriteIndirectObjectTest, WriteFailsAndReturnsWriterStatus) {
  IndirectObject indirect_obj = {
      .object_number = 1,
      .generation_number = 2,
      .object = {.value = Object::Name{.value = "Catalog"}}};
  absl::Status expected_status =
      absl::InternalError("Simulated indirect write failure");
  FailingWriter writer(expected_status);

  absl::Status status = WriteIndirectObject(indirect_obj, writer);

  EXPECT_THAT(status, StatusIs(absl::StatusCode::kInternal,
                               HasSubstr("Simulated indirect write failure")));
}

TEST(WriteObjectTest, ArraySpaceWriteFailure) {
  Object obj = {.value = Object::Array{
                    .objects = {Object{.value = Object::Integer{.value = 1}},
                                Object{.value = Object::Integer{.value = 2}}}}};
  StepFailingWriter writer(2, absl::InternalError("Simulated space failure"));

  absl::Status status = WriteObject(obj, writer);

  EXPECT_THAT(status, StatusIs(absl::StatusCode::kInternal,
                               HasSubstr("Simulated space failure")));
}

TEST(WriteObjectTest, ArrayCloseBracketWriteFailure) {
  Object obj = {.value = Object::Array{
                    .objects = {Object{.value = Object::Integer{.value = 1}}}}};
  StepFailingWriter writer(
      2, absl::InternalError("Simulated close bracket failure"));

  absl::Status status = WriteObject(obj, writer);

  EXPECT_THAT(status, StatusIs(absl::StatusCode::kInternal,
                               HasSubstr("Simulated close bracket failure")));
}

TEST(WriteObjectTest, DictionarySpaceWriteFailure) {
  Object obj = {
      .value = Object::Dictionary{
          .entries = {{"K", Object{.value = Object::Integer{.value = 1}}}}}};
  StepFailingWriter writer(4, absl::InternalError("Simulated space failure"));

  absl::Status status = WriteObject(obj, writer);

  EXPECT_THAT(status, StatusIs(absl::StatusCode::kInternal,
                               HasSubstr("Simulated space failure")));
}

TEST(WriteObjectTest, DictionaryCloseBracketsWriteFailure) {
  Object obj = {
      .value = Object::Dictionary{
          .entries = {{"K", Object{.value = Object::Integer{.value = 1}}}}}};
  StepFailingWriter writer(
      6, absl::InternalError("Simulated close brackets failure"));

  absl::Status status = WriteObject(obj, writer);

  EXPECT_THAT(status, StatusIs(absl::StatusCode::kInternal,
                               HasSubstr("Simulated close brackets failure")));
}

TEST(WriteIndirectObjectTest, WriteStreamObjectFailure) {
  IndirectObject indirect_obj = {
      .object_number = 1,
      .generation_number = 2,
      .stream_dictionary =
          Object::Dictionary{
              .entries = {{"Length",
                           Object{.value = Object::Integer{.value = 10}}}}},
      .object = {.value = Object::Stream{.raw_value = "1234567890"}}};
  StepFailingWriter writer(23, absl::InternalError("Simulated stream failure"));

  absl::Status status = WriteIndirectObject(indirect_obj, writer);

  EXPECT_THAT(status, StatusIs(absl::StatusCode::kInternal,
                               HasSubstr("Simulated stream failure")));
}

}  // namespace
}  // namespace credentio