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}")