// 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/riff/extractor.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/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "formats/asset_box.h"
#include "formats/byte_range.h"
#include "formats/extractor_result.h"
#include "formats/riff/create_riff.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 ::credentio_riff::Chunk;
using ::credentio_riff::ListChunk;
using ::credentio_riff::RiffChunk;
using ::testing::ElementsAre;
using ::testing::HasSubstr;

constexpr absl::string_view kManifestStore = "1234567890";

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;
};

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 = "EmptyRiffChunk",
            .contents = RiffChunk("ABCD", ""),
            .result = absl::NotFoundError("No manifest store found"),
        },
        TestCase{
            .name = "VariousChunksNoC2pa",
            .contents = RiffChunk(
                "WEBP",
                absl::StrCat(Chunk("ABCD", "some data"), Chunk("HIJK", "etc"),
                             Chunk("DEFG", "more irrelevant data"))),
            .result = absl::NotFoundError("No manifest store found"),
        },
        TestCase{
            .name = "C2paChunkTooLarge",
            .contents = RiffChunk(
                "WEBP",
                absl::StrCat(Chunk("ABCD", std::string(100000, 'X')),
                             Chunk("C2PA",
                                   std::string((1024 * 1024 * 10) + 1, 'Z')))),
            .result = absl::InvalidArgumentError(
                "RIFF C2PA chunk is too large to extract"),
        },
        TestCase{
            .name = "SingleC2paChunk",
            .contents = RiffChunk("WEBP", Chunk("C2PA", kManifestStore)),
            .result =
                ExtractorResult{
                    .manifest_store = std::string(kManifestStore),
                    .asset_byte_info =
                        {
                            .manifest_store_location = ByteRange{.offset = 12,
                                                                 .length = 18},
                            .boxes = CreateBoxes({{"RIFF", /*size=*/12},
                                                  {"C2PA", /*size=*/18}}),
                        },
                },
        },
        TestCase{
            .name = "VariousChunksIncludingC2pa",
            .contents =
                RiffChunk(
                    "WEBP",
                    absl::StrCat(
                        Chunk("ABCD",
                              "123456789"),  // odd length, will have pad byte
                        Chunk("EFGH", std::string(100000, 'Z')),  // large chunk
                        Chunk("C2PA", kManifestStore), Chunk("DEFG", "12"))),
            .result =
                ExtractorResult{
                    .manifest_store = std::string(kManifestStore),
                    .asset_byte_info =
                        {
                            .manifest_store_location =
                                ByteRange{.offset = 100038, .length = 18},
                            .boxes = CreateBoxes({{"RIFF", /*size=*/12},
                                                  {"ABCD", /*size=*/18},
                                                  {"EFGH", /*size=*/100008},
                                                  {"C2PA", /*size=*/18},
                                                  {"DEFG", /*size=*/10}}),
                        }},
        },
        TestCase{
            .name = "ContainingListChunk",
            .contents = RiffChunk(
                "WEBP", absl::StrCat(Chunk("ABCD", "1234"),
                                     ListChunk("DATA", Chunk("HIJK", "1234")),
                                     Chunk("DEFG", "1234"),
                                     Chunk("C2PA", kManifestStore))),
            .result =
                ExtractorResult{
                    .manifest_store = std::string(kManifestStore),
                    .asset_byte_info =
                        {
                            .manifest_store_location = ByteRange{.offset = 60,
                                                                 .length = 18},
                            .boxes = CreateBoxes({{"RIFF", /*size=*/12},
                                                  {"ABCD", /*size=*/12},
                                                  {"LIST", /*size=*/24},
                                                  {"DEFG", /*size=*/12},
                                                  {"C2PA", /*size=*/18}}),
                        }},
        },
    }),
    [](const testing::TestParamInfo<ExtractorTest::ParamType>& info) {
      return info.param.name;
    });

