/*
MIT License

Copyright (c) 2026 Poul Sander

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

*/

#include "archive.h"

#include <zlib.h>

#include <array>
#include <cstring>

namespace {

constexpr size_t BLOCK_SIZE = 512;
constexpr size_t NAME_FIELD = 100;   // ustar name field width
constexpr size_t PREFIX_FIELD = 155; // ustar prefix field width
constexpr int GZIP_LEVEL = 6;
constexpr int GZIP_WINDOW_BITS = 15 + 16; // 16 selects the gzip wrapper

// GNU extension typeflags and the magic name their header carries.
constexpr char GNU_LONGNAME = 'L';
constexpr char GNU_LONGLINK = 'K';
const char *const GNU_LONG_MAGIC_NAME = "././@LongLink";

struct TarHeader {
    char name[100];
    char mode[8];
    char uid[8];
    char gid[8];
    char size[12];
    char mtime[12];
    char chksum[8];
    char typeflag;
    char linkname[100];
    char magic[6];
    char version[2];
    char uname[32];
    char gname[32];
    char devmajor[8];
    char devminor[8];
    char prefix[155];
    char pad[12];
};
static_assert(sizeof(TarHeader) == BLOCK_SIZE, "tar header must be 512 bytes");

// Writes `value` as zero-padded octal into a fixed-width field followed by a
// trailing NUL (the conventional ustar encoding).
void write_octal(char *field, size_t width, uint64_t value) {
    for (size_t i = width - 1; i-- > 0;) {
        field[i] = static_cast<char>('0' + (value & 7));
        value >>= 3;
    }
    field[width - 1] = '\0';
}

void set_field(char *field, size_t width, const std::string &value) {
    std::memcpy(field, value.data(), std::min(width, value.size()));
}

unsigned header_checksum(const TarHeader &h) {
    const unsigned char *bytes = reinterpret_cast<const unsigned char *>(&h);
    unsigned sum = 0;
    for (size_t i = 0; i < BLOCK_SIZE; ++i)
        sum += bytes[i];
    return sum;
}

// Fits `path` into the ustar name/prefix fields by splitting at a '/'. Returns
// false when no split keeps both parts within their limits, signalling that the
// GNU long-name extension is required.
bool split_ustar_path(const std::string &path, std::string &prefix, std::string &name) {
    if (path.size() <= NAME_FIELD) {
        prefix.clear();
        name = path;
        return true;
    }
    if (path.size() > NAME_FIELD + PREFIX_FIELD + 1)
        return false;
    const size_t earliest = path.size() - NAME_FIELD - 1;
    for (size_t i = earliest; i < path.size(); ++i) {
        if (path[i] != '/')
            continue;
        if (i > PREFIX_FIELD)
            return false;
        prefix = path.substr(0, i);
        name = path.substr(i + 1);
        return name.size() <= NAME_FIELD;
    }
    return false;
}

} // namespace

TarGzWriter::TarGzWriter(const std::string &path) {
    file_ = std::fopen(path.c_str(), "wb");
    if (!file_)
        throw ArchiveError("cannot open archive file for writing: " + path);

    z_stream *strm = new z_stream{};
    if (deflateInit2(strm, GZIP_LEVEL, Z_DEFLATED, GZIP_WINDOW_BITS, 8,
                     Z_DEFAULT_STRATEGY) != Z_OK) {
        delete strm;
        std::fclose(file_);
        file_ = nullptr;
        throw ArchiveError("failed to initialise gzip stream");
    }
    stream_ = strm;
}

TarGzWriter::~TarGzWriter() {
    if (stream_) {
        deflateEnd(static_cast<z_stream *>(stream_));
        delete static_cast<z_stream *>(stream_);
    }
    if (file_)
        std::fclose(file_);
}

void TarGzWriter::deflate_chunk(const void *data, size_t size, int flush) {
    z_stream *strm = static_cast<z_stream *>(stream_);
    strm->next_in = reinterpret_cast<Bytef *>(const_cast<void *>(data));
    strm->avail_in = static_cast<uInt>(size);

    std::array<unsigned char, 1 << 16> out{};
    do {
        strm->next_out = out.data();
        strm->avail_out = static_cast<uInt>(out.size());
        const int ret = deflate(strm, flush);
        if (ret == Z_STREAM_ERROR)
            throw ArchiveError("gzip compression failed");
        const size_t produced = out.size() - strm->avail_out;
        if (produced > 0 && std::fwrite(out.data(), 1, produced, file_) != produced)
            throw ArchiveError("failed to write archive data");
    } while (strm->avail_out == 0);
}

