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

#include <algorithm>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/base/nullability.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "absl/types/span.h"
#include "formats/zip/constants.h"
#include "riegeli/base/types.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/endian/endian_reading.h"

namespace credentio {
namespace {

// Holds information from the End of Central Directory Record (EOCD) needed to
// locate and read the central directory.
struct CentralDirectoryLocator {
  // Offset of the start of the central directory.
  uint32_t offset;
  // Total number of entries in the central directory.
  uint16_t entry_count;
};

// Reads the End of Central Directory (EOCD) record from the end of the input
// to locate the central directory.
absl::StatusOr<CentralDirectoryLocator> FindCentralDirectory(
    riegeli::Reader& input) {
  if (!input.SupportsSize() || !input.Size().has_value()) {
    return absl::InvalidArgumentError(
        "Input does not support size or size is unknown.");
  }
  int64_t file_size = *input.Size();

  if (file_size < kEocdRecordSize) {
    return absl::InvalidArgumentError("Input is too small to be a ZIP file.");
  }
  if (!input.Seek(file_size - kEocdRecordSize)) {
    return absl::InvalidArgumentError(
        "Failed to seek to the end of the input.");
  }

  CentralDirectoryLocator locator;
  std::string eocd_signature;
  if (!input.Read(kZipEndOfCentralDirectorySignature.size(), eocd_signature) ||
      eocd_signature != kZipEndOfCentralDirectorySignature) {
    return absl::InvalidArgumentError(absl::Substitute(
        "Invalid End of Central Directory record signature: $0",
        eocd_signature));
  }

  // Skip the following fields since they are not needed for locating the
  // central directory:
  //   - Number of this disk (2 bytes)
  //   - Number of the disk with the start of the central directory (2 bytes)
  //   - Total number of entries in the central directory on this disk (2 bytes)
  if (!input.Skip(6)) {
    return absl::InvalidArgumentError(
        "Failed to skip fields in the End of Central Directory record.");
  }

  // Total number of entries in the central directory (2 bytes, little-endian)
  if (!riegeli::ReadLittleEndian<uint16_t>(input, locator.entry_count)) {
    return absl::InvalidArgumentError(
        "Failed to read the total number of entries in the central directory.");
  }
  if (locator.entry_count == 0xffff) {
    return absl::UnimplementedError(
        "ZIP64 format is not supported. Central directory entry "
        "count must "
        "not be 0xffff.");
  }

  // Skip the size of the central directory (4 bytes)
  if (!input.Skip(4)) {
    return absl::InvalidArgumentError(
        "Failed to skip the size of the central directory.");
  }

  // Offset of the start of the central directory (4 bytes, little-endian)
  if (!riegeli::ReadLittleEndian<uint32_t>(input, locator.offset)) {
    return absl::InvalidArgumentError(
        "Failed to read the offset of the start of the central directory.");
  }
  if (locator.offset == 0xffffffff) {
    return absl::UnimplementedError(
        "ZIP64 format is not supported. Central Directory offset "
        "must not be "
        "0xffffffff.");
  }

  return locator;
}

// Reads the central directory entries, and extracts file header information.
absl::StatusOr<std::vector<ZipReader::FileEntry>>
ExtractFileEntriesFromCentralDirectory(riegeli::Reader& input,
                                       const CentralDirectoryLocator& locator) {
  if (!input.Seek(locator.offset) || input.pos() != locator.offset) {
    return absl::InvalidArgumentError(
        "Failed to seek to the start of the central directory.");
  }
  std::vector<ZipReader::FileEntry> file_entries;
  file_entries.reserve(locator.entry_count);

  for (int i = 0; i < locator.entry_count; ++i) {
    ZipReader::FileEntry entry;
    entry.central_directory_header_offset = input.pos();

    // Central directory signature (4 bytes)
    std::string central_directory_signature;
    if (!input.Read(kZipCentralDirectorySignature.size(),
                    central_directory_signature) ||
        central_directory_signature != kZipCentralDirectorySignature) {
      return absl::InvalidArgumentError(
          "Central directory signature not found.");
    }

    // Skip the following fields since they are not needed:
    //   - Version made by (2 bytes)
    //   - Version needed to extract (2 bytes)
    if (!input.Skip(4)) {
      return absl::InvalidArgumentError(
          "Failed to skip version fields in the central directory entry.");
    }

    // General purpose bit flag (2 bytes, little-endian)
    uint16_t general_purpose_bit_flag;
    if (!riegeli::ReadLittleEndian<uint16_t>(input, general_purpose_bit_flag)) {
      return absl::InvalidArgumentError(
          "Failed to read the general purpose bit flag.");
    }
    // Bit 3 in the general purpose bit flag indicates whether the file has a
    // data descriptor.
    entry.has_data_descriptor = general_purpose_bit_flag & 0x0008;
    entry.general_purpose_bit_flag = general_purpose_bit_flag;

    // Compression method (2 bytes, little-endian)
    if (!riegeli::ReadLittleEndian<uint16_t>(input, entry.compression_method)) {
      return absl::InvalidArgumentError(
          "Failed to read the compression method.");
    }

    // Skip the following fields since they are not needed:
    //   - Last mod file time (2 bytes)
    //   - Last mod file date (2 bytes)
    //   - CRC-32 (4 bytes)
    if (!input.Skip(8)) {
      return absl::InvalidArgumentError(
          "Failed to skip fields in the central directory entry.");
    }

    // Compressed size (4 bytes, little-endian)
    uint32_t compressed_size;
    if (!riegeli::ReadLittleEndian<uint32_t>(input, compressed_size)) {
      return absl::InvalidArgumentError("Failed to read the compressed size.");
    }
    if (compressed_size == 0xffffffff) {
      return absl::UnimplementedError(
          "ZIP64 format is not supported. Compressed size must "
          "not be "
          "0xffffffff.");
    }
    std::optional<riegeli::Position> input_size = input.Size();
    if (input_size.has_value() && compressed_size > *input_size) {
      return absl::InvalidArgumentError("Compressed size exceeds file limits.");
    }
    entry.file_range.length = compressed_size;

    // Uncompressed size (4 bytes, little-endian)
    if (!riegeli::ReadLittleEndian<uint32_t>(input, entry.uncompressed_size)) {
      return absl::InvalidArgumentError(
          "Failed to read the uncompressed size.");
    }

    // File name length (2 bytes, little-endian)
    uint16_t file_name_length;
    if (!riegeli::ReadLittleEndian<uint16_t>(input, file_name_length)) {
      return absl::InvalidArgumentError("Failed to read the file name length.");
    }

    // Extra field length (2 bytes, little-endian)
    uint16_t extra_field_length;
    if (!riegeli::ReadLittleEndian<uint16_t>(input, extra_field_length)) {
      return absl::InvalidArgumentError(
          "Failed to read the extra field length.");
    }
    entry.extra_field_length = extra_field_length;

    // File comment length (2 bytes, little-endian)
    uint16_t file_comment_length;
    if (!riegeli::ReadLittleEndian<uint16_t>(input, file_comment_length)) {
      return absl::InvalidArgumentError(
          "Failed to read the file comment length.");
    }
    entry.file_comment_length = file_comment_length;

    // Skip the following fields since they are not needed:
    //   - Disk number (2 bytes)
    //   - Internal file attributes (2 bytes)
    //   - External file attributes (4 bytes)
    if (!input.Skip(8)) {
      return absl::InvalidArgumentError(
          "Failed to skip fields in the central directory entry.");
    }

    // Local file header offset (4 bytes, little-endian)
    uint32_t local_file_header_offset;
    if (!riegeli::ReadLittleEndian<uint32_t>(input, local_file_header_offset)) {
      return absl::InvalidArgumentError(
          "Failed to read the local file header offset.");
    }
    entry.local_file_header_offset = local_file_header_offset;

    // File name (length is determined by `file_name_length`)
    if (!input.Read(file_name_length, entry.file_name)) {
      return absl::InvalidArgumentError("Failed to read the file name.");
    }

    // Skip extra field and file comment in central directory
    if (!input.Skip(extra_field_length + file_comment_length)) {
      return absl::InvalidArgumentError(
          "Failed to skip extra field and file comment in central directory.");
    }

    // The file data starts immediately after the local file header, which
    // consists of:
    //   - Fixed length fields (in total of 30 bytes)
    //   - File name (variable length)
    //   - Extra field (variable length)
    entry.file_range.offset = entry.local_file_header_offset +
                              kLocalFileHeaderFixedSize +
                              entry.file_name.size() + entry.extra_field_length;

    file_entries.push_back(entry);
  }

  // Reads and validates LFH signatures and ensures that the LFH
  // extra_field_length matches the Central Directory requirement.
  const int64_t cd_end_pos = input.pos();
  ABSL_ASSIGN_OR_RETURN(auto sorted_entries,
                        ZipReader::SortFileEntries(file_entries));

  for (const auto& entry : sorted_entries) {
    if (!input.Seek(entry.local_file_header_offset) ||
        input.pos() != entry.local_file_header_offset) {
      return absl::InvalidArgumentError(
          "Failed to seek to local file header offset.");
    }
    std::string lfh_sig;
    if (!input.Read(kZipLocalFileHeaderSignature.size(), lfh_sig) ||
        lfh_sig != kZipLocalFileHeaderSignature) {
      return absl::InvalidArgumentError("Invalid local file header signature.");
    }
    if (!input.Skip(2)) {
      return absl::InvalidArgumentError(
          "Failed to skip fields in local file header.");
    }
    uint16_t lfh_general_purpose_bit_flag;
    if (!riegeli::ReadLittleEndian<uint16_t>(input,
                                             lfh_general_purpose_bit_flag)) {
      return absl::InvalidArgumentError(
          "Failed to read general purpose bit flag from local file header.");
    }
    if ((entry.general_purpose_bit_flag & 0x0808) !=
        (lfh_general_purpose_bit_flag & 0x0808)) {
      return absl::InvalidArgumentError(
          "Local File Header bit flag does not match Central Directory.");
    }
    uint16_t lfh_compression_method;
    if (!riegeli::ReadLittleEndian<uint16_t>(input, lfh_compression_method)) {
      return absl::InvalidArgumentError(
          "Failed to read compression method from local file header.");
    }
    if (lfh_compression_method != entry.compression_method) {
      return absl::InvalidArgumentError(
          "Local File Header compression method does not match Central "
          "Directory.");
    }
    if (!input.Skip(8)) {
      return absl::InvalidArgumentError(
          "Failed to skip fields in local file header.");
    }
    uint32_t lfh_compressed_size;
    if (!riegeli::ReadLittleEndian<uint32_t>(input, lfh_compressed_size)) {
      return absl::InvalidArgumentError(
          "Failed to read compressed size from local file header.");
    }
    uint32_t lfh_uncompressed_size;
    if (!riegeli::ReadLittleEndian<uint32_t>(input, lfh_uncompressed_size)) {
      return absl::InvalidArgumentError(
          "Failed to read uncompressed size from local file header.");
    }
    if (!entry.has_data_descriptor) {
      if (lfh_compressed_size != entry.file_range.length ||
          lfh_uncompressed_size != entry.uncompressed_size) {
        return absl::InvalidArgumentError(
            "Local File Header sizes do not match Central Directory.");
      }
    } else {
      if ((lfh_compressed_size != 0 &&
           lfh_compressed_size != entry.file_range.length) ||
          (lfh_uncompressed_size != 0 &&
           lfh_uncompressed_size != entry.uncompressed_size)) {
        return absl::InvalidArgumentError(
            "Local File Header sizes do not match Central Directory.");
      }
    }
    uint16_t lfh_file_name_length;
    if (!riegeli::ReadLittleEndian<uint16_t>(input, lfh_file_name_length)) {
      return absl::InvalidArgumentError(
          "Failed to read file name length from local file header.");
    }
    if (lfh_file_name_length != entry.file_name.size()) {
      return absl::InvalidArgumentError(
          "Local File Header file name length does not match Central "
          "Directory.");
    }
    uint16_t lfh_extra_field_length;
    if (!riegeli::ReadLittleEndian<uint16_t>(input, lfh_extra_field_length)) {
      return absl::InvalidArgumentError(
          "Failed to read extra field length from local file header.");
    }
    if (lfh_extra_field_length != entry.extra_field_length) {
      return absl::InvalidArgumentError(
          "Local File Header extra field length does not match Central "
          "Directory.");
    }
  }

  if (!input.Seek(cd_end_pos) || input.pos() != cd_end_pos) {
    return absl::InvalidArgumentError(
        "Failed to seek back to central directory position.");
  }

  return sorted_entries;
}

}  // namespace

// Sorts ZIP file entries by header offset and checks that they do not have
// duplicate offsets or overlapping ranges to prevent ZIP bomb conditions.
absl::StatusOr<std::vector<ZipReader::FileEntry>> ZipReader::SortFileEntries(
    absl::Span<const ZipReader::FileEntry> file_entries) {
  std::vector<ZipReader::FileEntry> sorted_entries(file_entries.begin(),
                                                   file_entries.end());
  std::sort(sorted_entries.begin(), sorted_entries.end(),
            [](const ZipReader::FileEntry& a, const ZipReader::FileEntry& b) {
              return a.local_file_header_offset < b.local_file_header_offset;
            });

  int64_t previous_entry_end = 0;
  for (const auto& entry : sorted_entries) {
    if (entry.local_file_header_offset < previous_entry_end) {
      return absl::InvalidArgumentError(
          "Overlapping or duplicate ZIP entries detected.");
    }
    previous_entry_end = entry.file_range.offset + entry.file_range.length;
  }
  return sorted_entries;
}

absl::StatusOr<std::unique_ptr<ZipReader>> ZipReader::Create(
    riegeli::Reader* absl_nonnull input) {
  ABSL_ASSIGN_OR_RETURN(CentralDirectoryLocator locator,
                        FindCentralDirectory(*input));

  ABSL_ASSIGN_OR_RETURN(
      std::vector<FileEntry> file_entries,
      ExtractFileEntriesFromCentralDirectory(*input, locator));

  auto reader =
      absl::WrapUnique(new ZipReader(std::move(file_entries), locator.offset));
  return std::move(reader);
}

absl::StatusOr<ZipReader::FileEntry> ZipReader::Next() {
  if (!HasNext()) {
    return absl::OutOfRangeError("No more file entries to read.");
  }
  return *current_it_++;
}
}  // namespace credentio