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

#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#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/byte_range.h"
#include "formats/zip/constants.h"
#include "formats/zip/extractor.h"
#include "formats/zip/reader.h"
#include "formats/zip/zlib_util.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"
#include "testing/test_file_utils.h"

namespace credentio {
namespace {

using ::absl_testing::IsOk;
using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::testing::HasSubstr;

absl::StatusOr<std::vector<ZipReader::FileEntry>> GetFileEntries(
    riegeli::Reader& input) {
  ABSL_ASSIGN_OR_RETURN(auto reader, ZipReader::Create(&input));
  std::vector<ZipReader::FileEntry> file_entries;
  while (reader->HasNext()) {
    ABSL_ASSIGN_OR_RETURN(ZipReader::FileEntry entry, reader->Next());
    file_entries.push_back(entry);
  }
  return file_entries;
}

absl::StatusOr<std::string> GetContentTypesStream(riegeli::Reader& input) {
  ABSL_ASSIGN_OR_RETURN(auto reader, ZipReader::Create(&input));
  while (reader->HasNext()) {
    ABSL_ASSIGN_OR_RETURN(ZipReader::FileEntry entry, reader->Next());
    if (entry.file_name == kOoxmlContentTypesFileName) {
      if (!input.Seek(entry.file_range.offset)) {
        return input.StatusOrAnnotate(absl::DataLossError("failed to seek"));
      }
      std::string content;
      if (!input.Read(entry.file_range.length, content)) {
        return input.StatusOrAnnotate(absl::DataLossError("failed to read"));
      }
      ABSL_ASSIGN_OR_RETURN(std::string uncompressed, Uncompress(content));
      return uncompressed;
    }
  }
  return absl::NotFoundError("Content types stream not found.");
}

class ZipEmbedderTest : public testing::Test {
 protected:
  void InitWithZipFile(absl::string_view file_name,
                       int64_t manifest_store_size) {
    auto source = credentio_testing::GetFileReader(file_name);
    ABSL_ASSERT_OK(source);
    source_ = *std::move(source);
    auto file_entries = GetFileEntries(*source_);
    ABSL_ASSERT_OK(file_entries);
    file_entries_ = *std::move(file_entries);

    destination_ = std::make_unique<riegeli::StringWriter<>>(&output_s_);
    manifest_store_size_ = manifest_store_size;
  }

