// 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/embedder.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 "absl/types/span.h"
#include "constants/labels.h"
#include "formats/jpeg/extractor.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/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"
#include "riegeli/endian/endian_writing.h"

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::testing::AllOf;
using ::testing::HasSubstr;

absl::StatusOr<std::string> CreateFakeManifestStore(size_t content_size = 100) {
  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());
}

struct JpegSegment {
  uint16_t marker;
  std::string payload;
};

absl::StatusOr<std::string> CreateJpeg(absl::Span<const JpegSegment> segments) {
  std::string buffer;
  riegeli::StringWriter<> output(&buffer);
  for (auto segment : segments) {
    if (!riegeli::WriteBigEndian<uint16_t>(segment.marker, output)) {
      return output.status();
    }
    if (!segment.payload.empty()) {
      uint16_t seg_size = sizeof(uint16_t) + segment.payload.size();
      if (!riegeli::WriteBigEndian<uint16_t>(seg_size, output)) {
        return output.status();
      }
      if (!output.Write(segment.payload)) {
        return output.status();
      }
    }
  }
  output.Close();
  return buffer;
}

TEST(JpegEmbedderTest, CopyAndEmbedAfterSOI) {
  auto input_s_or = CreateJpeg({{.marker = /*SOI*/ 0xffd8, .payload = ""}});
  ABSL_ASSERT_OK(input_s_or);
  const auto& input_s = *input_s_or;

  riegeli::StringReader<absl::string_view> input(input_s);
  std::string output_s;
  riegeli::StringWriter output(&output_s);
  auto manifest_store_or = CreateFakeManifestStore();
  ABSL_ASSERT_OK(manifest_store_or);
  const auto& manifest_store = *manifest_store_or;

  JpegEmbedder embedder(/*segment_id=*/30);
  ABSL_ASSERT_OK(embedder.CopyAndEmbed(input, manifest_store,
                                       /*insert_at_offset=*/2, output));
  output.Close();

  riegeli::StringReader<absl::string_view> output_reader(output_s);
  EXPECT_THAT(JpegExtractor().ExtractManifestStore(output_reader),
              IsOkAndHolds(manifest_store));
}

TEST(JpegEmbedderTest, CopyAndEmbedBeforeSOS) {
  auto input_s_or = CreateJpeg({{.marker = /*SOI*/ 0xffd8, .payload = ""},
                                {.marker = /*SOS*/ 0xffda, .payload = "abcd"},
                                {.marker = /*EOI*/ 0xffd9, .payload = ""}});
  ABSL_ASSERT_OK(input_s_or);
  const auto& input_s = *input_s_or;

  riegeli::StringReader input(input_s);
  std::string output_s;
  riegeli::StringWriter output(&output_s);
  auto manifest_store_or = CreateFakeManifestStore();
  ABSL_ASSERT_OK(manifest_store_or);
  const auto& manifest_store = *manifest_store_or;

  JpegEmbedder embedder(/*segment_id=*/30);
  ABSL_ASSERT_OK(embedder.CopyAndEmbed(input, manifest_store,
                                       /*insert_at_offset=*/2, output));
  output.Close();

  riegeli::StringReader output_reader(output_s);
  EXPECT_THAT(JpegExtractor().ExtractManifestStore(output_reader),
              IsOkAndHolds(manifest_store));
}

TEST(JpegEmbedderTest, CopyAndEmbedInAppSegments) {
  auto input_s_or =
      CreateJpeg({{.marker = /*SOI*/ 0xffd8, .payload = ""},
                  {.marker = /*APP0*/ 0xFFE0, .payload = "abcd"},
                  {.marker = /*APP13*/ 0xFFED, .payload = "efgh"},
                  {.marker = /*APP0*/ 0xFFE0, .payload = "ijkl"}});
  ABSL_ASSERT_OK(input_s_or);
  const auto& input_s = *input_s_or;

  riegeli::StringReader input(input_s);
  std::string output_s;
  riegeli::StringWriter output(&output_s);
  auto manifest_store_or = CreateFakeManifestStore();
  ABSL_ASSERT_OK(manifest_store_or);
  const auto& manifest_store = *manifest_store_or;

  JpegEmbedder embedder(/*segment_id=*/30);
  ABSL_ASSERT_OK(embedder.CopyAndEmbed(input, manifest_store,
                                       /*insert_at_offset=*/10, output));
  output.Close();
  EXPECT_THAT(output_s,
              AllOf(HasSubstr("abcd"), HasSubstr("efgh"), HasSubstr("ijkl")));

  riegeli::StringReader output_reader(output_s);
  EXPECT_THAT(JpegExtractor().ExtractManifestStore(output_reader),
              IsOkAndHolds(manifest_store));
}

}  // namespace
}  // namespace credentio