#!/usr/bin/env qore
# -*- mode: qore; indent-tabs-mode: nil -*-

%new-style
%enable-all-warnings
%require-types
%strict-args

%requires ../../../../qlib/QUnit.qm

%exec-class CompressionStreamTest

class SrcStream inherits InputStream {
    public {
        binary data;
        int offset = 0;
        int chunk = 1;
    }

    constructor(binary d, int c = 1) {
        data = d;
        chunk = c;
    }

    *binary read(int limit) {
        if (limit > chunk) {
            limit = chunk;
        }
        if (limit > length(data) - offset) {
            limit = length(data) - offset;
        }
        if (limit == 0) {
            return NOTHING;
        }
        binary b = data.substr(offset, limit);
        offset += limit;
        return b;
    }

    int peek() {
        *binary b = data.substr(offset, 1);
        return ord(b.toString(b, "UTF-8"), 0);
    }
}

public class CompressionStreamTest inherits QUnit::Test {

    private {
        binary plain = File::readBinaryFile(get_script_dir() + "/../../data/lorem");
        binary zlib = File::readBinaryFile(get_script_dir() + "/../../data/lorem.zlib");
        binary gzip = File::readBinaryFile(get_script_dir() + "/../../data/lorem.gz");
        binary bzip2 = File::readBinaryFile(get_script_dir() + "/../../data/lorem.bz2");
    }

    constructor() : Test("CompressionStreamTest", "1.0") {
        addTestCase("zlib compression input stream", \zlibCompressInput());
        addTestCase("zlib compression output stream", \zlibCompressOutput());
        addTestCase("gzip compression input stream", \gzipCompressInput());
        addTestCase("gzip compression output stream", \gzipCompressOutput());
        addTestCase("bzip2 compression input stream", \bzip2CompressInput());
        addTestCase("bzip2 compression output stream", \bzip2CompressOutput());
        addTestCase("zlib decompression input stream", \zlibDecompressInput());
        addTestCase("zlib decompression output stream", \zlibDecompressOutput());
        addTestCase("gzip decompression input stream", \gzipDecompressInput());
        addTestCase("gzip decompression output stream", \gzipDecompressOutput());
        addTestCase("bzip2 decompression input stream", \bzip2DecompressInput());
        addTestCase("bzip2 decompression output stream", \bzip2DecompressOutput());
        addTestCase("decompression algorithm check", \decompressAlgCheck());

        # Return for compatibility with test harness that checks return value.
        set_return_value(main());
    }

    zlibCompressInput() {
        assertEq(zlib, compressInput(plain, COMPRESSION_ALG_ZLIB, 1, 100000));
        assertEq(zlib, compressInput(plain, COMPRESSION_ALG_ZLIB, 100000, 1));
        assertEq(zlib, compressInput(plain, COMPRESSION_ALG_ZLIB, 100000, 100000));
    }

    zlibCompressOutput() {
        assertEq(zlib, compressOutput(plain, COMPRESSION_ALG_ZLIB, 1));
        assertEq(zlib, compressOutput(plain, COMPRESSION_ALG_ZLIB, 100000));
    }

    gzipCompressInput() {
        assertEq(gzip, ignoreOs(compressInput(plain, COMPRESSION_ALG_GZIP, 1, 100000), gzip));
        assertEq(gzip, ignoreOs(compressInput(plain, COMPRESSION_ALG_GZIP, 100000, 1), gzip));
        assertEq(gzip, ignoreOs(compressInput(plain, COMPRESSION_ALG_GZIP, 100000, 100000), gzip));
    }

    gzipCompressOutput() {
        assertEq(gzip, ignoreOs(compressOutput(plain, COMPRESSION_ALG_GZIP, 1), gzip));
        assertEq(gzip, ignoreOs(compressOutput(plain, COMPRESSION_ALG_GZIP, 100000), gzip));
    }

    bzip2CompressInput() {
        assertEq(bzip2, compressInput(plain, COMPRESSION_ALG_BZIP2, 1, 100000));
        assertEq(bzip2, compressInput(plain, COMPRESSION_ALG_BZIP2, 100000, 1));
        assertEq(bzip2, compressInput(plain, COMPRESSION_ALG_BZIP2, 100000, 100000));
    }

    bzip2CompressOutput() {
        assertEq(bzip2, compressOutput(plain, COMPRESSION_ALG_BZIP2, 1));
        assertEq(bzip2, compressOutput(plain, COMPRESSION_ALG_BZIP2, 100000));
    }

