/*
 * Copyright (C) 2016-2022 Matthias Klumpp <matthias@tenstral.net>
 *
 * Licensed under the GNU Lesser General Public License Version 3
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the license, or
 * (at your option) any later version.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this software.  If not, see <http://www.gnu.org/licenses/>.
 */

module asgen.zarchive;

import core.stdc.string : strerror;
import std.stdio;
import std.string;
import std.regex;
import std.conv : to;
import std.path : buildNormalizedPath;
import std.array : appender;
import std.typecons : RefCounted, RefCountedAutoInitialize;
import std.concurrency : Generator, yield;
static import std.file;

import asgen.logging;
import asgen.utils : GENERIC_BUFFER_SIZE;

import asgen.bindings.libarchive;

private immutable DEFAULT_BLOCK_SIZE = 65_536;

enum ArchiveType {
    GZIP,
    XZ
}

private const(char)[] getArchiveErrorMessage (archive* ar)
{
    return fromStringz(archive_error_string(ar));
}

private string readArchiveData (archive* ar, string name = null)
{
    archive_entry* ae;
    int ret;
    char[GENERIC_BUFFER_SIZE] buff;
    auto data = appender!string;
    data.reserve(GENERIC_BUFFER_SIZE / 4);

    ret = archive_read_next_header(ar, &ae);

    if (ret == ARCHIVE_EOF)
        return data.data;

    if (ret != ARCHIVE_OK) {
        if (name is null)
            throw new Exception(format("Unable to read header of compressed data: %s", getArchiveErrorMessage(ar)));
        else
            throw new Exception(format("Unable to read header of compressed file '%s': %s", name, getArchiveErrorMessage(
                    ar)));
    }

    while (true) {
        immutable size = archive_read_data(ar, buff.ptr, buff.length);
        if (size < 0) {
            if (name is null)
                throw new Exception(format("Failed to read compressed data: %s", getArchiveErrorMessage(ar)));
            else
                throw new Exception(format("Failed to read data from '%s': %s", name, getArchiveErrorMessage(ar)));
        }

        if (size == 0)
            break;

        data ~= buff[0 .. size];
    }

    return data.data;
}

string decompressFile (string fname)
{
    int ret;

    archive* ar = archive_read_new();
    scope (exit)
        archive_read_free (ar);

    archive_read_support_format_raw(ar);
    archive_read_support_format_empty(ar);
    archive_read_support_filter_all(ar);

    ret = archive_read_open_filename(ar, toStringz(fname), DEFAULT_BLOCK_SIZE);
    if (ret != ARCHIVE_OK) {
        auto ret_errno = archive_errno(ar);
        auto ret_strerr = fromStringz(strerror(ret_errno));
        throw new Exception(format("Unable to open compressed file '%s': %s. error: %s (%d)",
                fname,
                getArchiveErrorMessage(ar),
                ret_strerr,
                ret_errno));
    }

    return readArchiveData(ar, fname);
}

string decompressData (ref ubyte[] data)
{
    int ret;

    archive* ar = archive_read_new();
    scope (exit)
        archive_read_free (ar);

    archive_read_support_filter_all(ar);
    archive_read_support_format_empty(ar);
    archive_read_support_format_raw(ar);

    auto dSize = ubyte.sizeof * data.length;
    ret = archive_read_open_memory(ar, data.ptr, dSize);
    if (ret != ARCHIVE_OK)
        throw new Exception(format("Unable to open compressed data: %s", getArchiveErrorMessage(ar)));

    return readArchiveData(ar);
}

struct ArchiveDecompressor {

private:
    string archive_fname;

    const(ubyte)[] readEntry (archive* ar)
    {
        const(void)* buff = null;
        size_t size = 0UL;
        long offset = 0;
        auto res = appender!(ubyte[]);
        res.reserve(GENERIC_BUFFER_SIZE / 4);

        while (archive_read_data_block(ar, &buff, &size, &offset) == ARCHIVE_OK) {
            res ~= cast(ubyte[]) buff[0 .. size];
        }

        return res.data;
    }

