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

#include <sys/types.h>

#include <cstdint>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "formats/asset_box.h"
#include "formats/byte_range.h"
#include "formats/extractor_result.h"
#include "formats/png/test_utils.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "riegeli/bytes/string_reader.h"
#include "testing/jumbf_utils.h"

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::testing::HasSubstr;
using ::testing::IsEmpty;

constexpr int64_t kMaxChunkBytes = 2 * 1048576;  // 2MiB

std::vector<AssetBox> CreateBoxes(
    std::vector<std::pair<std::string, uint32_t>> expected) {
  uint64_t offset = 0;
  std::vector<AssetBox> boxes;
  for (const auto& name_and_size : expected) {
    boxes.push_back(AssetBox{
        .identifier = name_and_size.first,
        .byte_range = {.offset = offset, .length = name_and_size.second},
    });
    offset += name_and_size.second;
  }
  return boxes;
}

struct TestCase {
  std::string name;
  std::string contents;
  absl::StatusOr<ExtractorResult> result;
  uint32_t max_chunk_bytes = kMaxChunkBytes;
};

class ExtractorTest : public testing::TestWithParam<TestCase> {};

INSTANTIATE_TEST_SUITE_P(
    ExtractorTests, ExtractorTest,
    testing::ValuesIn({
        TestCase{
            .name = "EmptyContent",
            .contents = "",
            .result = absl::NotFoundError("No manifest store found"),
        },
        TestCase{
            .name = "NoPngMarker",
            .contents = "this_is_clearly_not_a_png",
            .result = absl::InvalidArgumentError(
                "input does not start with PNG marker"),
        },
        TestCase{
            .name = "OnlyPngMarker",
            .contents =
                []() {
                  absl::StatusOr<std::string> png =
                      CreatePng({}, /*add_header=*/true, /*add_end=*/false);
                  CHECK_OK(png);
                  return *png;
                }(),
            .result = absl::NotFoundError("No manifest store found"),
        },
        TestCase{
            .name = "EmptyPng",
            .contents =
                []() {
                  absl::StatusOr<std::string> png = CreatePng({});
                  CHECK_OK(png);
                  return *png;
                }(),
            .result = absl::NotFoundError("No manifest store found"),
        },
        TestCase{
            .name = "NonC2paInvalidCrcIsIgnored",
            .contents =
                []() {
                  absl::StatusOr<std::string> png =
                      CreatePng({{.type = 'IEND', .payload = ""}},
                                /*add_header=*/true, /*add_end=*/false);
                  CHECK_OK(png);
                  return *png;
                }(),
            .result = absl::NotFoundError("No manifest store found"),
        },
        TestCase{
            .name = "C2paFound",
            .contents =
                []() {
                  absl::StatusOr<std::string> png = CreatePng({
                      {.type = 'IDAT', .payload = "ignored IDAT chunk"},
                      {.type = 'caBX',
                       .payload = "A C2PA manifest store",
                       .crc = 0xdb50c4b5},
                  });
                  CHECK_OK(png);
                  return *png;
                }(),
            .result =
                ExtractorResult{
                    .manifest_store = "A C2PA manifest store",
                    .asset_byte_info =
                        {
                            .manifest_store_location = {.offset = 38,
                                                        .length = 33},
                            .boxes = CreateBoxes({{"PNGh", /*size=*/8},
                                                  {"IDAT", /*size=*/30},
                                                  {"C2PA", /*size=*/33},
                                                  {"IEND", /*size=*/12}}),
                        },
                },
        },
        TestCase{
            .name = "TruncatedChunk",
            .contents =
                []() {
                  absl::StatusOr<std::string> data = CreatePng(
                      {{.type = 'IDAT', .payload = "ignored IDAT chunk"}},
                      /*add_header=*/true, /*add_end=*/false);
                  CHECK_OK(data);
                  data->resize(data->size() - 1);
                  return *data;
                }(),
            .result = absl::InvalidArgumentError(
                "PNG chunk extends beyond the end of the file"),
        },
        TestCase{
            .name = "C2paInvalidCrc",
            .contents =
                []() {
                  absl::StatusOr<std::string> png = CreatePng(
                      {{.type = 'caBX', .payload = "", .crc = 0x11111111}});
                  CHECK_OK(png);
                  return *png;
                }(),
            .result = absl::InvalidArgumentError("does not match computed CRC"),
        },
        TestCase{
            .name = "C2paChunkTooLarge",
            .contents =
                []() {
                  absl::StatusOr<std::string> png = CreatePng({
                      {.type = 'caBX',
                       .payload = std::string(10 * 1024 * 1024 + 1, 'a')},
                  });
                  CHECK_OK(png);
                  return *png;
                }(),
            .result = absl::InvalidArgumentError(
                "PNG C2PA chunk is too large to extract"),
        },
        TestCase{
            .name = "MultipleC2paChunks",
            .contents =
                []() {
                  absl::StatusOr<std::string> png = CreatePng({
                      {.type = 'caBX', .payload = "", .crc = 0xe63dd2a7},
                      {.type = 'caBX', .payload = "", .crc = 0xe63dd2a7},
                  });
                  CHECK_OK(png);
                  return *png;
                }(),
            .result = absl::NotFoundError("Multiple manifest stores found"),
        },
        TestCase{
            .name = "LargeIDATChunk",
            .contents =
                []() {
                  absl::StatusOr<std::string> png = CreatePng({
                      {.type = 'IDAT',
                       .payload = "this_is_a_large_idat_chunk",
                       .crc = 2230102072},
                  });
                  CHECK_OK(png);
                  return *png;
                }(),
            .result = absl::NotFoundError("No manifest store found"),
        },
        TestCase{
            .name = "ContainsAdditionalData",
            .contents =
                []() {
                  absl::StatusOr<std::string> png = CreatePng({
                      {.type = 'IDAT', .payload = "ignored IDAT chunk"},
                      {.type = 'caBX',
                       .payload = "A C2PA manifest store",
                       .crc = 0xdb50c4b5},
                  });
                  CHECK_OK(png);
                  return *png + *png;
                }(),
            .result =
                ExtractorResult{
                    .manifest_store = "A C2PA manifest store",
                    .asset_byte_info =
                        {
                            .manifest_store_location = {.offset = 38,
                                                        .length = 33},
                            .boxes = CreateBoxes({{"PNGh", /*size=*/8},
                                                  {"IDAT", /*size=*/30},
                                                  {"C2PA", /*size=*/33},
                                                  {"IEND", /*size=*/12},
                                                  {"c2pa.after", /*size=*/83}}),
                        },
                },
        },
    }),
    [](const testing::TestParamInfo<ExtractorTest::ParamType>& info) {
      return info.param.name;
    });

