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

#include <sys/types.h>

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

#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 "formats/asset_box.h"
#include "formats/byte_range.h"
#include "formats/extractor_result.h"
#include "formats/jpeg/testing/app_segment_creator.h"
#include "formats/jpeg/testing/jpeg_creator.h"
#include "formats/jpeg/testing/jumbf_creator.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::Eq;
using ::testing::HasSubstr;

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 PayloadInput {
  std::string payload;
  JumbfCreatorParams jumbf_creator_params;
  AppSegmentParams app_segment_params;
};

struct TestInput {
  std::string name;
  std::vector<PayloadInput> payload = {};
  JpegParams jpeg_params = {};
  void (*modify_jumbf_boxes)(std::vector<AppSegmentParams>*) = nullptr;
  void (*modify_app_segments)(JpegParams*) = nullptr;

  absl::StatusOr<ExtractorResult> expected =
      absl::InvalidArgumentError("Unsupported hard binding type");
};

class ExtractorTest : public testing::TestWithParam<TestInput> {
 private:
  std::vector<AppSegmentParams> PopulateJumbfBoxes(TestInput input) {
    std::vector<AppSegmentParams> jumbf_boxes;
    jumbf_boxes.reserve(input.payload.size());

    for (const auto& payload : input.payload) {
      JumbfCreatorParams jumbf_params = payload.jumbf_creator_params;
      jumbf_params.payload = payload.payload;

      AppSegmentParams app_segment_params = payload.app_segment_params;
      app_segment_params.payload = CreateJumbf(jumbf_params);

      jumbf_boxes.push_back(app_segment_params);
    }
    if (input.modify_jumbf_boxes != nullptr) {
      input.modify_jumbf_boxes(&jumbf_boxes);
    }
    return jumbf_boxes;
  }

  absl::StatusOr<JpegParams> PopulateAppSegments(
      TestInput input, std::vector<AppSegmentParams> jumbf_boxes) {
    JpegParams result = input.jpeg_params;

    result.payload.clear();
    result.payload.reserve(jumbf_boxes.size());

    uint16_t en = 0;
    for (auto& jumbf_box : jumbf_boxes) {
      if (!jumbf_box.en.has_value()) {
        jumbf_box.en = ++en;
      }

      ABSL_ASSIGN_OR_RETURN(auto segments, CreateAppSegments(jumbf_box));
      for (const auto& segment : segments) {
        result.payload.push_back(segment);
      }
    }

    if (input.modify_app_segments != nullptr) {
      input.modify_app_segments(&result);
    }
    return result;
  }

 protected:
  absl::StatusOr<std::string> CreateImage(TestInput input) {
    ABSL_ASSIGN_OR_RETURN(
        JpegParams params,
        PopulateAppSegments(input, PopulateJumbfBoxes(input)));
    return CreateJpeg(params);
  }
};

