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

#include <cstdint>
#include <optional>
#include <string>
#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_byte_info.h"
#include "formats/byte_range.h"
#include "formats/extractor_result.h"
#include "formats/png/extractor.h"
#include "formats/png/hasher.h"
#include "formats/png/test_utils.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"

namespace credentio {
namespace {

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

constexpr int64_t kChunkHeaderSize = 8;

absl::StatusOr<ExtractorResult> Extract(absl::string_view png) {
  riegeli::StringReader input(png);
  PngExtractor extractor;
  ExtractorResult result = {.manifest_store = "", .asset_byte_info = {}};

  ABSL_ASSIGN_OR_RETURN(result.manifest_store,
                        extractor.ExtractManifestStore(input));

  if (!input.Seek(0)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failure to seek"));
  }
  ABSL_ASSIGN_OR_RETURN(std::optional<ByteRange> manifest_store_location,
                        extractor.ExtractManifestStoreLocation(input, {}));
  result.asset_byte_info.manifest_store_location =
      manifest_store_location.value_or(ByteRange{.offset = 0, .length = 0});

  if (!input.Seek(0)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failure to seek"));
  }
  ABSL_ASSIGN_OR_RETURN(result.asset_byte_info.boxes,
                        extractor.ExtractBoxes(input, {}));
  return result;
}

std::vector<std::string> GetBoxNames(const ExtractorResult& result) {
  std::vector<std::string> box_names;
  for (const auto& box : result.asset_byte_info.boxes) {
    box_names.push_back(box.identifier);
  }
  return box_names;
}

TEST(PngEmbedderTest, CopyAndEmbedBeforeIdat) {
  absl::StatusOr<std::string> input_s =
      CreatePng({{.type = 'ABCD', .payload = "first chunk"},
                 {.type = 'EFGH', .payload = "second chunk"},
                 {.type = 'IDAT', .payload = "third chunk"},
                 {.type = 'XYZW', .payload = "fourth chunk"}});
  ABSL_ASSERT_OK(input_s);
  constexpr int64_t kInsertionOffset = 55;
  constexpr absl::string_view kManifestStore = "C2PA manifest store bytes";
  riegeli::StringReader input(*input_s);
  std::string output_s;
  riegeli::StringWriter output(&output_s);

  PngEmbedder embedder;
  ABSL_ASSERT_OK(embedder.CopyAndEmbed(
      input, kManifestStore, DigestPngResult{.insertion_offset = 55}, output));
  ASSERT_TRUE(output.Close());

  EXPECT_THAT(output_s.substr(kInsertionOffset + kChunkHeaderSize,
                              kManifestStore.size()),
              Eq(kManifestStore));
  EXPECT_THAT(output_s, SizeIs(input_s->size() + 12 + kManifestStore.size()));
  absl::StatusOr<ExtractorResult> extracted = Extract(output_s);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(extracted->manifest_store, Eq(kManifestStore));
  EXPECT_THAT(
      GetBoxNames(*extracted),
      ElementsAre("PNGh", "ABCD", "EFGH", "C2PA", "IDAT", "XYZW", "IEND"));
}

TEST(PngEmbedderTest, SeeksToBeginningBeforeCopying) {
  absl::StatusOr<std::string> input_s =
      CreatePng({{.type = 'ABCD', .payload = "first chunk"},
                 {.type = 'EFGH', .payload = "second chunk"},
                 {.type = 'IDAT', .payload = "third chunk"},
                 {.type = 'XYZW', .payload = "fourth chunk"}});
  ABSL_ASSERT_OK(input_s);
  constexpr int64_t kInsertionOffset = 55;
  constexpr absl::string_view kManifestStore = "C2PA manifest store bytes";
  riegeli::StringReader input(*input_s);
  std::string output_s;
  riegeli::StringWriter output(&output_s);

  // Seek to a non-zero offset.
  ASSERT_TRUE(input.Seek(10));

  PngEmbedder embedder;
  ABSL_ASSERT_OK(embedder.CopyAndEmbed(
      input, kManifestStore, DigestPngResult{.insertion_offset = 55}, output));
  ASSERT_TRUE(output.Close());

  EXPECT_THAT(output_s.substr(kInsertionOffset + kChunkHeaderSize,
                              kManifestStore.size()),
              Eq(kManifestStore));
  EXPECT_THAT(output_s, SizeIs(input_s->size() + 12 + kManifestStore.size()));
  absl::StatusOr<ExtractorResult> extracted = Extract(output_s);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(extracted->manifest_store, Eq(kManifestStore));
  EXPECT_THAT(
      GetBoxNames(*extracted),
      ElementsAre("PNGh", "ABCD", "EFGH", "C2PA", "IDAT", "XYZW", "IEND"));
}

TEST(PngEmbedderTest, CopyAndEmbedBeforeIend) {
  absl::StatusOr<std::string> input_s =
      CreatePng({{.type = 'ABCD', .payload = "first chunk"},
                 {.type = 'EFGH', .payload = "second chunk"},
                 {.type = 'XYZW', .payload = "third chunk"}});
  ABSL_ASSERT_OK(input_s);
  constexpr int64_t kInsertionOffset = 78;
  constexpr absl::string_view kManifestStore = "C2PA manifest store bytes";
  riegeli::StringReader input(*input_s);
  std::string output_s;
  riegeli::StringWriter output(&output_s);

  PngEmbedder embedder;
  ABSL_ASSERT_OK(embedder.CopyAndEmbed(
      input, kManifestStore,
      DigestPngResult{.insertion_offset = kInsertionOffset}, output));
  ASSERT_TRUE(output.Close());

  EXPECT_THAT(output_s.substr(kInsertionOffset + kChunkHeaderSize,
                              kManifestStore.size()),
              Eq(kManifestStore));
  EXPECT_THAT(output_s, SizeIs(input_s->size() + 12 + kManifestStore.size()));
  absl::StatusOr<ExtractorResult> extracted = Extract(output_s);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(extracted->manifest_store, Eq(kManifestStore));
  EXPECT_THAT(GetBoxNames(*extracted),
              ElementsAre("PNGh", "ABCD", "EFGH", "XYZW", "C2PA", "IEND"));
}

TEST(PngEmbedderTest, CopyAndEmbedC2paReplaceBeforeIdat) {
  absl::StatusOr<std::string> input_s =
      CreatePng({{.type = 'ABCD', .payload = "first chunk"},
                 {.type = 'caBX', .payload = "C2PA chunk"},
                 {.type = 'IDAT', .payload = "some data"}});
  ABSL_ASSERT_OK(input_s);
  constexpr int64_t kInsertionOffset = 31;
  constexpr int64_t kExistingC2paLength = 22;
  constexpr absl::string_view kManifestStore = "C2PA manifest store bytes";
  riegeli::StringReader input(*input_s);
  std::string output_s;
  riegeli::StringWriter output(&output_s);
  PngEmbedder embedder;
  ABSL_ASSERT_OK(embedder.CopyAndEmbed(
      input, kManifestStore,
      DigestPngResult{.insertion_offset = kInsertionOffset,
                      .existing_c2pa_length = kExistingC2paLength},
      output));
  ASSERT_TRUE(output.Close());

  EXPECT_THAT(output_s.substr(kInsertionOffset + kChunkHeaderSize,
                              kManifestStore.size()),
              Eq(kManifestStore));
  absl::StatusOr<ExtractorResult> extracted = Extract(output_s);
  ABSL_ASSERT_OK(extracted);
  EXPECT_THAT(extracted->manifest_store, Eq(kManifestStore));
  EXPECT_THAT(GetBoxNames(*extracted),
              ElementsAre("PNGh", "ABCD", "C2PA", "IDAT", "IEND"));
}

TEST(PngEmbedderTest, CopyAndEmbedC2paAlreadyExistsAfterIdat_Replace) {
  // Note that the digest operation would fail because there is a C2PA chunk
  // after the IDAT so this scenario would never actually occur.
  absl::StatusOr<std::string> input_s = CreatePng({
      {.type = 'ABCD', .payload = "first chunk"},
      {.type = 'IDAT', .payload = "some data"},
      {.type = 'caBX', .payload = "C2PA chunk"},
  });
  ABSL_ASSERT_OK(input_s);
  constexpr int64_t kInsertionOffset = 31;
  constexpr absl::string_view kManifestStore = "C2PA manifest store bytes";
  riegeli::StringReader input(*input_s);
  std::string output_s;
  riegeli::StringWriter output(&output_s);
  PngEmbedder embedder;
  ABSL_ASSERT_OK(embedder.CopyAndEmbed(
      input, kManifestStore,
      DigestPngResult{.insertion_offset = kInsertionOffset}, output));
  ASSERT_TRUE(output.Close());

  EXPECT_THAT(output_s.substr(kInsertionOffset + kChunkHeaderSize,
                              kManifestStore.size()),
              Eq(kManifestStore));
  EXPECT_THAT(Extract(output_s),
              StatusIs(absl::StatusCode::kNotFound,
                       HasSubstr("Multiple manifest stores found")));
}

}  // namespace
}  // namespace credentio