// 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 <functional>
#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 "absl/strings/string_view.h"
#include "crypto/default/default_crypto_read_handler.h"
#include "formats/bmff/test_utils.h"
#include "formats/gif/constants.h"
#include "formats/gif/test_builder.h"
#include "formats/jpeg/testing/jpeg_creator.h"
#include "formats/png/test_utils.h"
#include "formats/riff/create_riff.h"
#include "formats/tiff/constants.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/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;

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

// Signs with the claim signer and timestamps with the local TSA of `pki`.
LocalKeyCryptoWriteHandlerOptions CreateWriteOptions(const TestPki& pki) {
  LocalKeyCryptoWriteHandlerOptions options;
  options.claim_signer.key_pem = pki.claim_signer.key_pem;
  options.claim_signer.certificate_chain = {pki.claim_signer.cert_der};
  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),
  });
}

GenerateManifestStoreRequest CreateRequest() {
  GenerateManifestStoreRequest request;
  request.mutable_actions()->mutable_created();
  return request;
}

// Minimal unsigned test assets, one per supported container format.

absl::StatusOr<std::string> CreatePngAsset() {
  return CreatePng({Chunk{.type = 0x49444154, .payload = "pixel_data"}});
}

absl::StatusOr<std::string> CreateJpegAsset() { return CreateJpeg({}); }

absl::StatusOr<std::string> CreateGifAsset() {
  TestGifBuilder builder;
  builder.AddGifHeader();
  builder.AddLogicalDescriptor(/*has_color_table=*/false);
  builder.AddBlock(kGifExtensionIntroducer, kGifExtensionApplication,
                   "other_app_info", "misc_data");
  builder.AddBlock(kGifTrailerIntroducer);
  return builder.GetAsset();
}

absl::StatusOr<std::string> CreateBmffAsset() {
  return absl::StrCat(
      credentio_testing::Box("ftyp", {"isom\x00\x00\x02\x00isom", 12}),
      credentio_testing::Box("moov", ""),
      credentio_testing::Box("mdat", "video_payload"));
}

absl::StatusOr<std::string> CreateTiffAsset() {
  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
  if (!writer.Close()) return writer.status();
  return tiff_bytes;
}

absl::StatusOr<std::string> CreateWebpAsset() {
  return credentio_riff::RiffChunk("WEBP",
                                   credentio_riff::Chunk("VP8 ", "image_data"));
}

absl::StatusOr<std::string> CreateMp3Asset() {
  constexpr char kMp3Bytes[] = {'\xFF', '\xF3', '\x40', '\xC0',
                                '\x01', '\x02', '\x03', '\x04'};
  return std::string(kMp3Bytes, sizeof(kMp3Bytes));
}

absl::StatusOr<std::string> CreatePdfAsset() {
  return std::string(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
)");
}

absl::StatusOr<std::string> CreatePptxAsset() {
  return credentio_testing::GetContents(
      "credentio/formats/zip/testing/"
      "no_c2pa_032_asset.pptx");
}

struct FormatTestCase {
  std::string test_name;
  // File extension used by `WriteAssetFile` to infer the media type.
  std::string extension;
  std::string media_type;
  std::function<absl::StatusOr<std::string>()> create_asset;
};

class AssetWriterFormatTest : public testing::TestWithParam<FormatTestCase> {
 protected:
  // Checks that `signed_asset` has an active manifest that validates against
  // `pki` without failures.
  void ExpectSignedAndValid(riegeli::Reader& signed_asset, const TestPki& pki) {
    auto validation_result =
        CreateTestValidator(pki)->Validate(signed_asset, GetParam().media_type);
    ABSL_ASSERT_OK(validation_result);
    const ValidationStatusSet& validation =
        (*validation_result)->proto().active_manifest().validation();
    EXPECT_THAT(validation.failures(), IsEmpty());
    EXPECT_THAT(validation.successes(), Not(IsEmpty()));
  }
};

TEST_P(AssetWriterFormatTest, WriteAsset) {
  const FormatTestCase& test_case = GetParam();
  absl::StatusOr<std::string> asset_bytes = test_case.create_asset();
  ABSL_ASSERT_OK(asset_bytes);

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

  TestPki pki = CreateTestPki();
  ASSERT_THAT(WriteAsset(CreateWriteOptions(pki), CreateRequest(),
                         test_case.media_type, source, destination),
              IsOk());
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader<> result_reader(output_bytes);
  ExpectSignedAndValid(result_reader, pki);
}

// Unlike `WriteAsset` above, this writes to a real file through
// `riegeli::CFileWriter`, which e.g. `ZipWriter` must be able to read back.
TEST_P(AssetWriterFormatTest, WriteAssetFile) {
  const FormatTestCase& test_case = GetParam();
  absl::StatusOr<std::string> asset_bytes = test_case.create_asset();
  ABSL_ASSERT_OK(asset_bytes);

  const std::string input_path =
      absl::StrCat(testing::TempDir(), "/", test_case.test_name, "_input.",
                   test_case.extension);
  const std::string output_path =
      absl::StrCat(testing::TempDir(), "/", test_case.test_name, "_output.",
                   test_case.extension);
  {
    std::ofstream out(input_path, std::ios::binary);
    out << *asset_bytes;
    ASSERT_TRUE(out.good());
  }

  TestPki pki = CreateTestPki();
  ASSERT_THAT(WriteAssetFile(CreateWriteOptions(pki), CreateRequest(),
                             input_path, output_path),
              IsOk());

  riegeli::CFileReader<> result_reader(output_path);
  ASSERT_TRUE(result_reader.ok());
  ExpectSignedAndValid(result_reader, pki);
}

INSTANTIATE_TEST_SUITE_P(
    AllFormats, AssetWriterFormatTest,
    testing::ValuesIn(std::vector<FormatTestCase>{
        {"Png", "png", "image/png", CreatePngAsset},
        {"Jpeg", "jpg", "image/jpeg", CreateJpegAsset},
        {"Gif", "gif", "image/gif", CreateGifAsset},
        {"Bmff", "mp4", "video/mp4", CreateBmffAsset},
        {"Tiff", "tiff", "image/tiff", CreateTiffAsset},
        {"Webp", "webp", "image/webp", CreateWebpAsset},
        {"Mp3", "mp3", "audio/mpeg", CreateMp3Asset},
        {"Pdf", "pdf", "application/pdf", CreatePdfAsset},
        {"Pptx", "pptx", std::string(kPptxMediaType), CreatePptxAsset},
    }),
    [](const testing::TestParamInfo<FormatTestCase>& info) {
      return info.param.test_name;
    });

TEST(AssetWriterTest, RejectsInPlaceWrite) {
  absl::StatusOr<std::string> png_bytes = CreatePngAsset();
  ABSL_ASSERT_OK(png_bytes);

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

  TestPki pki = CreateTestPki();
  EXPECT_THAT(
      WriteAssetFile(CreateWriteOptions(pki), CreateRequest(), 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_bytes);
}

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

  TestPki pki = CreateTestPki();
  EXPECT_THAT(WriteAsset(CreateWriteOptions(pki), CreateRequest(),
                         "application/octet-stream", source, destination),
              StatusIs(absl::StatusCode::kUnimplemented));
}

}  // namespace
}  // namespace credentio