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

#include <cstdint>
#include <vector>

#include "absl/container/btree_set.h"
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "formats/jpeg/reader.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "jumbf/constants.h"
#include "riegeli/bytes/string_reader.h"

namespace credentio {
namespace {

using ::absl_testing::StatusIs;
using ::testing::ElementsAre;
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::SizeIs;

class TestC2paSegmentHandler : public C2paSegmentHandler {
 public:
  std::vector<std::vector<JpegSegment>> AllRecorded() const {
    absl::btree_set<uint16_t> segment_ids;
    for (const auto& [segment_id, segments] : GetSegmentMap()) {
      segment_ids.insert(segment_id);
    }

    std::vector<std::vector<JpegSegment>> result;
    for (uint16_t segment_id : segment_ids) {
      result.push_back(GetSegmentMap().at(segment_id));
    }
    return result;
  }

  absl::btree_set<uint16_t> AllIgnored() const {
    return absl::btree_set<uint16_t>(GetIgnoredSegmentIds().begin(),
                                     GetIgnoredSegmentIds().end());
  }
};

TEST(RecordSegmentTest, DoesNotRecordSegmentWithoutApp11Info) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), IsEmpty());
}

TEST(RecordSegmentTest, DoesNotRecordSegmentIdZero) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(
      JpegSegment{.app11_info = JpegApp11Segment{.segment_id = 0}});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), IsEmpty());
}

TEST(RecordSegmentTest, DoesNotRecordSegmentWithoutJumbfData) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(
      JpegSegment{.app11_info = JpegApp11Segment{.segment_id = 1}});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), IsEmpty());
}

TEST(RecordSegmentTest, DoesNotRecordSegmentIncorrectTbox) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .app11_info = JpegApp11Segment{
          .segment_id = 1, .jumbf_data = JpegJumbfData{.tbox = 0x1234}}});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), IsEmpty());
}

TEST(RecordSegmentTest, RecordSegment) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .app11_info = JpegApp11Segment{
          .segment_id = 1,
          .jumbf_data = JpegJumbfData{.tbox = jumbf::kSuperBoxType}}});
  EXPECT_THAT(
      handler.AllRecorded(),
      ElementsAre(ElementsAre(JpegSegment{
          .app11_info = JpegApp11Segment{
              .segment_id = 1,
              .jumbf_data = JpegJumbfData{.tbox = jumbf::kSuperBoxType}}})));
  EXPECT_THAT(handler.AllIgnored(), IsEmpty());
}

TEST(RecordSegmentTest, DoesNotRecordNextSegmentIncorrectSequence) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .jumbf_data = JpegJumbfData{
                                         .tbox = jumbf::kSuperBoxType,
                                     }}});
  handler.RecordSegment(JpegSegment{
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .jumbf_data = JpegJumbfData{
                                         .tbox = jumbf::kSuperBoxType,
                                     }}});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), ElementsAre(1));
}

TEST(RecordSegmentTest, DoesNotRecordNextSegmentNotContiguous) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .offset = 100,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .jumbf_data = JpegJumbfData{
                                         .tbox = jumbf::kSuperBoxType,
                                     }}});
  handler.RecordSegment(JpegSegment{
      .offset = 200,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 2,
                                     .jumbf_data = JpegJumbfData{
                                         .tbox = jumbf::kSuperBoxType,
                                     }}});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), ElementsAre(1));
}

TEST(RecordSegmentTest, DoesNotRecordNextSegmentDifferentLbox) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .offset = 100,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 10,
                                         .tbox = jumbf::kSuperBoxType,
                                     }}});
  handler.RecordSegment(JpegSegment{
      .offset = 110,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 2,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 20,
                                         .tbox = jumbf::kSuperBoxType,
                                     }}});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), ElementsAre(1));
}

TEST(RecordSegmentTest, DoesNotRecordNextSegmentDifferentXlbox) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .offset = 100,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 1,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 10,
                                     }}});
  handler.RecordSegment(JpegSegment{
      .offset = 110,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 2,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 1,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 20,
                                     }}});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), ElementsAre(1));
}

TEST(RecordSegmentTest, RecordsNextSegment) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .offset = 100,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 10,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 0,
                                     }}});
  handler.RecordSegment(JpegSegment{
      .offset = 110,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 2,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 10,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 0,
                                     }}});
  EXPECT_THAT(handler.AllRecorded(),
              ElementsAre(ElementsAre(
                  JpegSegment{
                      .offset = 100,
                      .length = 10,
                      .app11_info =
                          JpegApp11Segment{.segment_id = 1,
                                           .sequence_number = 1,
                                           .jumbf_data =
                                               JpegJumbfData{
                                                   .lbox = 10,
                                                   .tbox = jumbf::kSuperBoxType,
                                                   .xlbox = 0,
                                               }},
                  },
                  JpegSegment{
                      .offset = 110,
                      .length = 10,
                      .app11_info =
                          JpegApp11Segment{.segment_id = 1,
                                           .sequence_number = 2,
                                           .jumbf_data =
                                               JpegJumbfData{
                                                   .lbox = 10,
                                                   .tbox = jumbf::kSuperBoxType,
                                                   .xlbox = 0,
                                               }},
                  })));
  EXPECT_THAT(handler.AllIgnored(), IsEmpty());
}

