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

#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/log/absl_check.h"
#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/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/notification.h"
#include "async/async_context.h"
#include "formats/byte_range.h"
#include "formats/jpeg/extractor.h"
#include "formats/jpeg/reader.h"
#include "formats/jpeg/testing/jpeg_creator.h"
#include "formats/jpeg/testing/jumbf_creator.h"
#include "generator/generator.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "proto/generator.pb.h"
#include "proto/manifest.pb.h"
#include "proto/validation_status.pb.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"
#include "testing/fake_manifest_store_generator.h"
#include "testing/nonprod_generator.h"
#include "testing/nonprod_validator.h"
#include "testing/proto_test_utils.h"
#include "validator/asset_validator.h"
#include "validator/result.h"

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::credentio_testing::ParseTextProtoOrDie;
using ::credentio_testing::PartiallyEqualsProto;
using ::testing::AllOf;
using ::testing::Contains;
using ::testing::Field;
using ::testing::Gt;
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::Not;
using ::testing::Optional;
using ::testing::SizeIs;
using ::testing::StartsWith;

constexpr char kExifData[] = "Exif\0\0fake exif data";
constexpr absl::string_view kExifPayload(kExifData, sizeof(kExifData) - 1);

// Returns a JPEG segment with the given marker and payload.
std::string Segment(uint16_t marker, absl::string_view payload) {
  const uint16_t length = static_cast<uint16_t>(payload.size() + 2);
  return absl::StrCat(
      std::string({static_cast<char>(marker >> 8), static_cast<char>(marker),
                   static_cast<char>(length >> 8), static_cast<char>(length)}),
      payload);
}

// Returns a minimal JPEG image: SOI, APP0 (JFIF), APP1 (Exif), DQT, EOI.
std::string CreateJpegImage() {
  return CreateJpeg(JpegParams{
      .payload = {
          Segment(0xffe0,
                  absl::string_view("JFIF\0\x01\x02\0\0\x01\0\x01\0\0", 14)),
          Segment(0xffe1, kExifPayload),
          Segment(0xffdb, std::string(65, '\x01')),
      }});
}

GenerateManifestStoreRequest CreateRequest() {
  GenerateManifestStoreRequest request;
  request.mutable_actions()
      ->mutable_created()
      ->mutable_action_metadata()
      ->set_description("test only");
  return request;
}

// Returns a fake manifest store: a JUMBF box with a payload of the given size.
std::string CreateFakeManifestStore(size_t payload_size = 32) {
  return CreateJumbf({.payload = std::string(payload_size, 'm')});
}

// Returns the byte offset of the Exif payload in `jpeg`.
size_t ExifOffset(absl::string_view jpeg) {
  size_t offset = jpeg.find(kExifPayload);
  ABSL_CHECK_NE(offset, absl::string_view::npos);
  return offset;
}

class JpegWriterTest : public ::testing::Test {
 protected:
  void SetUp() override {
    auto generator = CreateNonProdManifestStoreGenerator();
    ABSL_CHECK_OK(generator);
    generator_ = *std::move(generator);
    writer_ = std::make_unique<JpegWriter>(generator_.get());

    auto validator = CreateNonProdValidator();
    ABSL_CHECK_OK(validator);
    validator_ = *std::move(validator);
  }