TEST(ExtractorOverflowTest, DeeplyNestedRiffWithoutC2pa) {
  std::string contents = Chunk("NMBR", "1234");
  for (int32_t i = 0; i < 200; ++i) {
    contents = ListChunk("LAYR", contents);
  }
  contents = RiffChunk("WEBP", contents);

  riegeli::StringReader<> input(contents);

  EXPECT_THAT(RiffExtractor().ExtractManifestStore(input),
              StatusIs(absl::StatusCode::kNotFound,
                       HasSubstr("No manifest store found")));
}

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

  absl::StatusOr<std::string> result =
      RiffExtractor().ExtractManifestStore(input);

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    ABSL_ASSERT_OK(result);
    EXPECT_EQ(*result, expected_result->manifest_store);
  } else {
    EXPECT_THAT(result.status(),
                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 =
      RiffExtractor().ExtractManifestStore(input);

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

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

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

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

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

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

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    ABSL_ASSERT_OK(result);
    EXPECT_THAT(*result, expected_result->asset_byte_info.boxes);
  } else {
    EXPECT_THAT(result.status(),
                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 + "padding8";
  riegeli::StringReader<> input(contents);

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

  const auto& expected_result = test_case.result;
  if (expected_result.ok()) {
    ABSL_ASSERT_OK(result);
    EXPECT_EQ(*result,
              expected_result->asset_byte_info.manifest_store_location);
  } else {
    EXPECT_THAT(result.status(),
                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 + "padding8";
  riegeli::StringReader<> input(contents);

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

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

TEST(RiffExtractorTest, IsManifestStore) {
  EXPECT_TRUE(RiffExtractor().MightBeC2paManifestStore(
      CreateStartOfManifestStorePayload()));
}

TEST(RiffExtractorTest, ExtractManifestStoreMultipleRiffChunks) {
  std::string contents =
      absl::StrCat(RiffChunk("WEBP", Chunk("C2PA", kManifestStore)),
                   RiffChunk("BLAH", Chunk("ABCD", "second riff chunk data")));
  riegeli::StringReader<> input(contents);

  EXPECT_THAT(RiffExtractor().ExtractManifestStore(input),
              IsOkAndHolds(kManifestStore));
}

TEST(RiffExtractorTest, ExtractManifestStoreC2paInSecondRiffChunk) {
  std::string contents =
      absl::StrCat(RiffChunk("WEBP", Chunk("ABCD", "first riff chunk data")),
                   RiffChunk("BLAH", Chunk("C2PA", kManifestStore)));
  riegeli::StringReader<> input(contents);

  EXPECT_THAT(RiffExtractor().ExtractManifestStore(input),
              StatusIs(absl::StatusCode::kNotFound, "No manifest store found"));
}

TEST(RiffExtractorTest, ExtractBoxesNonRiffTopLevelChunk) {
  std::string contents =
      absl::StrCat(RiffChunk("WEBP", Chunk("ABCD", "first riff chunk data")),
                   Chunk("BLAH", "non-riff top-level chunk data"));
  riegeli::StringReader<> input(contents);

  auto boxes_or = RiffExtractor().ExtractBoxes(input, {.requires_c2pa = false});
  ABSL_ASSERT_OK(boxes_or.status());
  std::vector<AssetBox> boxes = *boxes_or;

  std::vector<std::string> box_names;
  for (const auto& box : boxes) {
    box_names.push_back(box.identifier);
  }
  EXPECT_THAT(box_names, ElementsAre("RIFF", "ABCD", "BLAH"));
}

TEST(RiffExtractorTest, MultipleC2paChunks) {
  std::string contents =
      RiffChunk("WEBP", absl::StrCat(Chunk("C2PA", kManifestStore),
                                     Chunk("C2PA", kManifestStore)));
  riegeli::StringReader<> input(contents);

  EXPECT_THAT(
      RiffExtractor().ExtractBoxes(input, {.requires_c2pa = true}),
      StatusIs(absl::StatusCode::kNotFound, "Multiple manifest stores found"));
}

TEST(RiffExtractorTest, ExtractManifestStoreDoesNotStartWithRiff) {
  std::string contents = Chunk("ABCD", "some data");
  riegeli::StringReader<> input(contents);
  EXPECT_THAT(RiffExtractor().ExtractManifestStore(input),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       "Input does not start with RIFF chunk"));
}

TEST(RiffExtractorTest, ExtractBoxesDoesNotStartWithRiff) {
  std::string contents = Chunk("ABCD", "some data");
  riegeli::StringReader<> input(contents);
  EXPECT_THAT(
      RiffExtractor().ExtractBoxes(input, {}),
      StatusIs(absl::StatusCode::kInvalidArgument, "RIFF chunk not found"));
}

TEST(RiffExtractorTest, ExtractBoxesMultipleRiffChunks) {
  std::string contents =
      absl::StrCat(RiffChunk("WEBP", Chunk("ABCD", "first riff chunk data")),
                   RiffChunk("BLAH", Chunk("EFGH", "second riff chunk data")));
  riegeli::StringReader<> input(contents);

  auto boxes_or = RiffExtractor().ExtractBoxes(input, {.requires_c2pa = false});
  ABSL_ASSERT_OK(boxes_or.status());
  std::vector<AssetBox> boxes = *boxes_or;

  std::vector<std::string> box_names;
  for (const auto& box : boxes) {
    box_names.push_back(box.identifier);
  }
  EXPECT_THAT(box_names, ElementsAre("RIFF", "ABCD", "RIFF", "EFGH"));
}

}  // namespace
}  // namespace credentio