TEST(RecordSegmentTest, DoesNotRecordIgnoredSegment) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .offset = 100,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 10,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 0,
                                     }}});
  handler.RecordSegment(JpegSegment{
      .offset = 110,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 10,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 0,
                                     }}});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), ElementsAre(1));

  handler.RecordSegment(JpegSegment{
      .offset = 210,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 10,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 0,
                                     }}});

  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), ElementsAre(1));
}

TEST(RecordSegmentTest, DoesNotRecordMoreThanMaxSegmentsPerId) {
  TestC2paSegmentHandler handler;
  for (int i = 1; i <= 1000; ++i) {
    handler.RecordSegment(JpegSegment{
        .offset = static_cast<uint64_t>((i - 1) * 10),
        .length = 10,
        .app11_info =
            JpegApp11Segment{.segment_id = 1,
                             .sequence_number = static_cast<uint32_t>(i),
                             .jumbf_data = JpegJumbfData{
                                 .lbox = 10,
                                 .tbox = jumbf::kSuperBoxType,
                                 .xlbox = 0,
                             }}});
  }
  EXPECT_THAT(handler.AllRecorded(), SizeIs(1));
  EXPECT_THAT(handler.AllRecorded().front(), SizeIs(1000));
  EXPECT_THAT(handler.AllIgnored(), IsEmpty());

  handler.RecordSegment(JpegSegment{
      .offset = 10000,
      .length = 10,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1001,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 10,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 0,
                                     }}});
  EXPECT_THAT(handler.AllRecorded(), IsEmpty());
  EXPECT_THAT(handler.AllIgnored(), ElementsAre(1));
}

TEST(GetC2paSegmentTest, RejectsDeclaredPayloadSizeTooLarge) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .offset = 0,
      .length = 100,
      .app11_info = JpegApp11Segment{
          .segment_id = 1,
          .sequence_number = 1,
          .payload_offset = 0,
          .payload_length = 10,
          .jumbf_data = JpegJumbfData{
              .lbox = 1024 * 1024 * 10 + 1,  // > kMaxPayloadSize (10 MiB)
              .tbox = jumbf::kSuperBoxType,
              .xlbox = 0,
          }}});
  std::string data(10, 'A');
  riegeli::StringReader input(data);
  EXPECT_THAT(handler.GetC2paSegment(input),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("Declared JUMBF payload size is invalid")));
}

TEST(GetC2paSegmentTest, RejectsDeclaredPayloadSizeZero) {
  TestC2paSegmentHandler handler;
  handler.RecordSegment(JpegSegment{
      .offset = 0,
      .length = 100,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .payload_offset = 0,
                                     .payload_length = 10,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = 0,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 0,
                                     }}});
  std::string data(10, 'A');
  riegeli::StringReader input(data);
  EXPECT_THAT(handler.GetC2paSegment(input),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("Declared JUMBF payload size is invalid")));
}

TEST(GetC2paSegmentTest, RejectsCumulativePayloadSizeExceedingMax) {
  TestC2paSegmentHandler handler;
  const uint64_t kSixMiB = 6 * 1024 * 1024;
  const uint64_t kFiveMiB = 5 * 1024 * 1024;
  const uint64_t kTenMiB = 10 * 1024 * 1024;
  handler.RecordSegment(JpegSegment{
      .offset = 0,
      .length = 100,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 1,
                                     .payload_offset = 0,
                                     .payload_length = kSixMiB,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = kTenMiB,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 0,
                                         .payload_offset = 0,
                                         .payload_length = kSixMiB,
                                     }}});
  handler.RecordSegment(JpegSegment{
      .offset = 100,
      .length = 100,
      .app11_info = JpegApp11Segment{.segment_id = 1,
                                     .sequence_number = 2,
                                     .payload_offset = kSixMiB,
                                     .payload_length = kFiveMiB,
                                     .jumbf_data = JpegJumbfData{
                                         .lbox = kTenMiB,
                                         .tbox = jumbf::kSuperBoxType,
                                         .xlbox = 0,
                                         .payload_offset = kSixMiB,
                                         .payload_length = kFiveMiB,
                                     }}});
  std::string data(kSixMiB, 'A');
  riegeli::StringReader input(data);
  EXPECT_THAT(handler.GetC2paSegment(input),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("JPEG segment is too large to extract")));
}

}  // namespace
}  // namespace credentio