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

#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/string_view.h"
#include "absl/synchronization/notification.h"
#include "async/async_context.h"
#include "formats/byte_range.h"
#include "formats/png/extractor.h"
#include "formats/png/reader.h"
#include "formats/png/test_utils.h"
#include "generator/generator.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "proto/generator.pb.h"
#include "proto/ingredient_assertion.pb.h"
#include "proto/ingredient_validation_result.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/nonprod_generator.h"
#include "testing/nonprod_validator.h"
#include "validator/asset_validator.h"
#include "validator/ingredient_result.h"
#include "validator/ingredient_result_converter.h"
#include "validator/result.h"

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::testing::AllOf;
using ::testing::Field;
using ::testing::Gt;
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::Optional;

// Returns a minimal 13x17 RGB PNG image: signature, IHDR, IDAT, IEND.
//
// The IDAT payload is not a valid zlib stream; C2PA processing never decodes
// the pixel data.
absl::StatusOr<std::string> CreatePngImage() {
  return CreatePng({
      {.type = 'IHDR',
       .payload = std::string("\x00\x00\x00\x0d"       // Width: 13.
                              "\x00\x00\x00\x11"       // Height: 17.
                              "\x08\x02\x00\x00\x00",  // 8-bit RGB.
                              13)},
      {.type = 'IDAT', .payload = std::string(64, '\x40')},
  });
}

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

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

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

  // Validates `asset` and returns the ingredient validation result describing
  // it, as needed to reference it as an ingredient.
  absl::StatusOr<std::unique_ptr<IngredientResult>> ValidateIngredient(
      absl::string_view asset) {
    riegeli::StringReader reader(asset);
    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<ValidationResult> result,
                          validator_->Validate(reader, "image/png"));
    return IngredientResultConverter().Convert(std::move(*result));
  }

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

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

TEST_F(PngWriterTest, Write) {
  absl::StatusOr<std::string> image = CreatePngImage();
  ABSL_ASSERT_OK(image);
  PngWriter::OutputMetadata output_metadata;
  absl::StatusOr<std::string> result =
      WriteCopy(*writer_, CreateRequest(), *image, &output_metadata);
  ABSL_ASSERT_OK(result);
  EXPECT_TRUE(output_metadata.timestamp_included);

  riegeli::StringReader result_reader(*result);
  EXPECT_THAT(PngExtractor().ExtractManifestStoreLocation(result_reader, {}),
              IsOkAndHolds(Optional(AllOf(Field(&ByteRange::offset, Gt(0)),
                                          Field(&ByteRange::length, Gt(0))))));
}

TEST_F(PngWriterTest, WriteWithoutTimestamp) {
  auto generator =
      CreateNonProdManifestStoreGenerator({.disable_timestamping = true});
  ABSL_ASSERT_OK(generator);
  PngWriter writer(generator->get());
  absl::StatusOr<std::string> image = CreatePngImage();
  ABSL_ASSERT_OK(image);

  PngWriter::OutputMetadata output_metadata;
  ABSL_ASSERT_OK(WriteCopy(writer, CreateRequest(), *image, &output_metadata));
  EXPECT_FALSE(output_metadata.timestamp_included);
}

TEST_F(PngWriterTest, ExistingC2paWithoutOpenedActionFails) {
  absl::StatusOr<std::string> image = CreatePngImage();
  ABSL_ASSERT_OK(image);
  absl::StatusOr<std::string> png_with_c2pa =
      WriteCopy(*writer_, CreateRequest(), *image);
  ABSL_ASSERT_OK(png_with_c2pa);

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

TEST_F(PngWriterTest, ExistingC2paWithOpenedActionSucceeds) {
  absl::StatusOr<std::string> image = CreatePngImage();
  ABSL_ASSERT_OK(image);
  absl::StatusOr<std::string> png_with_c2pa =
      WriteCopy(*writer_, CreateRequest(), *image);
  ABSL_ASSERT_OK(png_with_c2pa);

  absl::StatusOr<std::unique_ptr<IngredientResult>>
      png_with_c2pa_validation_results = ValidateIngredient(*png_with_c2pa);
  ABSL_ASSERT_OK(png_with_c2pa_validation_results);
  EXPECT_THAT((*png_with_c2pa_validation_results)
                  ->proto()
                  .ingredient_validation_results()
                  .active_manifest()
                  .failures(),
              IsEmpty());

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_opened()->set_ingredient_index(0);
  *request.add_ingredients()->mutable_ingredient_validation_results() =
      (*png_with_c2pa_validation_results)->proto();
  absl::StatusOr<std::string> final_png =
      WriteCopy(*writer_, std::move(request), *png_with_c2pa);
  ABSL_ASSERT_OK(final_png);

  riegeli::StringReader source(*final_png);
  absl::StatusOr<std::unique_ptr<ValidationResult>> result =
      validator_->Validate(source, "image/png");
  ABSL_ASSERT_OK(result);
  EXPECT_THAT((*result)->proto().active_manifest().validation().failures(),
              IsEmpty());
}

TEST_F(PngWriterTest, WriteWithHardBindingFails) {
  GenerateManifestStoreRequest request = CreateRequest();
  request.mutable_hard_binding();
  absl::StatusOr<std::string> image = CreatePngImage();
  ABSL_ASSERT_OK(image);
  EXPECT_THAT(WriteCopy(*writer_, std::move(request), *image),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("must not include `hard_binding`")));
}

TEST_F(PngWriterTest, WriteAsync) {
  absl::StatusOr<std::string> image = CreatePngImage();
  ABSL_ASSERT_OK(image);
  riegeli::StringReader source(*image);
  std::string result;
  riegeli::StringWriter destination(&result);

  absl::Notification done;
  absl::Status write_status;
  PngWriter::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(source.Close());
  ASSERT_TRUE(destination.Close());
  EXPECT_TRUE(output_metadata.timestamp_included);

  riegeli::StringReader result_reader(result);
  EXPECT_THAT(PngExtractor().ExtractManifestStoreLocation(result_reader, {}),
              IsOkAndHolds(Optional(AllOf(Field(&ByteRange::offset, Gt(0)),
                                          Field(&ByteRange::length, Gt(0))))));
}

TEST_F(PngWriterTest, WriteWithCustomExclusionProcessor) {
  int custom_processor_call_count = 0;
  PngWriter writer(
      generator_.get(),
      {.exclusion_processor = [&custom_processor_call_count](
                                  riegeli::Reader&, const PngChunk&,
                                  const GenerateManifestStoreRequest&)
           -> absl::StatusOr<std::vector<ByteRange>> {
        ++custom_processor_call_count;
        return std::vector<ByteRange>();
      }});
  absl::StatusOr<std::string> image = CreatePngImage();
  ABSL_ASSERT_OK(image);

  ABSL_ASSERT_OK(WriteCopy(writer, CreateRequest(), *image));
  EXPECT_GT(custom_processor_call_count, 0);
}

}  // namespace
}  // namespace credentio