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

#include <sys/types.h>

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

#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "riegeli/bytes/string_reader.h"

namespace credentio {
namespace {

using ::absl_testing::IsOk;
using ::absl_testing::StatusIs;
using ::testing::HasSubstr;
using ::testing::PrintToString;
using ::testing::SizeIs;

class MarkerTest : public testing::TestWithParam<uint16_t> {};

INSTANTIATE_TEST_SUITE_P(MarkerBasedTests, MarkerTest,
                         ::testing::Range(uint16_t{0xff00}, uint16_t{0xffff}),
                         [](const testing::TestParamInfo<uint16_t>& marker) {
                           return absl::StrFormat("%04x", marker.param);
                         });

TEST_P(MarkerTest, HandlesMarker) {
  std::string data("\0\0\0\4\1\1\2\2\1\1\1\1\1\1\1\1", 16);
  data[0] = static_cast<uint8_t>((GetParam() & 0x0000ff00) >> 8);
  data[1] = static_cast<uint8_t>((GetParam() & 0x000000ff));
  data[6] = 0xff;
  data[7] = 0xff;

  riegeli::StringReader<> input(data);

  JpegSegment segment;
  bool called = false;
  EXPECT_THAT(
      IterateOverJpegSegments(input,
                              [&segment, &called](const JpegSegment& chunk) {
                                segment = std::move(chunk);
                                called = true;
                                return false;  // stop after first
                              }),
      IsOk());
  ASSERT_TRUE(called);

  EXPECT_EQ(segment.marker, GetParam());

  EXPECT_EQ(segment.offset, 0);

  if (GetParam() == 0xff01 || GetParam() == 0xffd8 || GetParam() == 0xffd9) {
    // No payload.
    EXPECT_EQ(segment.length, 2);
    EXPECT_EQ(segment.payload_offset, 2);
    EXPECT_EQ(segment.payload_length, 0);
    return;
  }

  EXPECT_EQ(segment.length, 6);

  if ((GetParam() >= 0xffd0 && GetParam() <= 0xffd7)) {
    // No defined length, so length bytes are part of payload
    EXPECT_EQ(segment.payload_offset, 2);
    EXPECT_EQ(segment.payload_length, 4);
    return;
  }

  EXPECT_EQ(segment.payload_offset, 4);
  EXPECT_EQ(segment.payload_length, 2);
}

TEST(ReaderTest, IterateOverMultipleSegments) {
  // SOI (2 bytes), APP0 (6 bytes with 2 bytes payload), EOI (2 bytes)
  std::string data("\xff\xd8\xff\xe0\x00\x04\x01\x02\xff\xd9", 10);
  riegeli::StringReader<> input(data);

  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;  // continue
                                      }),
              IsOk());

  // SOI is processed, APP0 is processed, EOI is processed (and breaks the loop)
  ASSERT_EQ(segments.size(), 3);
  EXPECT_EQ(segments[0].marker, 0xffd8);
  EXPECT_EQ(segments[0].marker_label, "SOI");
  EXPECT_EQ(segments[1].marker, 0xffe0);
  EXPECT_EQ(segments[1].marker_label, "APP0");
  EXPECT_EQ(segments[2].marker, 0xffd9);
  EXPECT_EQ(segments[2].marker_label, "EOI");
}

TEST(ReaderTest, ParsesValidApp11Segment) {
  std::string data(
      "\xff\xd8"                          // SOI
      "\xff\xeb\x00\x12"                  // APP11, length 18
      "\x4a\x50\x00\x01\x00\x00\x00\x01"  // APP11 header (type 0x4a50, id 1,
                                          // seq 1)
      "\x00\x00\x00\x08\x6a\x75\x6d\x62"  // JUMBF header (lbox 8, tbox 'jumb')
      "\xff\xd9",                         // EOI
      24);
  riegeli::StringReader<> input(data);

  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;
                                      }),
              IsOk());

  ASSERT_EQ(segments.size(), 3);
  EXPECT_EQ(segments[1].marker, 0xffeb);
  EXPECT_EQ(segments[1].marker_label, "APP11");
  ASSERT_TRUE(segments[1].app11_info.has_value());
  EXPECT_EQ(segments[1].app11_info->extension_type, 0x4a50);
  EXPECT_EQ(segments[1].app11_info->segment_id, 1);
  EXPECT_EQ(segments[1].app11_info->sequence_number, 1);
  ASSERT_TRUE(segments[1].app11_info->jumbf_data.has_value());
  EXPECT_EQ(segments[1].app11_info->jumbf_data->lbox, 8);
  EXPECT_EQ(segments[1].app11_info->jumbf_data->tbox, 0x6a756d62);
  EXPECT_EQ(segments[1].app11_info->jumbf_data->xlbox, 0);
}

