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

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

#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "formats/byte_range.h"
#include "formats/zip/constants.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/string_writer.h"
#include "riegeli/endian/endian_writing.h"
#include "testing/jumbf_utils.h"
#include "testing/test_file_utils.h"

// Compatibility macros for OSS
#ifndef ASSERT_OK
#define ASSERT_OK(expr) ASSERT_THAT(expr, ::absl_testing::IsOk())
#endif

#ifndef ASSERT_OK_AND_ASSIGN
#define ASSERT_OK_AND_ASSIGN_CONCAT2(x, y) x##y
#define ASSERT_OK_AND_ASSIGN_CONCAT(x, y) ASSERT_OK_AND_ASSIGN_CONCAT2(x, y)

#define ASSERT_OK_AND_ASSIGN(lhs, rexpr) \
  ASSERT_OK_AND_ASSIGN_IMPL(lhs, rexpr, __COUNTER__)

#define ASSERT_OK_AND_ASSIGN_IMPL(lhs, rexpr, id)             \
  auto ASSERT_OK_AND_ASSIGN_CONCAT(status_or_, id) = (rexpr); \
  ASSERT_THAT(ASSERT_OK_AND_ASSIGN_CONCAT(status_or_, id),    \
              ::absl_testing::IsOk());                        \
  lhs = std::move(*ASSERT_OK_AND_ASSIGN_CONCAT(status_or_, id))
#endif

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::testing::AllOf;
using ::testing::Eq;
using ::testing::Field;
using ::testing::HasSubstr;

// A Zip file entry.
struct ZipEntry {
  std::string name;
  std::string content;
};

// Creates a minimal Zip file in memory.
// This is not a full-fledged Zip creator, but enough to satisfy ZipReader.
absl::StatusOr<std::string> CreateFakeZip(absl::Span<const ZipEntry> entries) {
  std::string buffer;
  riegeli::StringWriter writer(&buffer);

  // Write local file headers and data.
  std::vector<uint32_t> offsets;
  for (const auto& entry : entries) {
    offsets.push_back(writer.pos());

    writer.Write(kZipLocalFileHeaderSignature);
    // Version needed to extract (2.0)
    riegeli::WriteLittleEndian<uint16_t>(20, writer);
    // General purpose bit flag
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // Compression method (stored)
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // Last mod file time
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // Last mod file date
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // CRC-32 (0 for now, we don't check it)
    riegeli::WriteLittleEndian<uint32_t>(0, writer);
    // Compressed size
    riegeli::WriteLittleEndian<uint32_t>(entry.content.size(), writer);
    // Uncompressed size
    riegeli::WriteLittleEndian<uint32_t>(entry.content.size(), writer);
    // File name length
    riegeli::WriteLittleEndian<uint16_t>(entry.name.size(), writer);
    // Extra field length
    riegeli::WriteLittleEndian<uint16_t>(0, writer);

    // File name
    writer.Write(entry.name);
    // File data
    writer.Write(entry.content);
  }

  uint32_t central_directory_start_offset = writer.pos();

  // Write central directory.
  for (size_t i = 0; i < entries.size(); ++i) {
    const auto& entry = entries[i];
    writer.Write(kZipCentralDirectorySignature);
    // Version made by
    riegeli::WriteLittleEndian<uint16_t>(20, writer);
    // Version needed to extract
    riegeli::WriteLittleEndian<uint16_t>(20, writer);
    // General purpose bit flag
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // Compression method
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // Last mod file time
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // Last mod file date
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // CRC-32
    riegeli::WriteLittleEndian<uint32_t>(0, writer);
    // Compressed size
    riegeli::WriteLittleEndian<uint32_t>(entry.content.size(), writer);
    // Uncompressed size
    riegeli::WriteLittleEndian<uint32_t>(entry.content.size(), writer);
    // File name length
    riegeli::WriteLittleEndian<uint16_t>(entry.name.size(), writer);
    // Extra field length
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // File comment length
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // Disk number start
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // Internal file attributes
    riegeli::WriteLittleEndian<uint16_t>(0, writer);
    // External file attributes
    riegeli::WriteLittleEndian<uint32_t>(0, writer);
    // Relative offset of local header
    riegeli::WriteLittleEndian<uint32_t>(offsets[i], writer);

    // File name
    writer.Write(entry.name);
  }

  uint32_t central_directory_size =
      writer.pos() - central_directory_start_offset;

  // Write End of Central Directory (EOCD) record.
  writer.Write(kZipEndOfCentralDirectorySignature);
  // Number of this disk
  riegeli::WriteLittleEndian<uint16_t>(0, writer);
  // Number of the disk with the start of the central directory
  riegeli::WriteLittleEndian<uint16_t>(0, writer);
  // Total number of entries in the central directory on this disk
  riegeli::WriteLittleEndian<uint16_t>(entries.size(), writer);
  // Total number of entries in the central directory
  riegeli::WriteLittleEndian<uint16_t>(entries.size(), writer);
  // Size of the central directory
  riegeli::WriteLittleEndian<uint32_t>(central_directory_size, writer);
  // Offset of start of central directory with respect to the starting disk
  // number
  riegeli::WriteLittleEndian<uint32_t>(central_directory_start_offset, writer);
  // .ZIP file comment length
  riegeli::WriteLittleEndian<uint16_t>(0, writer);

  if (!writer.Close()) {
    return writer.status();
  }
  return buffer;
}