void TarGzWriter::write_padding(uint64_t size) {
    const size_t remainder = static_cast<size_t>(size % BLOCK_SIZE);
    if (remainder == 0)
        return;
    std::array<char, BLOCK_SIZE> zeros{};
    deflate_chunk(zeros.data(), BLOCK_SIZE - remainder, Z_NO_FLUSH);
}

void TarGzWriter::write_gnu_long(const std::string &value, char typeflag) {
    TarHeader h{};
    set_field(h.name, sizeof(h.name), GNU_LONG_MAGIC_NAME);
    write_octal(h.mode, sizeof(h.mode), 0);
    write_octal(h.uid, sizeof(h.uid), 0);
    write_octal(h.gid, sizeof(h.gid), 0);
    write_octal(h.size, sizeof(h.size), value.size() + 1);
    write_octal(h.mtime, sizeof(h.mtime), 0);
    h.typeflag = typeflag;
    // Old GNU magic: "ustar  \0" spanning the magic and version fields.
    std::memcpy(h.magic, "ustar ", sizeof(h.magic));
    h.version[0] = ' ';
    h.version[1] = '\0';
    std::memset(h.chksum, ' ', sizeof(h.chksum));
    write_octal(h.chksum, 7, header_checksum(h));
    h.chksum[7] = ' ';
    deflate_chunk(&h, BLOCK_SIZE, Z_NO_FLUSH);

    deflate_chunk(value.c_str(), value.size() + 1, Z_NO_FLUSH);
    write_padding(value.size() + 1);
}

void TarGzWriter::write_header(const std::string &path, uint32_t mode,
                               std::time_t mtime, uint64_t size, char typeflag,
                               const std::string &linkname) {
    std::string prefix;
    std::string name;
    const bool fits = split_ustar_path(path, prefix, name);
    if (!fits) {
        write_gnu_long(path, GNU_LONGNAME);
        prefix.clear();
        name = path.substr(0, NAME_FIELD);
    }
    if (linkname.size() > sizeof(TarHeader::linkname))
        write_gnu_long(linkname, GNU_LONGLINK);

    TarHeader h{};
    set_field(h.name, sizeof(h.name), name);
    write_octal(h.mode, sizeof(h.mode), mode & 07777);
    write_octal(h.uid, sizeof(h.uid), 0);
    write_octal(h.gid, sizeof(h.gid), 0);
    write_octal(h.size, sizeof(h.size), size);
    write_octal(h.mtime, sizeof(h.mtime), static_cast<uint64_t>(mtime));
    h.typeflag = typeflag;
    set_field(h.linkname, sizeof(h.linkname), linkname);
    std::memcpy(h.magic, "ustar", 5); // trailing NUL stays from zero-init
    h.version[0] = '0';
    h.version[1] = '0';
    set_field(h.prefix, sizeof(h.prefix), prefix);

    std::memset(h.chksum, ' ', sizeof(h.chksum));
    write_octal(h.chksum, 7, header_checksum(h));
    h.chksum[7] = ' ';
    deflate_chunk(&h, BLOCK_SIZE, Z_NO_FLUSH);
}

void TarGzWriter::add_file(const std::string &path, uint32_t mode,
                           std::time_t mtime, const char *data, uint64_t size) {
    write_header(path, mode, mtime, size, '0', "");
    if (size > 0)
        deflate_chunk(data, static_cast<size_t>(size), Z_NO_FLUSH);
    write_padding(size);
}

void TarGzWriter::add_symlink(const std::string &path, const std::string &target,
                              std::time_t mtime) {
    write_header(path, 0777, mtime, 0, '2', target);
}

void TarGzWriter::add_dir(const std::string &path, std::time_t mtime) {
    const std::string with_slash = path.empty() || path.back() == '/' ? path : path + "/";
    write_header(with_slash, 0755, mtime, 0, '5', "");
}

void TarGzWriter::finish() {
    if (finished_)
        return;
    std::array<char, BLOCK_SIZE * 2> trailer{}; // two zero blocks end the archive
    deflate_chunk(trailer.data(), trailer.size(), Z_NO_FLUSH);
    deflate_chunk(nullptr, 0, Z_FINISH);
    if (std::fflush(file_) != 0)
        throw ArchiveError("failed to flush archive");
    finished_ = true;
}