    zlibDecompressInput() {
        assertEq(plain, decompressInput(zlib, COMPRESSION_ALG_ZLIB, 1, 100000));
        assertEq(plain, decompressInput(zlib, COMPRESSION_ALG_ZLIB, 100000, 1));
        assertEq(plain, decompressInput(zlib, COMPRESSION_ALG_ZLIB, 100000, 100000));
        assertThrows("ZLIB-ERROR", "extra bytes",
                sub() { decompressInput(zlib + <00>, COMPRESSION_ALG_ZLIB, 100000, 100000); });
        assertThrows("ZLIB-ERROR", "Unexpected end",
                sub() { decompressInput(zlib.substr(0, zlib.size() - 1), COMPRESSION_ALG_ZLIB, 100000, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressInput(modify(zlib), COMPRESSION_ALG_ZLIB, 100000, 100000); });
    }

    zlibDecompressOutput() {
        assertEq(plain, decompressOutput(zlib, COMPRESSION_ALG_ZLIB, 1));
        assertEq(plain, decompressOutput(zlib, COMPRESSION_ALG_ZLIB, 100000));
        assertThrows("ZLIB-ERROR", "extra bytes",
                sub() { decompressOutput(zlib + <00>, COMPRESSION_ALG_ZLIB, 100000); });
        assertThrows("ZLIB-ERROR", "Unexpected end",
                sub() { decompressOutput(zlib.substr(0, zlib.size() - 1), COMPRESSION_ALG_ZLIB, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressOutput(modify(zlib), COMPRESSION_ALG_ZLIB, 100000); });
    }

    gzipDecompressInput() {
        assertEq(plain, decompressInput(gzip, COMPRESSION_ALG_GZIP, 1, 100000));
        assertEq(plain, decompressInput(gzip, COMPRESSION_ALG_GZIP, 100000, 1));
        assertEq(plain, decompressInput(gzip, COMPRESSION_ALG_GZIP, 100000, 100000));
        assertThrows("ZLIB-ERROR", "extra bytes",
                sub() { decompressInput(gzip + <00>, COMPRESSION_ALG_GZIP, 100000, 100000); });
        assertThrows("ZLIB-ERROR", "Unexpected end",
                sub() { decompressInput(gzip.substr(0, gzip.size() - 1), COMPRESSION_ALG_GZIP, 100000, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressInput(modify(gzip), COMPRESSION_ALG_GZIP, 100000, 100000); });
    }

    gzipDecompressOutput() {
        assertEq(plain, decompressOutput(gzip, COMPRESSION_ALG_GZIP, 1));
        assertEq(plain, decompressOutput(gzip, COMPRESSION_ALG_GZIP, 100000));
        assertThrows("ZLIB-ERROR", "extra bytes",
                sub() { decompressOutput(gzip + <00>, COMPRESSION_ALG_GZIP, 100000); });
        assertThrows("ZLIB-ERROR", "Unexpected end",
                sub() { decompressOutput(gzip.substr(0, gzip.size() - 1), COMPRESSION_ALG_GZIP, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressOutput(modify(gzip), COMPRESSION_ALG_GZIP, 100000); });
    }

    bzip2DecompressInput() {
        assertEq(plain, decompressInput(bzip2, COMPRESSION_ALG_BZIP2, 1, 100000));
        assertEq(plain, decompressInput(bzip2, COMPRESSION_ALG_BZIP2, 100000, 1));
        assertEq(plain, decompressInput(bzip2, COMPRESSION_ALG_BZIP2, 100000, 100000));
        assertThrows("BZIP2-ERROR", "extra bytes",
                sub() { decompressInput(bzip2 + <00>, COMPRESSION_ALG_BZIP2, 100000, 100000); });
        assertThrows("BZIP2-ERROR", "Unexpected end",
                sub() { decompressInput(bzip2.substr(0, bzip2.size() - 1), COMPRESSION_ALG_BZIP2, 100000, 100000); });
        assertThrows("BZIP2-ERROR", "corrupted",
                sub() { decompressInput(modify(bzip2), COMPRESSION_ALG_BZIP2, 100000, 100000); });
    }

    bzip2DecompressOutput() {
        assertEq(plain, decompressOutput(bzip2, COMPRESSION_ALG_BZIP2, 1));
        assertEq(plain, decompressOutput(bzip2, COMPRESSION_ALG_BZIP2, 100000));
        assertThrows("BZIP2-ERROR", "extra bytes",
                sub() { decompressOutput(bzip2 + <00>, COMPRESSION_ALG_BZIP2, 100000); });
        assertThrows("BZIP2-ERROR", "Unexpected end",
                sub() { decompressOutput(bzip2.substr(0, bzip2.size() - 1), COMPRESSION_ALG_BZIP2, 100000); });
        assertThrows("BZIP2-ERROR", "corrupted",
                sub() { decompressOutput(modify(bzip2), COMPRESSION_ALG_BZIP2, 100000); });
    }

    decompressAlgCheck() {
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressInput(gzip, COMPRESSION_ALG_ZLIB, 100000, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressInput(bzip2, COMPRESSION_ALG_ZLIB, 100000, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressOutput(gzip, COMPRESSION_ALG_ZLIB, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressOutput(bzip2, COMPRESSION_ALG_ZLIB, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressInput(zlib, COMPRESSION_ALG_GZIP, 100000, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressInput(bzip2, COMPRESSION_ALG_GZIP, 100000, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressOutput(zlib, COMPRESSION_ALG_GZIP, 100000); });
        assertThrows("ZLIB-ERROR", "data corrupted",
                sub() { decompressOutput(bzip2, COMPRESSION_ALG_GZIP, 100000); });
        assertThrows("BZIP2-ERROR", "corrupted",
                sub() { decompressInput(zlib, COMPRESSION_ALG_BZIP2, 100000, 100000); });
        assertThrows("BZIP2-ERROR", "corrupted",
                sub() { decompressInput(gzip, COMPRESSION_ALG_BZIP2, 100000, 100000); });
        assertThrows("BZIP2-ERROR", "corrupted",
                sub() { decompressOutput(zlib, COMPRESSION_ALG_BZIP2, 100000); });
        assertThrows("BZIP2-ERROR", "corrupted",
                sub() { decompressOutput(gzip, COMPRESSION_ALG_BZIP2, 100000); });
    }

    /*
        issue 1565: the gzip compressed data sometimes differs in the OS byte in the header

        reference: https://tools.ietf.org/html/rfc1952#page-5

        the header looks like:
        +---+---+---+---+---+---+---+---+---+---+
        |ID1|ID2|CM |FLG|     MTIME     |XFL|OS | (more-->)
        +---+---+---+---+---+---+---+---+---+---+

        OS (Operating System)
            This identifies the type of file system on which compression
            took place.  This may be useful in determining end-of-line
            convention for text files.  The currently defined values are
            as follows:
                 0 - FAT filesystem (MS-DOS, OS/2, NT/Win32)
                 1 - Amiga
                 2 - VMS (or OpenVMS)
                 3 - Unix
                 4 - VM/CMS
                 5 - Atari TOS
                 6 - HPFS filesystem (OS/2, NT)
                 7 - Macintosh
                 8 - Z-System
                 9 - CP/M
                10 - TOPS-20
                11 - NTFS filesystem (NT)
                12 - QDOS
                13 - Acorn RISCOS
               255 - unknown
    */
    private binary ignoreOs(binary b, binary expected) {
        splice b, 9, 1, chr(expected[9]);
        return b;
    }

    private binary compressInput(binary src, string alg, int chunk, int readSize) {
        return processInput(src, get_compressor(alg), chunk, readSize);
    }

    private binary compressOutput(binary src, string alg, int writeSize) {
        return processOutput(src, get_compressor(alg), writeSize);
    }

    private binary decompressInput(binary src, string alg, int chunk, int readSize) {
        return processInput(src, get_decompressor(alg), chunk, readSize);
    }

    private binary decompressOutput(binary src, string alg, int writeSize) {
        return processOutput(src, get_decompressor(alg), writeSize);
    }

    private binary processInput(binary src, Transform t, int chunk, int readSize) {
        TransformInputStream tis(new SrcStream(src, chunk), t);
        binary out = binary();
        while (True) {
            *binary b = tis.read(readSize);
            if (!b) {
                break;
            }
            out = out + b;
        }
        return out;
    }

    private binary processOutput(binary src, Transform t, int writeSize) {
        BinaryOutputStream bos();
        TransformOutputStream tos(bos, t);
        int o = 0;
        while (o < src.size()) {
            int w = src.size() - o;
            if (w > writeSize) {
                w = writeSize;
            }
            tos.write(src.substr(o, w));
            o += w;
        }
        tos.close();
        return bos.getData();
    }

    private binary modify(binary b) {
        int i = (b[b.size() - 1] + 1) & 255;
        return b.substr(b.size() - 1) + parse_hex_string(sprintf("%02X", i));
    }
}