    void extractEntryTo (archive* ar, string fname)
    {
        const(void)* buff = null;
        size_t size = 0UL;
        long offset = 0;
        long output_offset = 0;

        auto f = File(fname, "w"); // open for writing

        while (archive_read_data_block(ar, &buff, &size, &offset) == ARCHIVE_OK) {
            if (offset > output_offset) {
                f.seek(offset - output_offset, SEEK_CUR);
                output_offset = offset;
            }
            while (size > 0) {
                auto bytes_to_write = size;
                if (bytes_to_write > DEFAULT_BLOCK_SIZE)
                    bytes_to_write = DEFAULT_BLOCK_SIZE;

                f.rawWrite(buff[0 .. bytes_to_write]);
                output_offset += bytes_to_write;
                size -= bytes_to_write;
            }
        }
    }

    archive* openArchive()
    {
        archive* ar = archive_read_new();

        archive_read_support_filter_all(ar);
        archive_read_support_format_all(ar);

        auto ret = archive_read_open_filename(ar, archive_fname.toStringz, DEFAULT_BLOCK_SIZE);
        if (ret != ARCHIVE_OK) {
            auto ret_errno = archive_errno(ar);
            auto ret_strerr = fromStringz(strerror(ret_errno));
            throw new Exception(format("Unable to open compressed file '%s': %s. error: %s (%d)",
                    archive_fname,
                    getArchiveErrorMessage(ar),
                    ret_strerr,
                    ret_errno));
        }

        return ar;
    }

    bool pathMatches (string path1, string path2)
    {
        import std.path;

        if (path1 == path2)
            return true;

        immutable path1Abs = buildNormalizedPath("/", path1);
        immutable path2Abs = buildNormalizedPath("/", path2);

        if (path1Abs == path2Abs)
            return true;

        return false;
    }

public:

    struct ArchiveEntry {
        string fname;
        const(ubyte)[] data;
    }

    void open (string fname)
    {
        archive_fname = fname;
    }

    bool isOpen ()
    {
        return archive_fname !is null;
    }

    void close ()
    {
        archive_fname = null;
    }

    bool extractFileTo (string fname, string fdest)
    {
        archive_entry* en;

        auto ar = openArchive();
        scope (exit)
            archive_read_free (ar);

        while (archive_read_next_header(ar, &en) == ARCHIVE_OK) {
            auto pathname = fromStringz(archive_entry_pathname(en));

            if (pathMatches(fname, to!string(pathname))) {
                this.extractEntryTo(ar, fdest);
                return true;
            } else {
                archive_read_data_skip(ar);
            }
        }

        return false;
    }

    void extractArchive (const string dest)
    in {
        assert(std.file.isDir(dest));
    }
    do {
        import std.path;

        archive_entry* en;

        auto ar = openArchive();
        scope (exit)
            archive_read_free (ar);

        while (archive_read_next_header(ar, &en) == ARCHIVE_OK) {
            auto pathname = buildPath(dest, archive_entry_pathname(en).fromStringz);
            /* at the moment we only handle directories and files */
            if (archive_entry_filetype(en) == AE_IFDIR) {
                if (!std.file.exists(pathname))
                    std.file.mkdir(pathname);
                continue;
            }

            if (archive_entry_filetype(en) == AE_IFREG) {
                this.extractEntryTo(ar, pathname);
            }
        }
    }

    const(ubyte)[] readData (string fname)
    {
        import core.sys.posix.sys.stat;
        import std.path;

        archive_entry* en;

        auto ar = openArchive();
        scope (exit)
            archive_read_free (ar);

        while (archive_read_next_header(ar, &en) == ARCHIVE_OK) {
            immutable pathname = to!string(fromStringz(archive_entry_pathname(en)));

            if (pathMatches(fname, pathname)) {
                immutable filetype = archive_entry_filetype(en);

                if (filetype == S_IFDIR) {
                    /* we don't extract directories explicitly */
                    throw new Exception(format("Path '%s' is a directory and can not be extracted.", fname));
                }

                /* check if we are dealing with a symlink */
                if (filetype == S_IFLNK) {
                    immutable fnameAbs = absolutePath(fname, "/");
                    string linkTarget = to!string(fromStringz(archive_entry_symlink(en)));
                    if (linkTarget is null)
                        throw new Exception(format("Unable to read destination of symbolic link for '%s'.", fname));

                    try {
                        if (!isAbsolute(linkTarget))
                            linkTarget = absolutePath(linkTarget, dirName(fnameAbs));

                        return readData(buildNormalizedPath(linkTarget));
                    } catch (Exception e) {
                        logError("Unable to read destination data of symlink in archive: %s", e.msg);
                        return null;
                    }
                }

                if (filetype != S_IFREG) {
                    // we really don't want to extract special files from a tarball - usually, those shouldn't
                    // be present anyway.
                    // This should probably be an error, but return nothing for now.
                    logError("Tried to extract non-regular file '%s' from the archive", fname);
                    return null;
                }

                return this.readEntry(ar);
            } else {
                archive_read_data_skip(ar);
            }
        }

        throw new Exception(format("File '%s' was not found in the archive.", fname));
    }