  std::string zip_file_contents_;
  std::unique_ptr<riegeli::Reader> source_;
  std::vector<ZipReader::FileEntry> file_entries_;
  std::string output_s_;
  std::unique_ptr<riegeli::StringWriter<>> destination_;
  int64_t manifest_store_size_;
};

TEST_F(ZipEmbedderTest, CopyAndAllocateFailsIfManifestStoreAlreadyExists) {
  InitWithZipFile(
      "credentio/formats/zip/testing/"
      "asset_fake_c2pa.zip",
      /*manifest_store_size=*/100);

  EXPECT_THAT(
      ZipEmbedder().CopyAndAllocate(source_.get(), file_entries_,
                                    manifest_store_size_, destination_.get()),
      StatusIs(absl::StatusCode::kAlreadyExists,
               HasSubstr("Manifest store already exists.")));
}

TEST_F(ZipEmbedderTest, CopyAndAllocateFailsIfNoContentTypesStream) {
  InitWithZipFile("credentio/formats/zip/testing/asset_no_c2pa.zip",
                  /*manifest_store_size=*/100);

  EXPECT_THAT(
      ZipEmbedder().CopyAndAllocate(source_.get(), file_entries_,
                                    manifest_store_size_, destination_.get()),
      StatusIs(absl::StatusCode::kNotFound,
               HasSubstr("Content types stream not found.")));
}

TEST_F(ZipEmbedderTest, CopyAndAllocate) {
  InitWithZipFile("credentio/formats/zip/testing/no_c2pa_032_asset.pptx",
                  /*manifest_store_size=*/100);

  auto copy_and_allocate_result = ZipEmbedder().CopyAndAllocate(
      source_.get(), file_entries_, manifest_store_size_, destination_.get());
  ABSL_ASSERT_OK(copy_and_allocate_result);
  ASSERT_TRUE(destination_->Close());

  riegeli::StringReader<absl::string_view> output_reader(output_s_);

  EXPECT_THAT(ZipExtractor().ExtractManifestStore(output_reader),
              IsOkAndHolds(std::string(100, '\0')));
  ASSERT_TRUE(output_reader.Seek(0));
  EXPECT_THAT(ZipExtractor().ExtractManifestStoreLocation(output_reader, {}),
              IsOkAndHolds(ByteRange{
                  .offset = copy_and_allocate_result->manifest_store_offset,
                  .length = 100,
              }));
  ASSERT_TRUE(output_reader.Seek(0));
  auto content_types_stream = GetContentTypesStream(output_reader);
  ABSL_ASSERT_OK(content_types_stream);
  EXPECT_THAT(
      *content_types_stream,
      HasSubstr(
          R"(<Override ContentType="application/c2pa" PartName="/META-INF/content_credential.c2pa"/>)"));
}

TEST_F(ZipEmbedderTest, Embed) {
  std::string manifest_store = "c2pa manifest store";
  InitWithZipFile("credentio/formats/zip/testing/no_c2pa_032_asset.pptx",
                  manifest_store.size());

  auto copy_and_allocate_result = ZipEmbedder().CopyAndAllocate(
      source_.get(), file_entries_, manifest_store_size_, destination_.get());
  ABSL_ASSERT_OK(copy_and_allocate_result);

  EXPECT_THAT(ZipEmbedder().Embed(manifest_store, *copy_and_allocate_result,
                                  destination_.get()),
              IsOk());
  ASSERT_TRUE(destination_->Close());

  riegeli::StringReader output_reader(&output_s_);

  EXPECT_THAT(ZipExtractor().ExtractManifestStore(output_reader),
              IsOkAndHolds(manifest_store));
}

TEST_F(ZipEmbedderTest, CopyAndAllocateFailsWithDuplicateEntries) {
  InitWithZipFile("credentio/formats/zip/testing/no_c2pa_032_asset.pptx", 100);
  ASSERT_GE(file_entries_.size(), 1);

  std::vector<ZipReader::FileEntry> duplicate_entries = file_entries_;
  duplicate_entries.push_back(file_entries_[0]);

  EXPECT_THAT(
      ZipEmbedder().CopyAndAllocate(source_.get(), duplicate_entries,
                                    manifest_store_size_, destination_.get()),
      StatusIs(absl::StatusCode::kInvalidArgument,
               HasSubstr("Overlapping or duplicate ZIP entries detected")));
}

TEST_F(ZipEmbedderTest, CopyAndAllocatePreservesExtraFieldAndCommentData) {
  auto write_le16 = [](std::string& s, uint16_t val) {
    s.push_back(static_cast<char>(val & 0xff));
    s.push_back(static_cast<char>((val >> 8) & 0xff));
  };
  auto write_le32 = [](std::string& s, uint32_t val) {
    s.push_back(static_cast<char>(val & 0xff));
    s.push_back(static_cast<char>((val >> 8) & 0xff));
    s.push_back(static_cast<char>((val >> 16) & 0xff));
    s.push_back(static_cast<char>((val >> 24) & 0xff));
  };

  std::string zip;
  std::string file1_name = "[Content_Types].xml";
  std::string file1_data = "<Types></Types>";
  std::string file1_extra = "EX1";
  std::string file1_cmt = "CM1";

  std::string file2_name = "doc.xml";
  std::string file2_data = "<Document/>";
  std::string file2_extra = "EXTRA2";
  std::string file2_cmt = "CMT2";

  // LFH 1
  uint32_t lfh1_offset = zip.size();
  zip.append("PK\x03\x04", 4);
  write_le16(zip, 10);
  write_le16(zip, 0x0008);  // Has data descriptor
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, 0);
  write_le16(zip, file1_name.size());
  write_le16(zip, file1_extra.size());
  zip.append(file1_name);
  zip.append(file1_extra);
  zip.append(file1_data);
  // Data descriptor with signature
  zip.append("PK\x07\x08", 4);
  write_le32(zip, 0x12345678);
  write_le32(zip, file1_data.size());
  write_le32(zip, file1_data.size());

