#include "libarchive.h"

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

#include <algorithm>
#include <chrono>
#include <cctype>
#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;
}
