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

#include <cstdint>
#include <fstream>
#include <ios>
#include <iterator>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/log/absl_check.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "crypto/default/default_crypto_read_handler.h"
#include "formats/bmff/extractor.h"
#include "formats/bmff/test_utils.h"
#include "formats/gif/constants.h"
#include "formats/gif/extractor.h"
#include "formats/gif/test_builder.h"
#include "formats/id3/extractor.h"
#include "formats/jpeg/extractor.h"
#include "formats/jpeg/testing/jpeg_creator.h"
#include "formats/pdf/extractor.h"
#include "formats/png/extractor.h"
#include "formats/png/test_utils.h"
#include "formats/riff/create_riff.h"
#include "formats/riff/extractor.h"
#include "formats/tiff/constants.h"
#include "formats/tiff/extractor.h"
#include "formats/zip/extractor.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/cfile_reader.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"
#include "riegeli/endian/endian_writing.h"
#include "testing/crypto_testing.h"
#include "testing/local_key_crypto_write_handler.h"
#include "testing/test_file_utils.h"
#include "validator/asset_validator_impl.h"
#include "validator/validator_options.h"

namespace credentio {
namespace {

using ::absl_testing::IsOk;
using ::absl_testing::StatusIs;
using ::testing::IsEmpty;
using ::testing::Not;

LocalKeyCryptoWriteHandlerOptions CreateWriteOptions(
    const TestPki& pki, bool disable_timestamping) {
  LocalKeyCryptoWriteHandlerOptions options;
  options.claim_signer.key_pem = pki.claim_signer.key_pem;
  options.claim_signer.certificate_chain = {pki.claim_signer.cert_der};
  options.disable_timestamping = disable_timestamping;
  if (!disable_timestamping) {
    options.tsa.key_pem = pki.tsa.key_pem;
    options.tsa.cert_pem = pki.tsa.cert_pem;
  }
  return options;
}

// Returns a validator that enforces claim signer and TSA trust against the
// root CA of `pki`.
std::unique_ptr<AssetValidatorImpl> CreateTestValidator(const TestPki& pki) {
  DefaultCryptoReadHandlerOptions crypto_options;
  crypto_options.claim_signer_trust_anchors_pem = pki.root_ca.cert_pem;
  crypto_options.tsa_trust_anchors_pem = pki.root_ca.cert_pem;
  auto crypto_read_handler = CreateDefaultCryptoReadHandler(crypto_options);
  ABSL_CHECK_OK(crypto_read_handler);
  return std::make_unique<AssetValidatorImpl>(ValidatorOptions{
      .crypto_read_handler = *std::move(crypto_read_handler),
  });
}

TEST(AssetWriterTest, WritesPngManifestStore) {
  auto png_or = CreatePng({Chunk{.type = 0x49444154, .payload = "pixel_data"}});
  ABSL_ASSERT_OK(png_or);

  riegeli::StringReader<> source(*png_or);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                         std::move(request), "image/png", source, destination),
              IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  PngExtractor extractor;
  auto extracted = extractor.ExtractManifestStore(result_reader);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(*extracted, Not(IsEmpty()));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, "image/png");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WritesJpegManifestStore) {
  std::string jpeg_bytes = CreateJpeg({});
  riegeli::StringReader<> source(jpeg_bytes);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                         std::move(request), "image/jpeg", source, destination),
              IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  JpegExtractor extractor;
  auto extracted = extractor.ExtractManifestStore(result_reader);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(*extracted, Not(IsEmpty()));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, "image/jpeg");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WritesGifManifestStore) {
  TestGifBuilder builder;
  builder.AddGifHeader();
  builder.AddLogicalDescriptor(/*has_color_table=*/false);
  builder.AddBlock(kGifExtensionIntroducer, kGifExtensionApplication,
                   "other_app_info", "misc_data");
  builder.AddBlock(kGifTrailerIntroducer);
  std::string gif_bytes = builder.GetAsset();

  riegeli::StringReader<> source(gif_bytes);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                         std::move(request), "image/gif", source, destination),
              IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  GifExtractor extractor;
  auto extracted = extractor.ExtractManifestStore(result_reader);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(*extracted, Not(IsEmpty()));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, "image/gif");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WritesBmffManifestStore) {
  std::string bmff_bytes = absl::StrCat(
      credentio_testing::Box("ftyp", {"isom\x00\x00\x02\x00isom", 12}),
      credentio_testing::Box("moov", ""),
      credentio_testing::Box("mdat", "video_payload"));

  riegeli::StringReader<> source(bmff_bytes);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                         std::move(request), "video/mp4", source, destination),
              IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  BmffExtractor extractor;
  auto extracted = extractor.ExtractManifestStore(result_reader);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(*extracted, Not(IsEmpty()));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, "video/mp4");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WritesTiffManifestStore) {
  std::string tiff_bytes;
  riegeli::StringWriter<> writer(&tiff_bytes);
  riegeli::WriteBigEndian<uint32_t>(kTiffBigEndianHeader, writer);
  riegeli::WriteBigEndian<uint32_t>(8, writer);  // Offset to first IFD
  riegeli::WriteBigEndian<uint16_t>(1, writer);  // 1 entry
  riegeli::WriteBigEndian<uint16_t>(0x0100, writer);
  riegeli::WriteBigEndian<uint16_t>(kTiffTypeShort, writer);
  riegeli::WriteBigEndian<uint32_t>(1, writer);
  riegeli::WriteBigEndian<uint32_t>(10, writer);
  riegeli::WriteBigEndian<uint32_t>(0, writer);  // next IFD offset
  ASSERT_TRUE(writer.Close());

  riegeli::StringReader<> source(tiff_bytes);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                         std::move(request), "image/tiff", source, destination),
              IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  TiffExtractor extractor;
  auto extracted = extractor.ExtractManifestStore(result_reader);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(*extracted, Not(IsEmpty()));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, "image/tiff");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WritesRiffManifestStore) {
  std::string riff_bytes = credentio_riff::RiffChunk(
      "WEBP", credentio_riff::Chunk("VP8 ", "image_data"));

  riegeli::StringReader<> source(riff_bytes);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                         std::move(request), "image/webp", source, destination),
              IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  RiffExtractor extractor;
  auto extracted = extractor.ExtractManifestStore(result_reader);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(*extracted, Not(IsEmpty()));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, "image/webp");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WritesId3ManifestStore) {
  constexpr char kMp3Bytes[] = {'\xFF', '\xF3', '\x40', '\xC0',
                                '\x01', '\x02', '\x03', '\x04'};
  std::string id3_bytes(kMp3Bytes, sizeof(kMp3Bytes));

  riegeli::StringReader<> source(id3_bytes);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                         std::move(request), "audio/mpeg", source, destination),
              IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  Id3Extractor extractor;
  auto extracted = extractor.ExtractManifestStore(result_reader);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(*extracted, Not(IsEmpty()));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, "audio/mpeg");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WritesPdfManifestStore) {
  constexpr absl::string_view kPdfBytes = 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
)";

  riegeli::StringReader<> source(kPdfBytes);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(
      WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                 std::move(request), "application/pdf", source, destination),
      IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  PdfExtractor extractor;
  auto extracted = extractor.ExtractManifestStore(result_reader);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(*extracted, Not(IsEmpty()));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validator = CreateTestValidator(pki);
  auto validation_result =
      validator->Validate(result_reader, "application/pdf");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WritesZipManifestStore) {
  absl::StatusOr<std::string> pptx_bytes = credentio_testing::GetContents(
      "credentio/formats/zip/testing/"
      "no_c2pa_032_asset.pptx");
  ABSL_ASSERT_OK(pptx_bytes);
  riegeli::StringReader<> source(*pptx_bytes);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  constexpr absl::string_view kPptxMediaType =
      "application/"
      "vnd.openxmlformats-officedocument.presentationml.presentation";

  TestPki pki = CreateTestPki();
  EXPECT_THAT(
      WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                 std::move(request), kPptxMediaType, source, destination),
      IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  ZipExtractor extractor;
  auto extracted = extractor.ExtractManifestStore(result_reader);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(*extracted, Not(IsEmpty()));

  ASSERT_TRUE(result_reader.Seek(0));
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, kPptxMediaType);
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WritesPngWithTimestamp) {
  auto png_or = CreatePng({Chunk{.type = 0x49444154, .payload = "pixel_data"}});
  ABSL_ASSERT_OK(png_or);

  riegeli::StringReader<> source(*png_or);
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(
      WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/false),
                 std::move(request), "image/png", source, destination),
      IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, "image/png");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, WriteAssetFileWritesAndValidatesPng) {
  auto png_or = CreatePng({Chunk{.type = 0x49444154, .payload = "pixel_data"}});
  ABSL_ASSERT_OK(png_or);

  const std::string input_path =
      absl::StrCat(testing::TempDir(), "/test_input.png");
  const std::string output_path =
      absl::StrCat(testing::TempDir(), "/test_output.png");

  {
    std::ofstream out(input_path, std::ios::binary);
    out << *png_or;
  }

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(
      WriteAssetFile(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                     std::move(request), input_path, output_path),
      IsOk());

  riegeli::CFileReader<> result_reader(output_path);
  ASSERT_TRUE(result_reader.ok());
  auto validator = CreateTestValidator(pki);
  auto validation_result = validator->Validate(result_reader, "image/png");
  ABSL_ASSERT_OK(validation_result);
  EXPECT_THAT(
      (*validation_result)->proto().active_manifest().validation().failures(),
      IsEmpty());
}