    string[] extractFilesByRegex (Regex!char re, string destdir)
    {
        import std.path;

        archive_entry* en;
        auto matches = appender!(string[]);

        auto ar = openArchive();
        scope (exit)
            archive_read_free (ar);

        while (archive_read_next_header(ar, &en) == ARCHIVE_OK) {
            auto pathname = fromStringz(archive_entry_pathname(en));

            const m = matchFirst(pathname, re);
            if (!m.empty) {
                auto fdest = buildPath(destdir, baseName(pathname));
                this.extractEntryTo(ar, fdest);
                matches ~= fdest;
            } else {
                archive_read_data_skip(ar);
            }
        }

        return matches.data;
    }

    string[] readContents ()
    {
        import std.conv : to;

        archive_entry* en;

        auto ar = openArchive();
        scope (exit)
            archive_read_free (ar);

        auto contents = appender!(string[]);
        contents.reserve(20);
        while (archive_read_next_header(ar, &en) == ARCHIVE_OK) {
            auto pathname = fromStringz(archive_entry_pathname(en));

            // ignore directories
            if (pathname.endsWith("/"))
                continue;

            auto path = buildNormalizedPath("/", to!string(pathname));
            contents ~= path;
        }

        return contents.data;
    }

    /**
     * Returns a generator to iterate over the contents of this tarball.
     */
    auto read()
    {
        import core.sys.posix.sys.stat;
        import std.path;

        // we set a high fiber stack size, because having a lower size
        // may make functions crash that need a bigger stack (and libarchive
        // might do that). See https://github.com/ximion/appstream-generator/issues/101
        const size_t fiberStackSize = 65536;

        auto gen = new Generator!ArchiveEntry(
        {
            auto ar = openArchive();
            scope (exit)
                archive_read_free (ar);

            archive_entry* en;
            while (archive_read_next_header(ar, &en) == ARCHIVE_OK) {
                auto pathname = fromStringz(archive_entry_pathname(en));

                // ignore directories
                if (pathname.endsWith("/"))
                    continue;

                auto path = std.path.buildNormalizedPath("/", to!string(pathname));

                ArchiveEntry aentry;
                aentry.fname = path;
                aentry.data = null;

                auto filetype = archive_entry_filetype(en);
                /* check if we are dealing with a symlink */
                if (filetype == S_IFLNK) {
                    auto linkTarget = to!string(fromStringz(archive_entry_symlink(en)));
                    if (linkTarget is null)
                        throw new Exception(format("Unable to read destination of symbolic link for '%s'.", path));

                    // we cheat here and set the link target as data
                    // TODO: Proper handling of symlinks, e.g. by adding a filetype property to ArchiveEntry.
                    aentry.data = cast(ubyte[]) linkTarget;
                    yield(aentry);
                    continue;
                }

                if (filetype != S_IFREG) {
                    yield(aentry);
                    continue;
                }

                aentry.data = this.readEntry(ar);
                yield(aentry);
            }
        }, fiberStackSize);

        return gen;
    }
}

/**
 * Save data to a compressed file.
 *
 * Params:
 *      data = The data to save.
 *      fname = The filename the data should be saved to.
 *      atype = The archive type (GZ or XZ).
 */
void compressAndSave (ref ubyte[] data, const string fname, ArchiveType atype)
{
    import std.file : isFile, remove, rename;

    auto ar = archive_write_new();
    scope (exit)
        archive_write_free (ar);

    archive_write_set_format_raw(ar);
    if (atype == ArchiveType.GZIP) {
        archive_write_add_filter_gzip(ar);
        archive_write_set_filter_option(ar, "gzip", "timestamp", null);
    } else {
        archive_write_add_filter_xz(ar);
    }

    // don't write to the new file directly, we create a temporary file and
    // rename it when we successfully saved the data.
    auto tmpFname = fname ~ ".new";

    auto ret = archive_write_open_filename(ar, tmpFname.toStringz);
    if (ret != ARCHIVE_OK)
        throw new Exception(format("Unable to open file '%s': %s", tmpFname, getArchiveErrorMessage(ar)));

    archive_entry* entry;
    entry = archive_entry_new();
    scope (exit)
        archive_entry_free (entry);

    archive_entry_set_filetype(entry, AE_IFREG);
    archive_entry_set_size(entry, ubyte.sizeof * data.length);
    archive_write_header(ar, entry);

    archive_write_data(ar, data.ptr, ubyte.sizeof * data.length);
    archive_write_close(ar);

    // delete old file if it exists
    if (std.file.exists(fname))
        fname.remove();
    // rename temporary file to actual file
    std.file.rename(tmpFname, fname); // we need to use std.file explicitly, because otherwise core.stdc gets used
}