  // LFH 2
  uint32_t lfh2_offset = zip.size();
  zip.append("PK\x03\x04", 4);
  write_le16(zip, 10);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0x87654321);
  write_le32(zip, file2_data.size());
  write_le32(zip, file2_data.size());
  write_le16(zip, file2_name.size());
  write_le16(zip, file2_extra.size());
  zip.append(file2_name);
  zip.append(file2_extra);
  zip.append(file2_data);

  // CD Start
  uint32_t cd_offset = zip.size();
  // CD 1
  zip.append("PK\x01\x02", 4);
  write_le16(zip, 20);
  write_le16(zip, 10);
  write_le16(zip, 0x0008);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0x12345678);
  write_le32(zip, file1_data.size());
  write_le32(zip, file1_data.size());
  write_le16(zip, file1_name.size());
  write_le16(zip, file1_extra.size());
  write_le16(zip, file1_cmt.size());
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, lfh1_offset);
  zip.append(file1_name);
  zip.append(file1_extra);
  zip.append(file1_cmt);

  // CD 2
  zip.append("PK\x01\x02", 4);
  write_le16(zip, 20);
  write_le16(zip, 10);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0x87654321);
  write_le32(zip, file2_data.size());
  write_le32(zip, file2_data.size());
  write_le16(zip, file2_name.size());
  write_le16(zip, file2_extra.size());
  write_le16(zip, file2_cmt.size());
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, lfh2_offset);
  zip.append(file2_name);
  zip.append(file2_extra);
  zip.append(file2_cmt);

  uint32_t cd_size = zip.size() - cd_offset;
  // EOCD
  zip.append("PK\x05\x06", 4);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 2);
  write_le16(zip, 2);
  write_le32(zip, cd_size);
  write_le32(zip, cd_offset);
  write_le16(zip, 0);

  riegeli::StringReader source_reader(zip);
  auto source_zip_reader = ZipReader::Create(&source_reader);
  ABSL_ASSERT_OK(source_zip_reader);
  std::vector<ZipReader::FileEntry> file_entries;
  while ((*source_zip_reader)->HasNext()) {
    auto entry = (*source_zip_reader)->Next();
    ABSL_ASSERT_OK(entry);
    file_entries.push_back(*entry);
  }
  ASSERT_EQ(file_entries.size(), 2);

  std::string output_s;
  riegeli::StringWriter<> destination(&output_s);

  auto copy_and_allocate_result = ZipEmbedder().CopyAndAllocate(
      &source_reader, file_entries, /*manifest_store_size=*/100, &destination);
  ABSL_ASSERT_OK(copy_and_allocate_result);
  (void)*copy_and_allocate_result;
  ASSERT_TRUE(destination.Close());

  riegeli::StringReader output_reader(output_s);
  auto out_zip_reader = ZipReader::Create(&output_reader);
  ABSL_ASSERT_OK(out_zip_reader);
  std::vector<ZipReader::FileEntry> out_file_entries;
  while ((*out_zip_reader)->HasNext()) {
    auto entry = (*out_zip_reader)->Next();
    ABSL_ASSERT_OK(entry);
    out_file_entries.push_back(*entry);
  }

  // Expect 3 entries in output: doc.xml, [Content_Types].xml, and manifest
  // store
  EXPECT_EQ(out_file_entries.size(), 3);

  // Check doc.xml retains extra_field_length and file_comment_length
  const ZipReader::FileEntry* doc_entry = nullptr;
  for (const auto& entry : out_file_entries) {
    if (entry.file_name == file2_name) {
      doc_entry = &entry;
    }
  }
  ASSERT_NE(doc_entry, nullptr);
  EXPECT_EQ(doc_entry->extra_field_length, file2_extra.size());
  EXPECT_EQ(doc_entry->file_comment_length, file2_cmt.size());

  // Check extra field data and comment data are present in output string
  EXPECT_THAT(output_s, HasSubstr(file1_extra));
  EXPECT_THAT(output_s, HasSubstr(file1_cmt));
  EXPECT_THAT(output_s, HasSubstr(file2_extra));
  EXPECT_THAT(output_s, HasSubstr(file2_cmt));
}

}  // namespace
}  // namespace credentio