  // Validates `asset` and returns the active manifest validation statuses.
  absl::StatusOr<ValidationStatusSet> Validate(absl::string_view asset) {
    riegeli::StringReader reader(asset);
    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<ValidationResult> result,
                          validator_->Validate(reader, "image/jpeg"));
    return result->proto().active_manifest().validation();
  }

  // Expects `asset` to contain a manifest store and to be valid.
  void ExpectValid(absl::string_view asset) {
    riegeli::StringReader reader(asset);
    EXPECT_THAT(
        JpegExtractor().ExtractManifestStoreLocation(reader, {}),
        IsOkAndHolds(Optional(AllOf(Field(&ByteRange::offset, Gt(0)),
                                    Field(&ByteRange::length, Gt(0))))));

    absl::StatusOr<ValidationStatusSet> validation = Validate(asset);
    ABSL_ASSERT_OK(validation);
    EXPECT_THAT(validation->failures(), IsEmpty());
    EXPECT_THAT(validation->successes(),
                Contains(PartiallyEqualsProto(
                    ParseTextProtoOrDie<ValidationStatus>(R"pb(
                      code: "assertion.dataHash.match"
                    )pb"))));
  }

  // Writes a manifest store into a copy of `image` using `writer`.
  static absl::StatusOr<std::string> WriteCopy(
      const JpegWriter& writer, GenerateManifestStoreRequest request,
      absl::string_view image,
      JpegWriter::OutputMetadata* output_metadata = nullptr) {
    riegeli::StringReader source(image);
    std::string result;
    riegeli::StringWriter destination(&result);
    absl::Status status = writer.Write(std::move(request), &source,
                                       &destination, output_metadata);
    if (!status.ok()) return status;
    if (!destination.Close()) return destination.status();
    return result;
  }

  std::unique_ptr<ManifestStoreGenerator> generator_;
  std::unique_ptr<JpegWriter> writer_;
  std::unique_ptr<AssetValidator> validator_;
};

TEST_F(JpegWriterTest, Write) {
  std::string image = CreateJpegImage();
  JpegWriter::OutputMetadata output_metadata;
  absl::StatusOr<std::string> result =
      WriteCopy(*writer_, CreateRequest(), image, &output_metadata);
  ABSL_ASSERT_OK(result);
  EXPECT_TRUE(output_metadata.timestamp_included);
  EXPECT_GT(result->size(), image.size());
  ExpectValid(*result);

  absl::StatusOr<ValidationStatusSet> validation = Validate(*result);
  ABSL_ASSERT_OK(validation);
  EXPECT_THAT(validation->informationals(),
              Not(Contains(PartiallyEqualsProto(
                  ParseTextProtoOrDie<ValidationStatus>(R"pb(
                    code: "assertion.dataHash.additionalExclusionsPresent"
                  )pb")))));
}

TEST_F(JpegWriterTest, WriteInsertsManifestStoreBeforeFirstNonAppSegment) {
  FakeManifestStoreGenerator generator(CreateFakeManifestStore());
  JpegWriter writer(&generator);
  std::string image = CreateJpegImage();
  size_t dqt_offset = image.find("\xff\xdb");
  ASSERT_NE(dqt_offset, std::string::npos);

  absl::StatusOr<std::string> result =
      WriteCopy(writer, CreateRequest(), image);
  ABSL_ASSERT_OK(result);

  // Bytes before the insertion point are unchanged.
  EXPECT_THAT(*result, StartsWith(image.substr(0, dqt_offset)));
  // An APP11 segment containing the manifest store is inserted.
  EXPECT_EQ(result->substr(dqt_offset, 2), "\xff\xeb");
  EXPECT_THAT(*result, HasSubstr(generator.manifest_store()));
  // Bytes after the insertion point are unchanged.
  EXPECT_EQ(result->substr(result->size() - (image.size() - dqt_offset)),
            image.substr(dqt_offset));
}

TEST_F(JpegWriterTest, WriteWithoutOutputMetadata) {
  absl::StatusOr<std::string> result =
      WriteCopy(*writer_, CreateRequest(), CreateJpegImage());
  ABSL_ASSERT_OK(result);
  ExpectValid(*result);
}

TEST_F(JpegWriterTest, WriteWithoutTimestamp) {
  auto generator =
      CreateNonProdManifestStoreGenerator({.disable_timestamping = true});
  ABSL_ASSERT_OK(generator);
  JpegWriter writer(generator->get());

  JpegWriter::OutputMetadata output_metadata;
  absl::StatusOr<std::string> result =
      WriteCopy(writer, CreateRequest(), CreateJpegImage(), &output_metadata);
  ABSL_ASSERT_OK(result);
  EXPECT_FALSE(output_metadata.timestamp_included);
  ExpectValid(*result);
}

TEST_F(JpegWriterTest, WriteInPlace) {
  std::string image = CreateJpegImage();
  const size_t original_size = image.size();
  riegeli::StringWriter<> asset(
      &image, riegeli::StringWriterBase::Options().set_append(true));
  ABSL_ASSERT_OK(asset.status());

  JpegWriter::OutputMetadata output_metadata;
  ABSL_ASSERT_OK(writer_->Write(CreateRequest(), &asset, &output_metadata));
  ASSERT_TRUE(asset.Close());
  EXPECT_TRUE(output_metadata.timestamp_included);
  EXPECT_GT(image.size(), original_size);
  ExpectValid(image);
}

TEST_F(JpegWriterTest, WriteInPlaceWithoutOutputMetadata) {
  std::string image = CreateJpegImage();
  riegeli::StringWriter<> asset(
      &image, riegeli::StringWriterBase::Options().set_append(true));
  ABSL_ASSERT_OK(asset.status());

  ABSL_ASSERT_OK(writer_->Write(CreateRequest(), &asset));
  ASSERT_TRUE(asset.Close());
  ExpectValid(image);
}

TEST_F(JpegWriterTest, WriteInPlaceFailsWhenRandomAccessNotSupported) {
  class NonRandomAccessWriter : public riegeli::StringWriter<> {
   public:
    using StringWriter::StringWriter;
    bool SupportsRandomAccess() override { return false; }
  };
  std::string image = CreateJpegImage();
  NonRandomAccessWriter asset(&image);
  EXPECT_THAT(writer_->Write(CreateRequest(), &asset),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("asset does not support random access")));
}

TEST_F(JpegWriterTest, WriteInPlaceFailsWhenReadModeNotSupported) {
  class NonReadModeWriter : public riegeli::StringWriter<> {
   public:
    using StringWriter::StringWriter;
    bool SupportsReadMode() override { return false; }
  };
  std::string image = CreateJpegImage();
  NonReadModeWriter asset(
      &image, riegeli::StringWriterBase::Options().set_append(true));
  EXPECT_THAT(writer_->Write(CreateRequest(), &asset),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("asset does not support read mode")));
}

TEST_F(JpegWriterTest, WriteWithHardBindingFails) {
  GenerateManifestStoreRequest request = CreateRequest();
  request.mutable_hard_binding();
  EXPECT_THAT(WriteCopy(*writer_, request, CreateJpegImage()),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("must not include `hard_binding`")));
}

TEST_F(JpegWriterTest, WriteInPlaceWithHardBindingFails) {
  GenerateManifestStoreRequest request = CreateRequest();
  request.mutable_hard_binding();
  std::string image = CreateJpegImage();
  const std::string original = image;
  riegeli::StringWriter<> asset(
      &image, riegeli::StringWriterBase::Options().set_append(true));
  EXPECT_THAT(writer_->Write(request, &asset),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("must not include `hard_binding`")));
  ASSERT_TRUE(asset.Close());
  EXPECT_EQ(image, original);
}

TEST_F(JpegWriterTest, WriteWithMultiAssetHash) {
  GenerateManifestStoreRequest request = CreateRequest();
  request.mutable_options()->set_generate_multi_asset_hash(true);
  absl::StatusOr<std::string> result =
      WriteCopy(*writer_, request, CreateJpegImage());
  ABSL_ASSERT_OK(result);
  ExpectValid(*result);
}

TEST_F(JpegWriterTest, WriteFailsWhenNoInsertionPointFound) {
  EXPECT_THAT(
      WriteCopy(*writer_, CreateRequest(),
                absl::string_view("\xff\xd8\xff\xe0\x00\x02", 6)),
      StatusIs(absl::StatusCode::kInvalidArgument,
               HasSubstr("No valid C2PA insertion point found in JPEG")));
}

TEST_F(JpegWriterTest, WriteFailsWhenManifestStoreAlreadyExists) {
  absl::StatusOr<std::string> result =
      WriteCopy(*writer_, CreateRequest(), CreateJpegImage());
  ABSL_ASSERT_OK(result);

  EXPECT_THAT(WriteCopy(*writer_, CreateRequest(), *result),
              StatusIs(absl::StatusCode::kAlreadyExists,
                       HasSubstr("C2PA chunk already exists")));
}

TEST_F(JpegWriterTest, ModifyingHashedBytesInvalidatesAsset) {
  absl::StatusOr<std::string> result =
      WriteCopy(*writer_, CreateRequest(), CreateJpegImage());
  ABSL_ASSERT_OK(result);

  (*result)[ExifOffset(*result) + kExifPayload.size() - 1] ^= 0xff;
  absl::StatusOr<ValidationStatusSet> validation = Validate(*result);
  ABSL_ASSERT_OK(validation);
  EXPECT_THAT(
      validation->failures(),
      Contains(PartiallyEqualsProto(ParseTextProtoOrDie<ValidationStatus>(R"pb(
        code: "assertion.dataHash.mismatch"
      )pb"))));
}

TEST_F(JpegWriterTest, WriteAsync) {
  std::string image = CreateJpegImage();
  riegeli::StringReader source(image);
  std::string result;
  riegeli::StringWriter destination(&result);

  absl::Notification done;
  absl::Status write_status;
  JpegWriter::OutputMetadata output_metadata;
  writer_->Write(CreateRequest(), &source, &destination, &output_metadata,
                 AsyncContext(), [&](absl::Status status) {
                   write_status = std::move(status);
                   done.Notify();
                 });
  done.WaitForNotification();
  ABSL_ASSERT_OK(write_status);
  ASSERT_TRUE(destination.Close());
  EXPECT_TRUE(output_metadata.timestamp_included);
  ExpectValid(result);
}

TEST_F(JpegWriterTest, WriteInPlaceAsync) {
  std::string image = CreateJpegImage();
  riegeli::StringWriter<> asset(
      &image, riegeli::StringWriterBase::Options().set_append(true));

  absl::Notification done;
  absl::Status write_status;
  writer_->Write(CreateRequest(), &asset, /*output_metadata=*/nullptr,
                 AsyncContext(), [&](absl::Status status) {
                   write_status = std::move(status);
                   done.Notify();
                 });
  done.WaitForNotification();
  ABSL_ASSERT_OK(write_status);
  ASSERT_TRUE(asset.Close());
  ExpectValid(image);
}

TEST_F(JpegWriterTest, WriteAsyncReportsErrorThroughCallback) {
  GenerateManifestStoreRequest request = CreateRequest();
  request.mutable_hard_binding();
  std::string image = CreateJpegImage();
  riegeli::StringReader source(image);
  std::string result;
  riegeli::StringWriter destination(&result);

  absl::Notification done;
  absl::Status write_status;
  writer_->Write(request, &source, &destination, /*output_metadata=*/nullptr,
                 AsyncContext(), [&](absl::Status status) {
                   write_status = std::move(status);
                   done.Notify();
                 });
  done.WaitForNotification();
  EXPECT_THAT(write_status, StatusIs(absl::StatusCode::kInvalidArgument));
}

TEST_F(JpegWriterTest, ExclusionProcessorExcludesBytesFromHardBinding) {
  std::vector<std::string> visited_segments;
  JpegWriter writer(
      generator_.get(),
      {.exclusion_processor = [&visited_segments](
                                  riegeli::Reader&, const JpegSegment& segment,
                                  const GenerateManifestStoreRequest& request)
           -> absl::StatusOr<std::vector<ByteRange>> {
        EXPECT_EQ(request.actions().created().action_metadata().description(),
                  "test only");
        visited_segments.push_back(segment.marker_label);
        if (segment.marker_label != "APP1") return std::vector<ByteRange>();
        return std::vector<ByteRange>{{.offset = segment.payload_offset,
                                       .length = segment.payload_length}};
      }});

  absl::StatusOr<std::string> result =
      WriteCopy(writer, CreateRequest(), CreateJpegImage());
  ABSL_ASSERT_OK(result);
  EXPECT_THAT(visited_segments, Contains("APP0"));
  EXPECT_THAT(visited_segments, Contains("APP1"));
  EXPECT_THAT(visited_segments, Contains("DQT"));
  ExpectValid(*result);

  absl::StatusOr<ValidationStatusSet> validation = Validate(*result);
  ABSL_ASSERT_OK(validation);
  EXPECT_THAT(
      validation->informationals(),
      Contains(PartiallyEqualsProto(ParseTextProtoOrDie<ValidationStatus>(R"pb(
        code: "assertion.dataHash.additionalExclusionsPresent"
      )pb"))));

  // Modifying excluded bytes does not invalidate the asset.
  (*result)[ExifOffset(*result) + kExifPayload.size() - 1] ^= 0xff;
  ExpectValid(*result);
}

TEST_F(JpegWriterTest, ExclusionProcessorErrorIsPropagated) {
  JpegWriter writer(
      generator_.get(),
      {.exclusion_processor = [](riegeli::Reader&, const JpegSegment&,
                                 const GenerateManifestStoreRequest&)
           -> absl::StatusOr<std::vector<ByteRange>> {
        return absl::DataLossError("exclusion failure");
      }});

  EXPECT_THAT(WriteCopy(writer, CreateRequest(), CreateJpegImage()),
              StatusIs(absl::StatusCode::kDataLoss, "exclusion failure"));
}

TEST_F(JpegWriterTest, ExclusionProcessorRangeOutsideSegmentFails) {
  JpegWriter writer(
      generator_.get(),
      {.exclusion_processor = [](riegeli::Reader&, const JpegSegment& segment,
                                 const GenerateManifestStoreRequest&)
           -> absl::StatusOr<std::vector<ByteRange>> {
        return std::vector<ByteRange>{
            {.offset = segment.offset, .length = segment.length + 1}};
      }});

  EXPECT_THAT(WriteCopy(writer, CreateRequest(), CreateJpegImage()),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("Invalid or overlapping exclusion range")));
}

TEST_F(JpegWriterTest, GeneratorErrorIsPropagated) {
  FakeManifestStoreGenerator generator;
  generator.set_status(absl::UnavailableError("generator failure"));
  JpegWriter writer(&generator);

  EXPECT_THAT(WriteCopy(writer, CreateRequest(), CreateJpegImage()),
              StatusIs(absl::StatusCode::kUnavailable, "generator failure"));
}

TEST_F(JpegWriterTest, InPlaceGeneratorErrorLeavesAssetUnchanged) {
  FakeManifestStoreGenerator generator;
  generator.set_status(absl::UnavailableError("generator failure"));
  JpegWriter writer(&generator);
  std::string image = CreateJpegImage();
  const std::string original = image;
  riegeli::StringWriter<> asset(
      &image, riegeli::StringWriterBase::Options().set_append(true));

  EXPECT_THAT(writer.Write(CreateRequest(), &asset),
              StatusIs(absl::StatusCode::kUnavailable, "generator failure"));
  ASSERT_TRUE(asset.Close());
  EXPECT_EQ(image, original);
}

TEST_F(JpegWriterTest, WriteFailsWithEmptyManifestStoreTemplate) {
  FakeManifestStoreGenerator generator;
  generator.mutable_manifest_store_template()->set_manifest_store_size(0);
  JpegWriter writer(&generator);

  EXPECT_THAT(WriteCopy(writer, CreateRequest(), CreateJpegImage()),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("invalid manifest store template size: 0")));
}

TEST_F(JpegWriterTest, WriteFailsWithTooLargeManifestStoreTemplate) {
  FakeManifestStoreGenerator generator;
  generator.mutable_manifest_store_template()->set_manifest_store_size(
      100 * 1024 * 1024 + 1);
  JpegWriter writer(&generator);

  EXPECT_THAT(WriteCopy(writer, CreateRequest(), CreateJpegImage()),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("invalid manifest store template size")));
}

TEST_F(JpegWriterTest, OutputMetadataReflectsFinalizeResponse) {
  FakeManifestStoreGenerator generator(CreateFakeManifestStore());
  JpegWriter writer(&generator);

  JpegWriter::OutputMetadata output_metadata;
  generator.mutable_finalize_manifest_store_response()->set_is_timestamped(
      true);
  ABSL_ASSERT_OK(
      WriteCopy(writer, CreateRequest(), CreateJpegImage(), &output_metadata));
  EXPECT_TRUE(output_metadata.timestamp_included);

  output_metadata.timestamp_included = true;
  generator.mutable_finalize_manifest_store_response()->set_is_timestamped(
      false);
  ABSL_ASSERT_OK(
      WriteCopy(writer, CreateRequest(), CreateJpegImage(), &output_metadata));
  EXPECT_FALSE(output_metadata.timestamp_included);
}

TEST_F(JpegWriterTest, LargeManifestStoreIsSplitAcrossApp11Segments) {
  // A manifest store larger than a single APP11 segment (64 KiB).
  FakeManifestStoreGenerator generator(CreateFakeManifestStore(200 * 1024));
  JpegWriter writer(&generator);

  absl::StatusOr<std::string> result =
      WriteCopy(writer, CreateRequest(), CreateJpegImage());
  ABSL_ASSERT_OK(result);

  std::vector<size_t> app11_offsets;
  for (size_t pos = result->find("\xff\xeb"); pos != std::string::npos;
       pos = result->find("\xff\xeb", pos + 1)) {
    app11_offsets.push_back(pos);
  }
  EXPECT_THAT(app11_offsets, SizeIs(Gt(3)));
}

}  // namespace
}  // namespace credentio