#!/usr/bin/env python3
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

"""This script compresses and decompresses data using the zstandard compression
format, as provided by the python-zstandard module.

Data is provided on stdin and output on stdout."""

import sys
import zstandard
from argparse import ArgumentParser


def main(argv=None):
    parser = ArgumentParser(description=__doc__)
    parser.set_defaults(mode="compress")
    parser.add_argument(
        "-z",
        "--compress",
        dest="mode",
        action="store_const",
        const="compress",
        help="compress the data (this is the default)",
    )
    parser.add_argument(
        "-d",
        "--decompress",
        dest="mode",
        action="store_const",
        const="decompress",
        help="decompress the data",
    )
    parser.add_argument(
        "-T",
        "--threads",
        dest="threads",
        default=1,
        type=int,
        help="Compress using # working threads. If 0, use number of CPUs on the system. (default 1)",
    )
    parser.add_argument(
        "-l",
        "--level",
        dest="level",
        default=3,
        type=int,
        help="Compression level from 1-22 (default 3)",
    )
    parser.add_argument(
        "file",
        nargs="?",
        help="File to compress/decompress. Default is stdin.",
    )

    args = parser.parse_args(argv)

    # The zstd commandline tool uses 0 to specify number of threads equal to
    # the number of CPUs whereas the python module uses negative numbers to
    # flag this behavior. Emulate the zstd commandline utility's behavior here
    if args.threads == 0:
        args.threads = -1

    if args.file:
        in_file = open(args.file, "rb")
    else:
        in_file = sys.stdin.buffer

    if args.mode == "compress":
        ctx = zstandard.ZstdCompressor(level=args.level, threads=args.threads)
    elif args.mode == "decompress":
        ctx = zstandard.ZstdDecompressor()

    ctx.copy_stream(in_file, sys.stdout.buffer)


if __name__ == "__main__":
    main()