import os
import zstandard as zstd

def unpack_zst(input_file, output_file=None):
    """
    Unpacks a .zst file using the Zstandard library.

    Args:
        input_file (str): Path to the input .zst file.
        output_file (str): Path to the output file. If None, the output file 
                           will have the same name as the input file but without the .zst extension.

    Returns:
        str: Path to the decompressed file.
    """
    if not input_file.endswith('.zst'):
        raise ValueError("Input file must have a .zst extension")

    if output_file is None:
        output_file = input_file.rsplit('.zst', 1)[0]

    try:
        with open(input_file, 'rb') as compressed_file:
            dctx = zstd.ZstdDecompressor()
            with open(output_file, 'wb') as decompressed_file:
                dctx.copy_stream(compressed_file, decompressed_file)
        print(f"Decompressed '{input_file}' to '{output_file}'")
        return output_file
    except Exception as e:
        print(f"Failed to decompress '{input_file}': {e}")
        raise

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Unpack a .zst file")
    parser.add_argument("input_file", help="Path to the .zst file")
    parser.add_argument(
        "-o", "--output_file", 
        help="Path to the output file (default: same name as input without .zst)"
    )
    args = parser.parse_args()

    try:
        unpack_zst(args.input_file, args.output_file)
    except Exception as e:
        print(f"Error: {e}")