final class ArchiveCompressor {

private:
    string archiveFname;
    archive* ar;
    bool closed;

public:

    this (ArchiveType type)
    {
        ar = archive_write_new();

        if (type == ArchiveType.GZIP) {
            archive_write_add_filter_gzip(ar);
            archive_write_set_filter_option(ar, "gzip", "timestamp", null);
        } else {
            archive_write_add_filter_xz(ar);
        }

        archive_write_set_format_pax_restricted(ar);
        closed = true;
    }

    ~this ()
    {
        close();
        archive_write_free(ar);
    }

    void open (string fname)
    {
        archiveFname = fname;
        auto ret = archive_write_open_filename(ar, fname.toStringz);
        if (ret != ARCHIVE_OK)
            throw new Exception(format("Unable to open file '%s'", fname, getArchiveErrorMessage(ar)));
        closed = false;
    }

    bool isOpen ()
    {
        return !closed;
    }

    void close ()
    {
        if (closed)
            return;
        archive_write_close(ar);
        closed = true;
    }

    void addFile (const string fname, const string dest = null)
    in {
        if (!std.file.exists(fname)) {
            logError("File %s does not exist!", fname);
            assert(0);
        }
    }
    do {
        import std.conv : octal;
        import std.path : baseName;
        import core.sys.posix.sys.stat;

        immutable BUFFER_SIZE = 8192;
        archive_entry* entry;
        stat_t st;
        ubyte[BUFFER_SIZE] buff;

        string destName = dest;
        if (destName is null)
            destName = baseName(fname);

        lstat(fname.toStringz, &st);
        entry = archive_entry_new();
        scope (exit)
            archive_entry_free (entry);

        archive_entry_set_pathname(entry, destName.toStringz);
        archive_entry_set_size(entry, st.st_size);
        archive_entry_set_filetype(entry, S_IFREG);
        archive_entry_set_perm(entry, octal!755);
        archive_entry_set_mtime(entry, st.st_mtime, 0);

        synchronized {
            archive_write_header(ar, entry);

            auto f = File(fname, "r");
            while (!f.eof) {
                auto data = f.rawRead(buff);
                archive_write_data(ar, data.ptr, ubyte.sizeof * data.length);
            }
        }
    }

}

version (unittest) {
    version (GNU) {
        extern (C) char* mkdtemp(char*) nothrow @nogc;
    } else {
        import core.sys.posix.stdlib : mkdtemp;
    }
}

unittest {
    writeln("TEST: ", "Compressed empty file");

    ubyte[] emptyGz = [
        0x1f, 0x8b, 0x08, 0x08, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x03, 0x65, 0x6d, 0x70, 0x74, 0x79, 0x00,
        0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00,
    ];
    assert(decompressData(emptyGz) == "");

    writeln("TEST: ", "Extracting a tarball");

    import std.file : tempDir;
    import std.path : buildPath;
    import asgen.utils : getTestSamplesDir;

    auto archive = buildPath(getTestSamplesDir(), "test.tar.xz");
    assert(std.file.exists(archive));
    ArchiveDecompressor ar;

    auto tmpdir = buildPath(tempDir, "asgenXXXXXX");
    auto ctmpdir = new char[](tmpdir.length + 1);
    ctmpdir[0 .. tmpdir.length] = tmpdir[];
    ctmpdir[$ - 1] = '\0';

    tmpdir = to!string(mkdtemp(ctmpdir.ptr));
    scope (exit)
        std.file.rmdirRecurse(tmpdir);

    ar.open(archive);
    ar.extractArchive(tmpdir);

    auto path = buildPath(tmpdir, "b", "a");
    assert(std.file.exists(path));
    assert(std.file.readText(path).chomp == "hello");
}
