#!/usr/bin/env python3
"""
Batch-decompress Steam .vz files in the current folder using valvevz.exe.

Handles both plain names:
    steam_win32.zip.vz
and Steam's real cache names (name.vz.<sha1>_<size>):
    steam_win32.zip.vz.5e2a876cfdc1020f0c8ebd9ba10bc1377a6bb3ae_631023

Usage:
    python decompress_vz.py
    python decompress_vz.py --folder "G:\Downloads\stem\package" --exe valvevz.exe
"""

import argparse
import re
import subprocess
import sys
from pathlib import Path

# Matches: <base>.vz.<40-char-hex-sha1>_<digits>   -> base is what we restore
PATTERN_WITH_HASH = re.compile(r"^(?P<base>.+)\.vz\.[0-9a-f]{40}_\d+$", re.IGNORECASE)
# Matches: <base>.vz   (plain, no trailing hash/size)
PATTERN_PLAIN = re.compile(r"^(?P<base>.+)\.vz$", re.IGNORECASE)


def find_output_name(filename: str) -> str | None:
    """Return the restored output filename, or None if this file doesn't match."""
    m = PATTERN_WITH_HASH.match(filename)
    if m:
        return m.group("base")
    m = PATTERN_PLAIN.match(filename)
    if m:
        return m.group("base")
    return None


def main():
    parser = argparse.ArgumentParser(description="Batch-decompress Steam .vz files.")
    parser.add_argument("--folder", default=".", help="Folder to scan (default: current folder)")
    parser.add_argument("--exe", default="valvevz.exe", help="Path to valvevz.exe (default: valvevz.exe, expected next to this script or on PATH)")
    parser.add_argument("--keep-original", action="store_true", help="Don't delete the .vz source file after successful decompression")
    parser.add_argument("--overwrite", action="store_true", help="Overwrite output files if they already exist")
    args = parser.parse_args()

    folder = Path(args.folder).resolve()
    exe_path = Path(args.exe)
    if not exe_path.is_absolute():
        # Try next to this script first, then fall back to bare name (PATH lookup)
        script_dir_candidate = Path(__file__).resolve().parent / args.exe
        if script_dir_candidate.exists():
            exe_path = script_dir_candidate

    if not folder.is_dir():
        print(f"Error: folder not found: {folder}")
        sys.exit(1)

    candidates = []
    for entry in sorted(folder.iterdir()):
        if not entry.is_file():
            continue
        out_name = find_output_name(entry.name)
        if out_name:
            candidates.append((entry, out_name))

    if not candidates:
        print(f"No .vz files found in {folder}")
        return

    print(f"Found {len(candidates)} .vz file(s) in {folder}\n")

    ok_count = 0
    fail_count = 0
    skip_count = 0

    for src_path, out_name in candidates:
        out_path = src_path.parent / out_name

        if out_path.exists() and not args.overwrite:
            print(f"[SKIP] {src_path.name} -> {out_name} (already exists, use --overwrite to replace)")
            skip_count += 1
            continue

        print(f"[....] {src_path.name} -> {out_name}")
        try:
            result = subprocess.run(
                [str(exe_path), "--decompress", str(src_path), str(out_path)],
                capture_output=True,
                text=True,
            )
        except FileNotFoundError:
            print(f"Error: could not run '{exe_path}'. Make sure valvevz.exe is next to this script, "
                  f"in the current folder, or pass --exe <path>.")
            sys.exit(1)

        if result.returncode == 0:
            print(f"[ OK ] {out_name}")
            ok_count += 1
            if not args.keep_original:
                try:
                    src_path.unlink()
                except OSError as e:
                    print(f"        (warning: could not delete source file: {e})")
        else:
            print(f"[FAIL] {src_path.name} (exit code {result.returncode})")
            if result.stdout.strip():
                print(f"        {result.stdout.strip()}")
            if result.stderr.strip():
                print(f"        {result.stderr.strip()}")
            fail_count += 1

    print(f"\nDone: {ok_count} succeeded, {fail_count} failed, {skip_count} skipped.")


if __name__ == "__main__":
    main()