TEST(AssetWriterTest, RejectsInPlaceWrite) {
  auto png_or = CreatePng({Chunk{.type = 0x49444154, .payload = "pixel_data"}});
  ABSL_ASSERT_OK(png_or);

  const std::string path =
      absl::StrCat(testing::TempDir(), "/test_in_place.png");
  {
    std::ofstream out(path, std::ios::binary);
    out << *png_or;
  }

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(
      WriteAssetFile(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                     std::move(request), path, path),
      StatusIs(absl::StatusCode::kInvalidArgument));

  // The input asset must be left untouched.
  std::ifstream in(path, std::ios::binary);
  const std::string contents((std::istreambuf_iterator<char>(in)),
                             std::istreambuf_iterator<char>());
  EXPECT_EQ(contents, *png_or);
}

TEST(AssetWriterTest, RejectsUnsupportedMediaType) {
  riegeli::StringReader<> source("fake_unknown_bytes");
  std::string output_bytes;
  riegeli::StringWriter<> destination(&output_bytes);

  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();

  TestPki pki = CreateTestPki();
  EXPECT_THAT(WriteAsset(CreateWriteOptions(pki, /*disable_timestamping=*/true),
                         std::move(request), "application/octet-stream", source,
                         destination),
              StatusIs(absl::StatusCode::kUnimplemented));
}

}  // namespace
}  // namespace credentio