#include "libarchive.h"

#include <archive.h>
#include <archive_entry.h>

#include <algorithm>
#include <cctype>
#include <chrono>
#include <cstring>
#include <fstream>
#include <memory>
#include <sstream>
#include <system_error>
#include <unordered_map>
#include <unordered_set>

namespace fs = std::filesystem;

namespace {

constexpr std::size_t kArchiveBlockSize = 64 * 1024;

[[noreturn]] void throw_invalid_archive(const std::string &reason) {
  throw Invalid_Archive_Error("invalid upgrade archive: " + reason);
}

std::string archive_error_text(archive *a, const std::string &stage) {
  const char *err = archive_error_string(a);
  if (err == nullptr || *err == '\0') {
    return stage;
  }
  return stage + ": " + err;
}

void require_archive_ok(archive *a, int code, const std::string &stage) {
  if (code != ARCHIVE_OK) {
    throw std::runtime_error(archive_error_text(a, stage));
  }
}

std::string path_to_utf8(const fs::path &path) {
  const auto u8 = path.generic_u8string();
  return std::string(u8.begin(), u8.end());
}

bool path_is_within(const fs::path &base, const fs::path &target) {
  const fs::path canonical_base = fs::weakly_canonical(base);
  const fs::path canonical_target = fs::weakly_canonical(target);

  auto base_it = canonical_base.begin();
  auto target_it = canonical_target.begin();
  for (; base_it != canonical_base.end(); ++base_it, ++target_it) {
    if (target_it == canonical_target.end() || *base_it != *target_it) {
      return false;
    }
  }
  return true;
}

std::string normalize_archive_path(const std::string &raw,
                                   bool allow_empty = false) {
  if (raw.empty()) {
    if (allow_empty)
      return {};
    throw_invalid_archive("empty entry path");
  }
  if (raw.front() == '/' || raw.find('\\') != std::string::npos ||
      raw.find(':') != std::string::npos) {
    throw_invalid_archive("unsafe entry path: " + raw);
  }

  std::string normalized;
  std::size_t start = 0;
  while (start < raw.size()) {
    const std::size_t slash = raw.find('/', start);
    const std::size_t end = slash == std::string::npos ? raw.size() : slash;
    const std::string component = raw.substr(start, end - start);
    if (component.empty() || component == "." || component == "..") {
      throw_invalid_archive("unsafe entry path: " + raw);
    }
    if (!normalized.empty())
      normalized.push_back('/');
    normalized += component;
    if (slash == std::string::npos)
      break;
    start = slash + 1;
    if (start == raw.size())
      break;
  }
  return normalized;
}

std::string entry_path(archive_entry *entry) {
  const char *name = archive_entry_pathname_utf8(entry);
  if (name == nullptr) {
    name = archive_entry_pathname(entry);
  }
  if (name == nullptr) {
    throw_invalid_archive("entry path is unavailable");
  }
  return normalize_archive_path(name);
}

bool is_directory_entry(archive_entry *entry) {
  return archive_entry_filetype(entry) == AE_IFDIR;
}

bool is_regular_entry(archive_entry *entry) {
  const auto type = archive_entry_filetype(entry);
  return type == AE_IFREG || type == 0;
}

void reject_unsupported_entry_type(archive_entry *entry,
                                   const std::string &path) {
  const auto type = archive_entry_filetype(entry);
  if (type == AE_IFLNK) {
    throw_invalid_archive("symbolic link is not allowed: " + path);
  }
  if (type != AE_IFDIR && type != AE_IFREG && type != 0) {
    throw_invalid_archive("special file is not allowed: " + path);
  }
}

void setup_archive_reader(archive *a) {
  require_archive_ok(a, archive_read_support_filter_all(a),
                     "archive_read_support_filter_all failed");
  require_archive_ok(a, archive_read_support_format_all(a),
                     "archive_read_support_format_all failed");
}

void consume_entry_data_for_validation(archive *a, const std::string &path,
                                       const Archive_Layout_Rules &rules,
                                       std::uint64_t &total_uncompressed_size) {
  std::uint64_t entry_size = 0;
  const void *block = nullptr;
  std::size_t block_size = 0;
  la_int64_t offset = 0;

  while (true) {
    const int result = archive_read_data_block(a, &block, &block_size, &offset);
    if (result == ARCHIVE_EOF) {
      break;
    }
    if (result != ARCHIVE_OK) {
      throw_invalid_archive(
          archive_error_text(a, "cannot read entry: " + path));
    }

    if (block_size > rules.max_entry_uncompressed_size ||
        entry_size > rules.max_entry_uncompressed_size - block_size) {
      throw_invalid_archive("entry size exceeds limit: " + path);
    }
    if (block_size > rules.max_total_uncompressed_size ||
        total_uncompressed_size >
            rules.max_total_uncompressed_size - block_size) {
      throw_invalid_archive("uncompressed size exceeds limit");
    }

    entry_size += block_size;
    total_uncompressed_size += block_size;
  }
}

void consume_entry_data_for_validation(archive *a, archive_entry *entry,
                                       const std::string &path,
                                       const Archive_Layout_Rules &rules,
                                       std::uint64_t &total_uncompressed_size) {
  if (archive_entry_size_is_set(entry)) {
    const auto expected_size = archive_entry_size(entry);
    if (expected_size < 0) {
      throw_invalid_archive("negative entry size: " + path);
    }
    const auto size = static_cast<std::uint64_t>(expected_size);
    if (size > rules.max_entry_uncompressed_size ||
        size > rules.max_total_uncompressed_size) {
      throw_invalid_archive("entry size exceeds limit: " + path);
    }
  }

  consume_entry_data_for_validation(a, path, rules, total_uncompressed_size);
}

std::unordered_set<std::string>
make_allowed_entries(const std::vector<std::string> &required,
                     const std::vector<std::string> &optional) {
  std::unordered_set<std::string> result;
  auto add = [&result](const std::vector<std::string> &entries) {
    for (const auto &entry : entries) {
      if (entry.empty() || entry.find('/') != std::string::npos ||
          entry.find('\\') != std::string::npos) {
        throw_invalid_archive("invalid layout rule entry: " + entry);
      }
      result.insert(entry);
    }
  };
  add(required);
  add(optional);
  return result;
}

struct Validated_Archive_Entry {
  std::string path;
  bool is_directory = false;
};

struct Archive_Read_Deleter {
  void operator()(archive *a) const noexcept {
    if (a != nullptr) {
      archive_read_free(a);
    }
  }
};

struct Archive_Write_Deleter {
  void operator()(archive *a) const noexcept {
    if (a != nullptr) {
      archive_write_free(a);
    }
  }
};

using Read_Archive_Ptr = std::unique_ptr<archive, Archive_Read_Deleter>;
using Write_Archive_Ptr = std::unique_ptr<archive, Archive_Write_Deleter>;

Read_Archive_Ptr make_archive_reader() {
  Read_Archive_Ptr reader(archive_read_new());
  if (!reader) {
    throw std::runtime_error("archive_read_new failed");
  }
  setup_archive_reader(reader.get());
  return reader;
}

bool is_safe_archive_name_token(const std::string &value) {
  if (value.empty()) {
    return false;
  }
  return std::ranges::all_of(value, [](unsigned char ch) {
    return std::isalnum(ch) != 0 || ch == '_' || ch == '-' || ch == '.';
  });
}

void validate_archive_write_options(const Archive_Write_Options &options) {
  if (!is_safe_archive_name_token(options.format)) {
    throw std::invalid_argument("invalid archive format: " + options.format);
  }
  if (!options.filter.empty() && options.filter != "none" &&
      !is_safe_archive_name_token(options.filter)) {
    throw std::invalid_argument("invalid archive filter: " + options.filter);
  }
  if (!is_safe_archive_name_token(options.extension) ||
      options.extension.find('/') != std::string::npos) {
    throw std::invalid_argument("invalid archive extension: " +
                                options.extension);
  }
}

void configure_archive_writer(archive *writer,
                              const Archive_Write_Options &options) {
  validate_archive_write_options(options);

  if (!options.filter.empty() && options.filter != "none") {
    require_archive_ok(
        writer,
        archive_write_add_filter_by_name(writer, options.filter.c_str()),
        "archive_write_add_filter_by_name failed: " + options.filter);
  }

  require_archive_ok(
      writer, archive_write_set_format_by_name(writer, options.format.c_str()),
      "archive_write_set_format_by_name failed: " + options.format);

  if (options.format == "zip") {
    require_archive_ok(writer,
                       archive_write_zip_set_compression_deflate(writer),
                       "archive_write_zip_set_compression_deflate failed");
  }
}

Write_Archive_Ptr make_archive_writer(const Archive_Write_Options &options) {
  Write_Archive_Ptr writer(archive_write_new());
  if (!writer) {
    throw std::runtime_error("archive_write_new failed");
  }
  configure_archive_writer(writer.get(), options);
  return writer;
}

void extract_archive(archive *a, const fs::path &out_dir) {
  fs::create_directories(out_dir);
  const fs::path base = fs::weakly_canonical(out_dir);

  archive_entry *entry = nullptr;
  while (true) {
    const int result = archive_read_next_header(a, &entry);
    if (result == ARCHIVE_EOF) {
      break;
    }
    if (result != ARCHIVE_OK) {
      throw std::runtime_error(
          archive_error_text(a, "archive_read_next_header failed"));
    }

    const std::string normalized = entry_path(entry);
    reject_unsupported_entry_type(entry, normalized);

    const fs::path out_path = (base / fs::path(normalized)).lexically_normal();
    if (!path_is_within(base, out_path)) {
      throw std::runtime_error("archive-slip detected: " + normalized);
    }

    if (is_directory_entry(entry)) {
      fs::create_directories(out_path);
      continue;
    }

    if (!is_regular_entry(entry)) {
      throw std::runtime_error("unsupported archive entry type: " + normalized);
    }

    fs::create_directories(out_path.parent_path());
    std::ofstream out(out_path, std::ios::binary);
    if (!out) {
      throw std::runtime_error("open output failed: " + path_to_utf8(out_path));
    }

    const void *block = nullptr;
    std::size_t block_size = 0;
    la_int64_t offset = 0;
    while (true) {
      const int read_result =
          archive_read_data_block(a, &block, &block_size, &offset);
      if (read_result == ARCHIVE_EOF) {
        break;
      }
      if (read_result != ARCHIVE_OK) {
        throw std::runtime_error(
            archive_error_text(a, "archive_read_data_block failed"));
      }
      out.write(static_cast<const char *>(block),
                static_cast<std::streamsize>(block_size));
      if (!out) {
        throw std::runtime_error("write output failed: " +
                                 path_to_utf8(out_path));
      }
    }
  }
}

struct Ignore_Rules {
  std::vector<fs::path> directories;
  std::unordered_set<std::string> files;
  std::unordered_set<std::string> extensions;
};

Ignore_Rules
parse_ignore_rules(const std::vector<std::string> &ignore_patterns) {
  Ignore_Rules rules;
  for (const auto &pattern : ignore_patterns) {
    if (pattern.empty())
      continue;
    if (pattern.ends_with('/')) {
      rules.directories.emplace_back(
          fs::path(pattern.substr(0, pattern.size() - 1)).lexically_normal());
    } else if (pattern[0] == '.') {
      rules.extensions.insert(pattern);
    } else {
      rules.files.insert(pattern);
    }
  }
  return rules;
}

bool should_skip_directory(const fs::path &rel_path,
                           const Ignore_Rules &rules) {
  for (const auto &ignored : rules.directories) {
    const auto rel_native = rel_path.native();
    const auto ignored_native = ignored.native();
    if (rel_path == ignored ||
        rel_native.starts_with(ignored_native +
                               fs::path::preferred_separator)) {
      return true;
    }
  }
  return false;
}

void add_file_to_archive(archive *writer, const fs::path &base,
                         const fs::path &abs_path, const fs::path &rel_path) {
  const std::uintmax_t file_size = fs::file_size(abs_path);
  const std::string archive_name = path_to_utf8(rel_path);
  if (archive_name.empty() || archive_name == ".." ||
      archive_name.starts_with("../") ||
      archive_name.find("/../") != std::string::npos ||
      archive_name.ends_with("/..")) {
    throw std::runtime_error("invalid path in archive: " + archive_name);
  }

  std::unique_ptr<archive_entry, decltype(&archive_entry_free)> entry(
      archive_entry_new(), archive_entry_free);
  if (!entry) {
    throw std::runtime_error("archive_entry_new failed");
  }
  archive_entry_set_pathname_utf8(entry.get(), archive_name.c_str());
  archive_entry_set_filetype(entry.get(), AE_IFREG);
  archive_entry_set_perm(entry.get(), 0644);
  archive_entry_set_size(entry.get(), static_cast<la_int64_t>(file_size));

  const auto now =
      std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
  archive_entry_set_mtime(entry.get(), now, 0);

  require_archive_ok(writer, archive_write_header(writer, entry.get()),
                     "archive_write_header failed");

  std::ifstream in(abs_path, std::ios::binary);
  if (!in) {
    throw std::runtime_error("failed to open input file: " +
                             path_to_utf8(abs_path));
  }

  char buffer[kArchiveBlockSize];
  while (in) {
    in.read(buffer, sizeof(buffer));
    const auto count = in.gcount();
    if (count <= 0) {
      break;
    }
    const la_ssize_t written =
        archive_write_data(writer, buffer, static_cast<std::size_t>(count));
    if (written < 0 || written != count) {
      throw std::runtime_error(
          archive_error_text(writer, "archive_write_data failed"));
    }
  }
  if (!in.eof()) {
    throw std::runtime_error("failed to read input file: " +
                             path_to_utf8(abs_path));
  }

  require_archive_ok(writer, archive_write_finish_entry(writer),
                     "archive_write_finish_entry failed");
  (void)base;
}

void write_directory_to_archive(
    archive *writer, const fs::path &dir_path,
    const std::vector<std::string> &ignore_patterns) {
  if (!fs::exists(dir_path)) {
    throw std::runtime_error("source path does not exist: " +
                             path_to_utf8(dir_path));
  }
  if (!fs::is_directory(dir_path)) {
    throw std::runtime_error("source path exists but is not a directory: " +
                             path_to_utf8(dir_path));
  }

  const fs::path base = fs::weakly_canonical(dir_path);
  const Ignore_Rules ignore_rules = parse_ignore_rules(ignore_patterns);

  fs::recursive_directory_iterator it(base);
  const fs::recursive_directory_iterator end;
  for (; it != end; ++it) {
    const fs::path abs_path = it->path();
    const fs::path rel_path = fs::relative(abs_path, base);

    if (it->is_directory()) {
      if (should_skip_directory(rel_path, ignore_rules)) {
        it.disable_recursion_pending();
      }
      continue;
    }

    if (!it->is_regular_file()) {
      continue;
    }

    const std::string filename = abs_path.filename().string();
    const std::string extension = abs_path.extension().string();
    if (ignore_rules.files.contains(filename) ||
        (!extension.empty() && ignore_rules.extensions.contains(extension))) {
      continue;
    }

    add_file_to_archive(writer, base, abs_path, rel_path);
  }
}

la_ssize_t archive_string_write_callback(archive *, void *client_data,
                                         const void *buffer,
                                         std::size_t length) {
  auto *out = static_cast<std::string *>(client_data);
  out->append(static_cast<const char *>(buffer), length);
  return static_cast<la_ssize_t>(length);
}

int archive_noop_open_close_callback(archive *, void *) { return ARCHIVE_OK; }

void validate_archive_layout_from_memory_or_throw(
    const char *data, std::size_t size, const Archive_Layout_Rules &rules) {
  if (data == nullptr || size == 0) {
    throw_invalid_archive("empty archive");
  }
  if (size > rules.max_archive_size) {
    throw_invalid_archive("archive size exceeds limit");
  }

  auto reader = make_archive_reader();
  require_archive_ok(reader.get(),
                     archive_read_open_memory(reader.get(), data, size),
                     "archive_read_open_memory failed");

  std::vector<Validated_Archive_Entry> entries;
  std::unordered_set<std::string> unique_paths;
  std::uint64_t total_uncompressed_size = 0;
  std::size_t entry_count = 0;

  archive_entry *entry = nullptr;
  while (true) {
    const int result = archive_read_next_header(reader.get(), &entry);
    if (result == ARCHIVE_EOF) {
      break;
    }
    if (result != ARCHIVE_OK) {
      throw_invalid_archive(
          archive_error_text(reader.get(), "cannot read archive header"));
    }

    ++entry_count;
    if (entry_count > rules.max_entry_count) {
      throw_invalid_archive("entry count exceeds limit");
    }

    const std::string path = entry_path(entry);
    if (!unique_paths.insert(path).second) {
      throw_invalid_archive("duplicate entry path: " + path);
    }

    if (archive_entry_is_encrypted(entry)) {
      throw_invalid_archive("encrypted entry is not allowed: " + path);
    }

    reject_unsupported_entry_type(entry, path);

    const bool directory = is_directory_entry(entry);
    if (!directory) {
      if (!is_regular_entry(entry)) {
        throw_invalid_archive("special file is not allowed: " + path);
      }
      consume_entry_data_for_validation(reader.get(), entry, path, rules,
                                        total_uncompressed_size);
    }

    entries.push_back({path, directory});
  }

  for (const auto &rule : rules.directories) {
    const std::string rule_directory =
        normalize_archive_path(rule.directory, true);
    const auto allowed_files =
        make_allowed_entries(rule.required_files, rule.optional_files);
    const auto allowed_directories = make_allowed_entries(
        rule.required_directories, rule.optional_directories);
    std::unordered_map<std::string, bool> direct_children;

    for (const auto &entry : entries) {
      std::string relative;
      if (rule_directory.empty()) {
        relative = entry.path;
      } else {
        const std::string prefix = rule_directory + "/";
        if (!entry.path.starts_with(prefix))
          continue;
        relative = entry.path.substr(prefix.size());
      }

      const std::size_t slash = relative.find('/');
      const std::string child = relative.substr(0, slash);
      const bool is_directory =
          slash != std::string::npos || entry.is_directory;
      const auto [it, inserted] = direct_children.emplace(child, is_directory);
      if (!inserted && it->second != is_directory) {
        throw_invalid_archive("entry is both file and directory: " + child);
      }
    }

    for (const auto &[child, is_directory] : direct_children) {
      const auto &allowed = is_directory ? allowed_directories : allowed_files;
      if (!allowed.contains(child)) {
        throw_invalid_archive("unexpected entry under " + rule_directory +
                              ": " + child);
      }
    }
    for (const auto &required : rule.required_files) {
      const auto it = direct_children.find(required);
      if (it == direct_children.end() || it->second) {
        throw_invalid_archive("missing required file under " + rule_directory +
                              ": " + required);
      }
    }
    for (const auto &required : rule.required_directories) {
      const auto it = direct_children.find(required);
      if (it == direct_children.end() || !it->second) {
        throw_invalid_archive("missing required directory under " +
                              rule_directory + ": " + required);
      }
    }
  }
}

} // namespace

