#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;
}