TEST(ZipExtractorTest, ExtractManifestStore) {
  ASSERT_OK_AND_ASSIGN(auto input,
                       credentio_testing::GetFileReader(
                           "c2pa/formats/zip/testing/asset_fake_c2pa.zip"));

  EXPECT_THAT(ZipExtractor().ExtractManifestStore(*input),
              IsOkAndHolds(HasSubstr("c2pa manifest")));
}

TEST(ZipExtractorTest,
     ExtractManifestStoreThrowsErrorIfMultipleC2paManifestStoresFound) {
  // Create a fake Zip file with duplicate manifest stores.
  ASSERT_OK_AND_ASSIGN(auto zip_file_contents,
                       CreateFakeZip({
                           {.name = "META-INF/content_credential.c2pa",
                            .content = "c2pa manifest 1"},
                           {.name = "META-INF/content_credential.c2pa",
                            .content = "c2pa manifest 2"},
                       }));
  riegeli::StringReader<> input(zip_file_contents);

  EXPECT_THAT(ZipExtractor().ExtractManifestStore(input),
              StatusIs(absl::StatusCode::kNotFound,
                       HasSubstr("Multiple manifest stores found")));
}

TEST(ZipExtractorTest, ExtractManifestStoreThrowsNotFoundErrorIfNoC2pa) {
  ASSERT_OK_AND_ASSIGN(auto input,
                       credentio_testing::GetFileReader(
                           "c2pa/formats/zip/testing/asset_no_c2pa.zip"));

  EXPECT_THAT(ZipExtractor().ExtractManifestStore(*input),
              StatusIs(absl::StatusCode::kNotFound,
                       HasSubstr("No manifest store found")));
}

TEST(ZipExtractorTest, ExtractManifestStoreLocation) {
  ASSERT_OK_AND_ASSIGN(auto input,
                       credentio_testing::GetFileReader(
                           "c2pa/formats/zip/testing/asset_fake_c2pa.zip"));
  ASSERT_OK_AND_ASSIGN(auto manifest_store_location,
                       ZipExtractor().ExtractManifestStoreLocation(
                           *input, {.requires_c2pa = true}));

  EXPECT_THAT(manifest_store_location,
              Optional(AllOf(Field(&ByteRange::offset, Eq(201)),
                             Field(&ByteRange::length, Eq(14)))));
}

TEST(ZipExtractorTest,
     ExtractManifestStoreLocationReturnsNulloptIfNoC2paAndNotRequired) {
  ASSERT_OK_AND_ASSIGN(auto input,
                       credentio_testing::GetFileReader(
                           "c2pa/formats/zip/testing/asset_no_c2pa.zip"));
  ASSERT_OK_AND_ASSIGN(auto manifest_store_location,
                       ZipExtractor().ExtractManifestStoreLocation(
                           *input, {.requires_c2pa = false}));

  EXPECT_EQ(manifest_store_location, std::nullopt);
}

