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

#include <cstddef>
#include <cstdint>
#include <string>
#include <utility>

#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/cord.h"
#include "absl/strings/string_view.h"
#include "constants/labels.h"
#include "formats/asset_box.h"
#include "formats/byte_range.h"
#include "formats/jpeg/extractor.h"
#include "formats/jpeg/testing/jumbf_creator.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "jumbf/box_builder.h"
#include "jumbf/constants.h"
#include "jumbf/encode_params.h"
#include "riegeli/base/object.h"
#include "riegeli/bytes/limiting_writer.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::testing::ElementsAre;
using ::testing::HasSubstr;
using ::testing::Optional;

// Creates a jumbf SuperBox that looks superficially like a Manifest Store.
absl::StatusOr<std::string> CreateFakeManifestStore(size_t content_size) {
  jumbf::SuperBoxBuilder manifest_store_builder(
      kManifestStoreUuid,
      jumbf::DescriptionBoxOptions{.label = absl::Cord(kManifestStoreLabel)});
  ABSL_RETURN_IF_ERROR(manifest_store_builder.AddContent(
      jumbf::kBinaryDataBoxType, absl::Cord(std::string(content_size, 'x'))));
  ABSL_ASSIGN_OR_RETURN(jumbf::BuiltSuperBox manifest_store_box,
                        std::move(manifest_store_builder).Finalize());
  return std::string(std::move(manifest_store_box).AsCord());
}

TEST(App11IOTest, WriteSmallPayload) {
  auto manifest_store_or = CreateFakeManifestStore(/*content_size=*/100);
  ABSL_ASSERT_OK(manifest_store_or);
  const auto& manifest_store = *manifest_store_or;

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(
      WriteJpegApp11Payload(/*segment_id=*/30, manifest_store, output));
  ASSERT_TRUE(output.Close());

  riegeli::StringReader<absl::string_view> input(output_s);
  JpegExtractor extractor;

  EXPECT_THAT(extractor.ExtractManifestStore(input),
              IsOkAndHolds(manifest_store));

  ASSERT_TRUE(input.Seek(0));

  EXPECT_THAT(extractor.ExtractManifestStoreLocation(input, {}),
              IsOkAndHolds(Optional(ByteRange{.offset = 0, .length = 158})));

  ASSERT_TRUE(input.Seek(0));

  EXPECT_THAT(
      extractor.ExtractBoxes(input, {}),
      IsOkAndHolds(ElementsAre(AssetBox{
          .identifier = "C2PA", .byte_range = {.offset = 0, .length = 158}})));
}

TEST(App11IOTest, WriteLargePayload) {
  auto manifest_store_or = CreateFakeManifestStore(/*content_size=*/250000);
  ABSL_ASSERT_OK(manifest_store_or);
  const auto& manifest_store = *manifest_store_or;

  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_ASSERT_OK(
      WriteJpegApp11Payload(/*segment_id=*/30, manifest_store, output));
  ASSERT_TRUE(output.Close());

  riegeli::StringReader<absl::string_view> input(output_s);

  JpegExtractor extractor;
  EXPECT_THAT(extractor.ExtractManifestStore(input),
              IsOkAndHolds(manifest_store));

  ASSERT_TRUE(input.Seek(0));

  EXPECT_THAT(extractor.ExtractManifestStoreLocation(input, {}),
              IsOkAndHolds(Optional(ByteRange{.offset = 0, .length = 250118})));

  ASSERT_TRUE(input.Seek(0));

  EXPECT_THAT(extractor.ExtractBoxes(input, {}),
              IsOkAndHolds(ElementsAre(
                  AssetBox{.identifier = "C2PA",
                           .byte_range = {.offset = 0, .length = 250118}})));
}

TEST(App11IOTest, WriteInvalidPayload) {
  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  EXPECT_THAT(
      WriteJpegApp11Payload(/*segment_id=*/30, std::string(6, '\0'), output),
      StatusIs(absl::StatusCode::kDataLoss, HasSubstr("Failed to read tbox")));
}

TEST(App11IOTest, WriteInvalidJumbfPayload) {
  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  EXPECT_THAT(
      WriteJpegApp11Payload(/*segment_id=*/30, std::string(60, '\0'), output),
      StatusIs(absl::StatusCode::kInvalidArgument,
               HasSubstr("invalid jumbf box")));
}

TEST(App11IOTest, WritePayloadWithXlbox) {
  std::string output_s;
  riegeli::StringWriter<> output(&output_s);
  ABSL_EXPECT_OK(WriteJpegApp11Payload(
      /*segment_id=*/30, CreateJumbf({.force_xlbox = true}), output));
  ASSERT_TRUE(output.Close());
}

TEST(App11IOTest, WritePayloadFails) {
  auto manifest_store_or = CreateFakeManifestStore(/*content_size=*/100);
  ABSL_ASSERT_OK(manifest_store_or);
  const auto& manifest_store = *manifest_store_or;

  std::string output_s;
  riegeli::StringWriter<> string_writer(&output_s);
  // Limit the writer to 1 byte, forcing a failure.
  riegeli::LimitingWriter<> limited_writer(
      &string_writer, riegeli::LimitingWriterBase::Options().set_max_length(1));
  EXPECT_THAT(
      WriteJpegApp11Payload(/*segment_id=*/30, manifest_store, limited_writer),
      StatusIs(absl::StatusCode::kResourceExhausted));
  EXPECT_FALSE(limited_writer.Close());
}

TEST(App11IOTest, WritePayloadFailsImmediately) {
  auto manifest_store_or = CreateFakeManifestStore(/*content_size=*/100);
  ABSL_ASSERT_OK(manifest_store_or);
  const auto& manifest_store = *manifest_store_or;

  riegeli::StringWriter<> failing_writer(riegeli::kClosed);
  failing_writer.Fail(absl::InternalError("Test failure"));
  EXPECT_THAT(
      WriteJpegApp11Payload(/*segment_id=*/30, manifest_store, failing_writer),
      StatusIs(absl::StatusCode::kInternal, HasSubstr("Test failure")));
  EXPECT_FALSE(failing_writer.Close());
}

TEST(App11IOTest, CalculateSizeOneSegment) {
  EXPECT_EQ(CalculateTotalJpegApp11SegmentsSize(1000), 1012);
}

TEST(App11IOTest, CalculateSizeExactSingleSegmentSize) {
  EXPECT_EQ(CalculateTotalJpegApp11SegmentsSize(65008), 65020);
}

TEST(App11IOTest, CalculateSizeMultipleSegments) {
  EXPECT_EQ(CalculateTotalJpegApp11SegmentsSize(65200), 65232);
}

TEST(App11IOTest, CalculateSizeOf64BitPayload) {
  EXPECT_EQ(CalculateTotalJpegApp11SegmentsSize(
                static_cast<uint64_t>(UINT32_MAX) + 1000),
            4296818435);
}

}  // namespace
}  // namespace credentio