TEST(ReaderTest, HandlesTruncatedBox) {
  // SOI (2 bytes), APP0 (18 bytes total: marker 2, length 2, payload 14)
  std::string data(
      "\xff\xd8"
      "\xff\xe0\x00\x10\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d"
      "\x0e",
      20);
  riegeli::StringReader<> input(data);

  // We pass end_offset = 15, which truncates the APP0 box.
  EXPECT_THAT(
      IterateOverJpegSegments(
          input, [](const JpegSegment& chunk) { return true; }, 15),
      StatusIs(absl::StatusCode::kInvalidArgument, "truncated JPEG box"));
}

TEST(ReaderTest, ParsesSosWithByteStuffing) {
  std::string data(
      "\xff\xd8"                          // SOI
      "\xff\xda\x00\x06\xaa\xbb\xcc\xdd"  // SOS header (size 8)
      "\x11\x22\xff\x00\x33\x44"  // Scan data (6 bytes, containing stuffed
                                  // \xff\x00)
      "\xff\xd9",                 // EOI
      18);
  riegeli::StringReader<> input(data);

  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;
                                      }),
              IsOk());

  ASSERT_EQ(segments.size(), 3);
  EXPECT_EQ(segments[1].marker, 0xffda);
  EXPECT_EQ(segments[1].marker_label, "SOS");
  EXPECT_EQ(segments[1].offset, 2);
  EXPECT_EQ(segments[1].length, 14);  // 8 (header) + 6 (scan data)
  EXPECT_EQ(segments[1].payload_offset, 6);
  EXPECT_EQ(segments[1].payload_length, 10);

  EXPECT_EQ(segments[2].marker, 0xffd9);
  EXPECT_EQ(segments[2].marker_label, "EOI");
  EXPECT_EQ(segments[2].offset, 16);
}

TEST(ReaderTest, ParsesC2paAfterSegment) {
  std::string data(
      "\xff\xd8"           // SOI
      "\xff\xd9"           // EOI
      "\x01\x02\x03\x04",  // Extra data
      8);
  riegeli::StringReader<> input(data);

  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;
                                      }),
              IsOk());

  ASSERT_EQ(segments.size(), 3);
  EXPECT_EQ(segments[0].marker, 0xffd8);
  EXPECT_EQ(segments[1].marker, 0xffd9);

  EXPECT_EQ(segments[2].marker, 0x0000);
  EXPECT_EQ(segments[2].marker_label, "c2pa.after");
  EXPECT_EQ(segments[2].offset, 4);
  EXPECT_EQ(segments[2].length, 4);
  EXPECT_EQ(segments[2].payload_offset, 0);
  EXPECT_EQ(segments[2].payload_length, 0);
}

TEST(ReaderTest, PropagatesProcessorError) {
  std::string data("\xff\xd8\xff\xe0\x00\x04\x01\x02\xff\xd9", 10);
  riegeli::StringReader<> input(data);

  EXPECT_THAT(
      IterateOverJpegSegments(input,
                              [](const JpegSegment& chunk) {
                                return absl::InternalError("test error");
                              }),
      StatusIs(absl::StatusCode::kInternal, "test error"));
}

TEST(ReaderTest, PrintToJpegJumbfData) {
  JpegJumbfData data = {.lbox = 1,
                        .tbox = 2,
                        .xlbox = 3,
                        .payload_offset = 4,
                        .payload_length = 5};
  EXPECT_EQ(
      PrintToString(data),
      "{lbox: 1, tbox: 2, xlbox: 3, payload_offset: 4, payload_length: 5}");
}

TEST(ReaderTest, PrintToJpegApp11Segment) {
  JpegApp11Segment segment = {
      .extension_type = 1,
      .segment_id = 2,
      .sequence_number = 3,
      .payload_offset = 4,
      .payload_length = 5,
      .jumbf_data = JpegJumbfData{.lbox = 10,
                                  .tbox = 20,
                                  .xlbox = 30,
                                  .payload_offset = 40,
                                  .payload_length = 50}};

  EXPECT_EQ(PrintToString(segment),
            "{extension_type: 1, segment_id: 2, sequence_number: 3, "
            "payload_offset: 4, payload_length: 5, jumbf_data: {lbox: 10, "
            "tbox: 20, xlbox: 30, payload_offset: 40, payload_length: 50}}");

  JpegApp11Segment segment_no_jumbf = {.extension_type = 1,
                                       .segment_id = 2,
                                       .sequence_number = 3,
                                       .payload_offset = 4,
                                       .payload_length = 5,
                                       .jumbf_data = std::nullopt};
  EXPECT_EQ(PrintToString(segment_no_jumbf),
            "{extension_type: 1, segment_id: 2, sequence_number: 3, "
            "payload_offset: 4, payload_length: 5, jumbf_data: nullopt}");
}