TEST(ZipExtractorTest,
     ExtractManifestStoreLocationReturnsErrorIfNoC2paButRequired) {
  ASSERT_OK_AND_ASSIGN(auto input,
                       credentio_testing::GetFileReader(
                           "c2pa/formats/zip/testing/asset_no_c2pa.zip"));
  EXPECT_THAT(ZipExtractor().ExtractManifestStoreLocation(
                  *input, {.requires_c2pa = true}),
              StatusIs(absl::StatusCode::kNotFound,
                       HasSubstr("No manifest store found")));
}

TEST(ZipExtractorTest, ExtractBoxesReturnsUnimplemented) {
  ASSERT_OK_AND_ASSIGN(auto zip, CreateFakeZip({}));
  riegeli::StringReader<> input(zip);

  EXPECT_THAT(ZipExtractor().ExtractBoxes(input, {}),
              StatusIs(absl::StatusCode::kUnimplemented,
                       HasSubstr("ZIP does not have a concept of boxes")));
}

TEST(ZipExtractorTest, IsManifestStore) {
  EXPECT_TRUE(ZipExtractor().MightBeC2paManifestStore(
      CreateStartOfManifestStorePayload()));
  EXPECT_FALSE(ZipExtractor().MightBeC2paManifestStore("not a manifest store"));
}

TEST(ZipExtractorTest,
     ExtractManifestStoreThrowsErrorIfManifestStoreExceedsLimit) {
  std::string large_content(17 * 1024 * 1024, 'a');
  ASSERT_OK_AND_ASSIGN(auto zip,
                       CreateFakeZip({
                           {.name = "META-INF/content_credential.c2pa",
                            .content = std::move(large_content)},
                       }));
  riegeli::StringReader<> input(zip);

  EXPECT_THAT(ZipExtractor().ExtractManifestStore(input),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("Manifest store size exceeds maximum limit")));
}

TEST(ZipExtractorTest, ExtractManifestStoreThrowsErrorIfCompressed) {
  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;
  zip.append("PK\x03\x04", 4);
  write_le16(zip, 20);
  write_le16(zip, 0);
  write_le16(zip, 8);  // LFH compression method = 8 (deflated)
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, 10);
  write_le32(zip, 20);
  std::string filename = "META-INF/content_credential.c2pa";
  write_le16(zip, filename.size());
  write_le16(zip, 0);
  zip.append(filename);
  zip.append("abcdefghij");

  uint32_t cd_offset = zip.size();
  zip.append("PK\x01\x02", 4);
  write_le16(zip, 20);
  write_le16(zip, 20);
  write_le16(zip, 0);
  write_le16(zip, 8);  // CD compression method = 8 (deflated)
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, 10);
  write_le32(zip, 20);
  write_le16(zip, filename.size());
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, 0);
  zip.append(filename);

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

  riegeli::StringReader<> input(zip);
  EXPECT_THAT(ZipExtractor().ExtractManifestStore(input),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("Manifest store must not be compressed")));
}

TEST(ZipExtractorTest,
     ExtractManifestStoreThrowsErrorIfUncompressedSizeExceedsLimit) {
  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;
  zip.append("PK\x03\x04", 4);
  write_le16(zip, 20);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, 10);
  write_le32(zip, 17 * 1024 * 1024);  // Uncompressed size = 17MB
  std::string filename = "META-INF/content_credential.c2pa";
  write_le16(zip, filename.size());
  write_le16(zip, 0);
  zip.append(filename);
  zip.append("abcdefghij");

  uint32_t cd_offset = zip.size();
  zip.append("PK\x01\x02", 4);
  write_le16(zip, 20);
  write_le16(zip, 20);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, 10);
  write_le32(zip, 17 * 1024 * 1024);
  write_le16(zip, filename.size());
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le16(zip, 0);
  write_le32(zip, 0);
  write_le32(zip, 0);
  zip.append(filename);

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

  riegeli::StringReader<> input(zip);
  EXPECT_THAT(ZipExtractor().ExtractManifestStore(input),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("Manifest store size exceeds maximum limit")));
}

}  // namespace
}  // namespace credentio