TEST_P(ExtractorTest, Extract) {
  const TestCase& test_case = GetParam();
  riegeli::StringReader<> input(test_case.contents);
  absl::StatusOr<std::string> result =
      PngExtractor().ExtractManifestStore(input);

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    EXPECT_THAT(result, IsOkAndHolds(expected_result->manifest_store));
  } else {
    EXPECT_THAT(result,
                StatusIs(expected_result.status().code(),
                         HasSubstr(expected_result.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractWorksWhenPrefixPadded) {
  const TestCase& test_case = GetParam();
  std::string contents = "padding" + test_case.contents;
  riegeli::StringReader<> input(contents);

  ASSERT_TRUE(input.Seek(7));
  absl::StatusOr<std::string> result =
      PngExtractor().ExtractManifestStore(input);

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    EXPECT_THAT(result, IsOkAndHolds(expected_result->manifest_store));
  } else {
    EXPECT_THAT(result,
                StatusIs(expected_result.status().code(),
                         HasSubstr(expected_result.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractManifestStoreLocationWithSuffixPadded) {
  const TestCase& test_case = GetParam();
  std::string contents = test_case.contents + "padding";
  riegeli::StringReader<> input(contents);

  absl::StatusOr<std::optional<ByteRange>> result =
      PngExtractor().ExtractManifestStoreLocation(
          input, {.end_offset = static_cast<int64_t>(contents.size()) - 7});

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    EXPECT_THAT(
        result,
        IsOkAndHolds(expected_result->asset_byte_info.manifest_store_location));
  } else {
    EXPECT_THAT(result,
                StatusIs(expected_result.status().code(),
                         HasSubstr(expected_result.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractBoxesWithSuffixPadded) {
  const TestCase& test_case = GetParam();
  std::string contents = test_case.contents + "padding";
  riegeli::StringReader<> input(contents);

  absl::StatusOr<std::vector<AssetBox>> result = PngExtractor().ExtractBoxes(
      input, {.end_offset = static_cast<int64_t>(contents.size()) - 7});

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    EXPECT_THAT(result, IsOkAndHolds(expected_result->asset_byte_info.boxes));
  } else {
    EXPECT_THAT(result,
                StatusIs(expected_result.status().code(),
                         HasSubstr(expected_result.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractManifestStoreLocationRequiresC2pa) {
  const TestCase& test_case = GetParam();
  riegeli::StringReader<> input(test_case.contents);

  absl::StatusOr<std::optional<ByteRange>> result =
      PngExtractor().ExtractManifestStoreLocation(input, {});

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    EXPECT_THAT(
        result,
        IsOkAndHolds(expected_result->asset_byte_info.manifest_store_location));
  } else {
    EXPECT_THAT(result,
                StatusIs(expected_result.status().code(),
                         HasSubstr(expected_result.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractManifestStoreLocationDoesNotRequireC2pa) {
  const TestCase& test_case = GetParam();
  riegeli::StringReader<> input(test_case.contents);

  absl::StatusOr<std::optional<ByteRange>> result =
      PngExtractor().ExtractManifestStoreLocation(input,
                                                  {.requires_c2pa = false});

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    EXPECT_THAT(
        result,
        IsOkAndHolds(expected_result->asset_byte_info.manifest_store_location));
  } else if (expected_result.status().code() == absl::StatusCode::kNotFound &&
             expected_result.status().message() !=
                 "Multiple manifest stores found") {
    EXPECT_THAT(result, IsOkAndHolds(std::nullopt));
  } else {
    EXPECT_THAT(result,
                StatusIs(expected_result.status().code(),
                         HasSubstr(expected_result.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractBoxes) {
  const TestCase& test_case = GetParam();
  riegeli::StringReader<> input(test_case.contents);

  absl::StatusOr<std::vector<AssetBox>> result =
      PngExtractor().ExtractBoxes(input, {});

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    EXPECT_THAT(result, IsOkAndHolds(expected_result->asset_byte_info.boxes));
  } else {
    EXPECT_THAT(result,
                StatusIs(expected_result.status().code(),
                         HasSubstr(expected_result.status().message())));
  }
}

TEST(PngExtractorTest, IsManifestStore) {
  EXPECT_TRUE(credentio::PngExtractor().MightBeC2paManifestStore(
      credentio::CreateStartOfManifestStorePayload()));
}

TEST(PngExtractorTest, ExtractManifestStoreLocationWithEndOffsetBeforeC2pa) {
  absl::StatusOr<std::string> png = CreatePng({
      {.type = 'IDAT', .payload = "ignored IDAT chunk"},
      {.type = 'caBX', .payload = "A C2PA manifest store", .crc = 0xdb50c4b5},
  });
  ASSERT_THAT(png, ::absl_testing::IsOk());
  riegeli::StringReader<> input(*png);

  // Set end_offset to 38 (right before caBX).
  // This should cause caBX to be considered past the asset window,
  // triggering the abort early return in ValidateWithinAssetWindow.
  auto result = PngExtractor().ExtractManifestStoreLocation(
      input, {.requires_c2pa = false, .end_offset = 38});
  EXPECT_THAT(result, IsOkAndHolds(std::nullopt));
}

TEST(PngExtractorTest, ExtractBoxesDoesNotRequireC2paAndAssetTooSmall) {
  riegeli::StringReader<> input("");
  auto result = PngExtractor().ExtractBoxes(input, {.requires_c2pa = false});
  EXPECT_THAT(result, IsOkAndHolds(IsEmpty()));
}

TEST(PngExtractorTest,
     ExtractManifestStoreLocationAssetJustTooSmallDoesNotRequireC2pa) {
  // kPngMinimumAssetSize is 16. We use 15 bytes.
  std::string contents = std::string(
      "\x89\x50\x4e\x47\x0d\x0a\x1a\x0a\x00\x00\x00\x00\x00\x00\x00", 15);
  riegeli::StringReader<> input(contents);
  auto result = PngExtractor().ExtractManifestStoreLocation(
      input, {.requires_c2pa = false});
  EXPECT_THAT(result, IsOkAndHolds(std::nullopt));
}

}  // namespace
}  // namespace credentio