Psc::expected<void, std::string>
validate_archive_layout_from_memory(const char *data, std::size_t size,
                                    const Archive_Layout_Rules &rules) {
  try {
    validate_archive_layout_from_memory_or_throw(data, size, rules);
    return {};
  } catch (const Invalid_Archive_Error &e) {
    return Psc::unexpected<std::string>(e.what());
  }
}

void unarchive_file(const fs::path &zipPath, const fs::path &outDir) {
  auto reader = make_archive_reader();
#ifdef _WIN32
  require_archive_ok(reader.get(),
                     archive_read_open_filename_w(reader.get(),
                                                  zipPath.wstring().c_str(),
                                                  kArchiveBlockSize),
                     "archive_read_open_filename_w failed");
#else
  require_archive_ok(reader.get(),
                     archive_read_open_filename(reader.get(), zipPath.c_str(),
                                                kArchiveBlockSize),
                     "archive_read_open_filename failed");
#endif
  extract_archive(reader.get(), outDir);
}

void unarchive_from_memory(const char *data, size_t size,
                           const fs::path &outDir) {
  if (data == nullptr || size == 0) {
    throw std::runtime_error("empty archive");
  }

  auto reader = make_archive_reader();
  require_archive_ok(reader.get(),
                     archive_read_open_memory(reader.get(), data, size),
                     "archive_read_open_memory failed");
  extract_archive(reader.get(), outDir);
}