TEST(ReaderTest, PrintToJpegSegment) {
  JpegSegment segment = {
      .offset = 1,
      .marker = 2,
      .marker_label = "LABEL",
      .length = 3,
      .payload_offset = 4,
      .payload_length = 5,
      .app11_info = JpegApp11Segment{.extension_type = 10,
                                     .segment_id = 20,
                                     .sequence_number = 30,
                                     .payload_offset = 40,
                                     .payload_length = 50,
                                     .jumbf_data = std::nullopt}};

  EXPECT_EQ(PrintToString(segment),
            "{offset: 1, marker: 2, marker_label: \"LABEL\", length: 3, "
            "payload_offset: 4, payload_length: 5, app11_info: "
            "{extension_type: 10, segment_id: 20, sequence_number: 30, "
            "payload_offset: 40, payload_length: 50, jumbf_data: nullopt}}");

  JpegSegment segment_no_app11 = {.offset = 1,
                                  .marker = 2,
                                  .marker_label = "LABEL",
                                  .length = 3,
                                  .payload_offset = 4,
                                  .payload_length = 5,
                                  .app11_info = std::nullopt};
  EXPECT_EQ(PrintToString(segment_no_app11),
            "{offset: 1, marker: 2, marker_label: \"LABEL\", length: 3, "
            "payload_offset: 4, payload_length: 5, app11_info: nullopt}");
}

TEST(ReaderTest, MalformedSegmentSize) {
  std::string data("\xff\xd8\xff\xe0\x00\x00\xff\xd9", 8);
  riegeli::StringReader<> input(data);
  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;
                                      }),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("Invalid JPEG segment size: 0")));
}

TEST(ReaderTest, App11SegmentShortLengthDesynchronization) {
  std::string data(
      "\xff\xd8"
      "\xff\xeb\x00\x04\x00\x00"
      "\xff\xeb\x00\x12\x4a\x50\x00\x01\x00\x00\x00\x01\x00\x00\x00\x08\x6a\x75"
      "\x6d\x62"
      "\xff\xd9",
      30);
  riegeli::StringReader<> input(data);
  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;
                                      }),
              IsOk());
  ASSERT_EQ(segments.size(), 4);
  EXPECT_EQ(segments[0].marker_label, "SOI");
  EXPECT_EQ(segments[1].marker, 0xffeb);
  EXPECT_EQ(segments[1].length, 6);
  EXPECT_FALSE(segments[1].app11_info.has_value());
  EXPECT_EQ(segments[2].marker, 0xffeb);
  EXPECT_EQ(segments[2].offset, 8);
  EXPECT_EQ(segments[2].length, 20);
  EXPECT_TRUE(segments[2].app11_info.has_value());
  EXPECT_EQ(segments[3].marker_label, "EOI");
}

TEST(ReaderTest, App11JumbfBoxXlboxTooShort) {
  std::string data(
      "\xff\xd8"
      "\xff\xeb\x00\x1c\x4a\x50\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x6a\x75"
      "\x6d\x62\x00\x00\x00\x00\x00\x00\x00\x0a\x00\x00"
      "\xff\xd9",
      34);
  riegeli::StringReader<> input(data);
  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;
                                      }),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("JUMBf box payload too short")));
}

TEST(ReaderTest, App11JumbfBoxLboxReserved) {
  std::string data(
      "\xff\xd8"
      "\xff\xeb\x00\x12\x4a\x50\x00\x01\x00\x00\x00\x01\x00\x00\x00\x05\x6a\x75"
      "\x6d\x62"
      "\xff\xd9",
      24);
  riegeli::StringReader<> input(data);
  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;
                                      }),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("JUMBf box payload too short")));
}

TEST(ReaderTest, App11JumbfBoxPayloadLengthTooShort) {
  // SOI (2 bytes)
  // APP11 (16 bytes)
  //   - header (4 bytes)
  //   - extension info (8 bytes)
  //   - JUMBF header (4 bytes) * Invalid: JUMBF needs a minimum of 8 bytes
  // EOI (2 bytes)
  std::string data(
      "\xff\xd8"
      "\xff\xeb\x00\x0e\x4a\x50\x00\x01\x00\x00\x00\x01\x00\x00\x00\x08"
      "\xff\xd9",
      20);
  riegeli::StringReader<> input(data);
  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;
                                      }),
              IsOk());
  EXPECT_THAT(segments, SizeIs(3));
  EXPECT_EQ(segments[0].marker_label, "SOI");
  EXPECT_EQ(segments[1].marker_label, "APP11");
  EXPECT_TRUE(segments[1].app11_info.has_value());
  EXPECT_FALSE(segments[1].app11_info->jumbf_data.has_value());
  EXPECT_EQ(segments[2].marker_label, "EOI");
}

TEST(ReaderTest, PayloadLengthDoesNotUnderflowOnShortSegment) {
  // SOI (2 bytes: \xff\xd8) followed by EOI (\xff\xd9)
  std::string data("\xff\xd8\xff\xd9", 4);
  riegeli::StringReader<> input(data);
  std::vector<JpegSegment> segments;
  EXPECT_THAT(IterateOverJpegSegments(input,
                                      [&segments](const JpegSegment& chunk) {
                                        segments.push_back(chunk);
                                        return true;
                                      }),
              IsOk());
  ASSERT_GE(segments.size(), 1);
  EXPECT_EQ(segments[0].marker, 0xffd8);
  EXPECT_EQ(segments[0].payload_length, 0);
}

}  // namespace
}  // namespace credentio