TEST_P(ExtractorTest, Extract) {
  absl::StatusOr<std::string> image = CreateImage(GetParam());
  ABSL_ASSERT_OK(image);
  riegeli::StringReader input(*image);

  absl::StatusOr<std::string> extractor_result =
      JpegExtractor().ExtractManifestStore(input);

  const absl::StatusOr<ExtractorResult>& expected = GetParam().expected;
  if (expected.ok()) {
    EXPECT_THAT(extractor_result, IsOkAndHolds(Eq(expected->manifest_store)));
  } else {
    EXPECT_THAT(extractor_result,
                StatusIs(expected.status().code(),
                         HasSubstr(expected.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractWorksWhenPrefixPadded) {
  absl::StatusOr<std::string> image = CreateImage(GetParam());
  ABSL_ASSERT_OK(image);
  std::string contents = "padding" + *image;
  riegeli::StringReader input(contents);

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

  const absl::StatusOr<ExtractorResult>& expected = GetParam().expected;
  if (expected.ok()) {
    EXPECT_THAT(extractor_result, IsOkAndHolds(Eq(expected->manifest_store)));
  } else {
    EXPECT_THAT(extractor_result,
                StatusIs(expected.status().code(),
                         HasSubstr(expected.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractManifestStoreLocation) {
  absl::StatusOr<std::string> image = CreateImage(GetParam());
  ABSL_ASSERT_OK(image);
  riegeli::StringReader input(*image);

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

  const absl::StatusOr<ExtractorResult>& expected = GetParam().expected;
  if (expected.ok()) {
    EXPECT_THAT(
        result,
        IsOkAndHolds(Eq(expected->asset_byte_info.manifest_store_location)));
  } else {
    EXPECT_THAT(result, StatusIs(expected.status().code(),
                                 HasSubstr(expected.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractBoxes) {
  absl::StatusOr<std::string> image = CreateImage(GetParam());
  ABSL_ASSERT_OK(image);
  riegeli::StringReader input(*image);

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

  const absl::StatusOr<ExtractorResult>& expected = GetParam().expected;
  if (expected.ok()) {
    EXPECT_THAT(result, IsOkAndHolds(Eq(expected->asset_byte_info.boxes)));
  } else {
    EXPECT_THAT(result, StatusIs(expected.status().code(),
                                 HasSubstr(expected.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractManifestStoreLocationWithSuffix) {
  absl::StatusOr<std::string> image = CreateImage(GetParam());
  ABSL_ASSERT_OK(image);
  std::string contents = *image + "suffix";
  riegeli::StringReader input(contents);

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

  const absl::StatusOr<ExtractorResult>& expected = GetParam().expected;
  if (expected.ok()) {
    EXPECT_THAT(
        result,
        IsOkAndHolds(Eq(expected->asset_byte_info.manifest_store_location)));
  } else {
    EXPECT_THAT(result, StatusIs(expected.status().code(),
                                 HasSubstr(expected.status().message())));
  }
}

TEST_P(ExtractorTest, ExtractBoxesWithSuffix) {
  absl::StatusOr<std::string> image = CreateImage(GetParam());
  ABSL_ASSERT_OK(image);
  std::string contents = *image + "suffix";
  riegeli::StringReader input(contents);

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

  const absl::StatusOr<ExtractorResult>& expected = GetParam().expected;
  if (expected.ok()) {
    EXPECT_THAT(result, IsOkAndHolds(Eq(expected->asset_byte_info.boxes)));
  } else {
    EXPECT_THAT(result, StatusIs(expected.status().code(),
                                 HasSubstr(expected.status().message())));
  }
}

INSTANTIATE_TEST_SUITE_P(
    ImageEncoding, ExtractorTest,
    testing::Values<TestInput>(
        TestInput{
            .name = "InvalidBeginning",
            .payload = {},
            .jpeg_params = {.start_marker = 0x0102},
            .expected = absl::InvalidArgumentError("Invalid JPEG marker: 258"),
        },
        TestInput{
            .name = "InvalidEnding",
            .payload = {},
            .jpeg_params = {.end_markers = 0x0304},
            .expected = absl::InvalidArgumentError("Invalid JPEG marker: 772"),
        },
        TestInput{
            .name = "NoSegments",
            .payload = {},
            .jpeg_params = {},
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "NoApp11Segments",
            .payload = {{.payload = "test",
                         .app_segment_params = {.marker = 0xffe2}}},
            .jpeg_params = {},
            .expected = absl::NotFoundError("No manifest store found"),
        }),
    [](const testing::TestParamInfo<ExtractorTest::ParamType>& info) {
      return info.param.name;
    });

INSTANTIATE_TEST_SUITE_P(
    Segments, ExtractorTest,
    testing::Values<TestInput>(
        TestInput{
            .name = "OnlyApp2Segment",
            .payload = {{
                .payload = "test",
                .app_segment_params = {.marker = 0xffe2},
            }},
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "IncorrectJpegExtensionType",
            .payload = {{
                .payload = "test",
                .app_segment_params = {.cl = 0101},
            }},
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "InvalidSegmentId",
            .payload = {{
                .payload = "test",
                .app_segment_params = {.en = 0x0000},
            }},
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "SegmentsOutOfOrder",
            .payload =
                {
                    {
                        .payload =
                            "This is a much longer string of test to ensure "
                            "that there is enough data to expect a C2PA "
                            "manifest store within the segment payload.",
                        .app_segment_params = {.max_size = 150},
                    },
                },
            .modify_app_segments =
                [](JpegParams* params) {
                  ASSERT_EQ(params->payload.size(), 2);

                  auto temp = params->payload[0];
                  params->payload[0] = params->payload[1];
                  params->payload[1] = temp;
                },
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "IncompleteSegment",
            .payload =
                {
                    {
                        .payload = "longer test",
                        .app_segment_params = {.max_size = 50},
                    },
                },
            .modify_app_segments =
                [](JpegParams* params) {
                  ASSERT_EQ(params->payload.size(), 2);
                  params->payload.pop_back();
                },
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "SplitSegment",
            .payload =
                {
                    {
                        .payload = "longer test",
                        .app_segment_params = {.max_size = 50},
                    },
                    {
                        .payload = "test",
                        .app_segment_params = {.marker = 0xffe6},
                    },
                },
            .modify_app_segments =
                [](JpegParams* params) {
                  ASSERT_EQ(params->payload.size(), 3);
                  auto temp = params->payload[1];
                  params->payload[1] = params->payload[2];
                  params->payload[2] = temp;
                },
            .expected = absl::NotFoundError("No manifest store found"),
        }),
    [](const testing::TestParamInfo<ExtractorTest::ParamType>& info) {
      return info.param.name;
    });

INSTANTIATE_TEST_SUITE_P(
    JumbfBoxes, ExtractorTest,
    testing::Values<TestInput>(
        TestInput{
            .name = "TooFewBytesToDetermineIfC2pa",
            .payload = {{.payload = "test"}},
            .modify_jumbf_boxes =
                [](std::vector<AppSegmentParams>* params) {
                  ASSERT_EQ(params->size(), 1);
                  params->at(0).payload = params->at(0).payload.substr(0, 30);
                },
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "IncorrectSuperBoxTbox",
            .payload = {{
                .payload = "test",
                .jumbf_creator_params = {.super_box_identifier = 0x12345678},
            }},
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "InvalidUuidHigh",
            .payload = {{
                .payload = "test",
                .jumbf_creator_params = {.uuid_high = 0x0000000000000000},
            }},
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "InvalidUuidLow",
            .payload = {{
                .payload = "test",
                .jumbf_creator_params = {.uuid_low = 0x0000000000000000},
            }},
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "InvalidToggles",
            .payload = {{
                .payload = "test",
                .jumbf_creator_params = {.toggles = 0b00110000},
            }},
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "InvalidLabel",
            .payload = {{
                .payload = "test",
                .jumbf_creator_params = {.label = "demo"},
            }},
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "ForceSizeIntoXlBox",
            .payload = {{
                .payload = "test",
                .jumbf_creator_params = {.force_xlbox = true},
            }},
            .expected =
                ExtractorResult{
                    .manifest_store = CreateJumbf({.payload = "test",
                                                   .force_xlbox = true}),
                    .asset_byte_info =
                        {
                            .manifest_store_location = ByteRange{.offset = 2,
                                                                 .length = 78},
                            .boxes = CreateBoxes({
                                {"SOI", /*size=*/2},
                                {"C2PA", /*size=*/78},
                                {"EOI", /*size=*/2},
                            }),
                        },
                },
        },
        TestInput{
            .name = "IncorrectLboxSize",
            .payload = {{
                .payload = "longer test",
                .app_segment_params = {.max_size = 50},
            }},
            .modify_app_segments =
                [](JpegParams* params) {
                  ASSERT_EQ(params->payload.size(), 2);
                  params->payload.pop_back();
                },
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "IncorrectXlboxSize",
            .payload = {{
                .payload = "This is a much longer string of test to ensure "
                           "that there is enough data to expect a C2PA "
                           "manifest store within the segment payload.",
                .jumbf_creator_params = {.force_xlbox = true},
                .app_segment_params = {.max_size = 100},
            }},
            .modify_app_segments =
                [](JpegParams* params) {
                  ASSERT_EQ(params->payload.size(), 3);
                  params->payload.pop_back();
                },
            .expected = absl::NotFoundError("No manifest store found"),
        }),
    [](const testing::TestParamInfo<ExtractorTest::ParamType>& info) {
      return info.param.name;
    });

INSTANTIATE_TEST_SUITE_P(
    C2paManifestStore, ExtractorTest,
    testing::Values<TestInput>(
        TestInput{
            .name = "OnlyManifestStore",
            .payload = {{.payload = "test"}},
            .expected =
                ExtractorResult{
                    .manifest_store = CreateJumbf({.payload = "test"}),
                    .asset_byte_info =
                        {
                            .manifest_store_location = ByteRange{.offset = 2,
                                                                 .length = 62},
                            .boxes = CreateBoxes({
                                {"SOI", /*size=*/2},
                                {"C2PA", /*size=*/62},
                                {"EOI", /*size=*/2},
                            }),
                        },
                },
        },
        TestInput{
            .name = "SandwichedManifestStore",
            .payload =
                {
                    {
                        .payload = "app3",
                        .app_segment_params = {.marker = 0xffe3},
                    },
                    {.payload = "test"},
                    {
                        .payload = "app6",
                        .app_segment_params = {.marker = 0xffe6},
                    },
                },
            .expected =
                ExtractorResult{
                    .manifest_store = CreateJumbf({.payload = "test"}),
                    .asset_byte_info =
                        {
                            .manifest_store_location = ByteRange{.offset = 64,
                                                                 .length = 62},
                            .boxes = CreateBoxes({
                                {"SOI", /*size=*/2},
                                {"APP3", /*size=*/62},
                                {"C2PA", /*size=*/62},
                                {"APP6", /*size=*/62},
                                {"EOI", /*size=*/2},
                            }),
                        },
                },
        },
        TestInput{
            .name = "ManifestStoreAcrossTwoSegments",
            .payload =
                {
                    {
                        .payload = "longer test",
                        .app_segment_params = {.max_size = 50},
                    },
                },
            .expected =
                ExtractorResult{
                    .manifest_store = CreateJumbf({.payload = "longer test"}),
                    .asset_byte_info =
                        {
                            .manifest_store_location = ByteRange{.offset = 2,
                                                                 .length = 89},
                            .boxes = CreateBoxes({
                                {"SOI", /*size=*/2},
                                {"C2PA", /*size=*/89},
                                {"EOI", /*size=*/2},
                            }),
                        },
                },
        },
        TestInput{
            .name = "ManifestStoreAcrossTwoSegmentsDifferentTbox",
            .payload =
                {
                    {
                        .payload = "longer test",
                        .app_segment_params = {.max_size = 50},
                    },
                },
            .modify_app_segments =
                [](JpegParams* params) {
                  ASSERT_EQ(params->payload.size(), 2);
                  // lowest byte of tbox is at 19, normally 0x62, change it
                  params->payload[1][19] = 0x69;
                },
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "ManifestStoreAcrossTwoSegmentsDifferentLbox",
            .payload =
                {
                    {
                        .payload = "longer test",
                        .app_segment_params = {.max_size = 50},
                    },
                },
            .modify_app_segments =
                [](JpegParams* params) {
                  ASSERT_EQ(params->payload.size(), 2);
                  // lowest byte of lbox is at 15, normally 0x39, change it
                  params->payload[1][15] = 0x47;
                },
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "ManifestStoreAcrossTwoSegmentsDifferentXlbox",
            .payload =
                {
                    {
                        .payload = "longer test",
                        .jumbf_creator_params = {.force_xlbox = true},
                        .app_segment_params = {.max_size = 60},
                    },
                },
            .modify_app_segments =
                [](JpegParams* params) {
                  ASSERT_EQ(params->payload.size(), 2);
                  // lowest byte of xlbox is at 27, normally 0x49, change it
                  params->payload[1][27] = 0x57;
                },
            .expected = absl::NotFoundError("No manifest store found"),
        },
        TestInput{
            .name = "TwoManifestStores",
            .payload =
                {
                    {.payload = "store 1"},
                    {.payload = "store 2"},
                },
            .expected = absl::NotFoundError("Multiple manifest stores found"),
        },
        TestInput{
            .name = "TwoManifestStoresThenOthers",
            .payload =
                {{.payload = "store 1"},
                 {.payload = "store 2"},
                 {.payload = "app8", .app_segment_params = {.marker = 0xffe8}},
                 {.payload = "app9", .app_segment_params = {.marker = 0xffe9}}},
            .expected = absl::NotFoundError("Multiple manifest stores found"),
        },
        // If there are multiple manifest stores, but one of them is invalid
        // JUMBF we should still return the valid one.
        TestInput{
            .name = "TwoManifestStoresOneInvalid",
            .payload =
                {
                    {.payload = "store 1"},
                    {
                        .payload = "store 2",
                        .app_segment_params = {.max_size = 50},
                    },
                },
            .modify_app_segments =
                [](JpegParams* params) {
                  // Remove the last segment of the second manifest store to
                  // invalidate it
                  ASSERT_EQ(params->payload.size(), 3);
                  params->payload.pop_back();
                },
            .expected =
                ExtractorResult{
                    .manifest_store = CreateJumbf({.payload = "store 1"}),
                    .asset_byte_info =
                        {
                            .manifest_store_location = ByteRange{.offset = 2,
                                                                 .length = 65},
                            .boxes = CreateBoxes({
                                {"SOI", /*size=*/2},
                                {"C2PA", /*size=*/65},
                                {"APP11", /*size=*/52},
                                {"EOI", /*size=*/2},
                            }),
                        },
                },
        },
        TestInput{
            .name = "MultiSegmentC2paAndMultiSegmentMiscApp11",
            .payload =
                {
                    {
                        .payload =
                            "This is a much longer string of test to ensure "
                            "that there is enough data to expect a C2PA "
                            "manifest store within the segment payload.",
                        .app_segment_params = {.max_size = 100},
                    },
                    {
                        .payload =
                            "This is a much longer string of test to ensure "
                            "that there is enough data to expect a C2PA "
                            "manifest store within the segment payload.",
                        .app_segment_params = {.cl = 0x1111, .max_size = 100},
                    },
                },
            .expected =
                ExtractorResult{
                    .manifest_store = CreateJumbf(
                        {.payload =
                             "This is a much longer string of test to ensure "
                             "that there is enough data to expect a C2PA "
                             "manifest store within the segment payload."}),
                    .asset_byte_info =
                        {
                            .manifest_store_location = ByteRange{.offset = 2,
                                                                 .length = 230},
                            .boxes = CreateBoxes({
                                {"SOI", /*size=*/2},
                                {"C2PA", /*size=*/230},
                                {"APP11", /*size=*/102},
                                {"APP11", /*size=*/102},
                                {"APP11", /*size=*/26},
                                {"EOI", /*size=*/2},
                            }),
                        },
                },
        }),
    [](const testing::TestParamInfo<ExtractorTest::ParamType>& info) {
      return info.param.name;
    });

TEST(JpegExtractorTest, IsManifestStorePayloadMissingJpegElements) {
  EXPECT_FALSE(JpegExtractor().MightBeC2paManifestStore(
      CreateStartOfManifestStorePayload()));
}

TEST(JpegExtractorTest, IsManifestStoreIncorrectExtensionType) {
  EXPECT_FALSE(JpegExtractor().MightBeC2paManifestStore(
      CreateStartOfManifestStorePayload(
          {0x12, 0x34, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01})));
}

TEST(JpegExtractorTest, IsManifestStore) {
  EXPECT_TRUE(JpegExtractor().MightBeC2paManifestStore(
      CreateStartOfManifestStorePayload(
          {0x4a, 0x50, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01})));
}

}  // namespace
}  // namespace credentio