void archive_directory(const fs::path &dirPath, const fs::path &archive_path,
                       const Archive_Write_Options &options) {
  auto writer = make_archive_writer(options);
#ifdef _WIN32
  require_archive_ok(writer.get(),
                     archive_write_open_filename_w(
                         writer.get(), archive_path.wstring().c_str()),
                     "archive_write_open_filename_w failed");
#else
  require_archive_ok(
      writer.get(),
      archive_write_open_filename(writer.get(), archive_path.c_str()),
      "archive_write_open_filename failed");
#endif
  write_directory_to_archive(writer.get(), dirPath, {});
  require_archive_ok(writer.get(), archive_write_close(writer.get()),
                     "archive_write_close failed");
}

std::string
archive_directory_to_memory(const fs::path &dirPath,
                            const std::vector<std::string> &ignorePatterns,
                            const Archive_Write_Options &options) {
  std::string out;
  auto writer = make_archive_writer(options);
  require_archive_ok(writer.get(),
                     archive_write_open(writer.get(), &out,
                                        archive_noop_open_close_callback,
                                        archive_string_write_callback,
                                        archive_noop_open_close_callback),
                     "archive_write_open failed");
  write_directory_to_archive(writer.get(), dirPath, ignorePatterns);
  require_archive_ok(writer.get(), archive_write_close(writer.get()),
                     "archive_write_close failed");
  return out;
}