#!/usr/bin/env python3
"""SAND .wbt save-file toolkit.

Envelope (verified from GameAssembly.dll, XorCryptography.Encrypt @ RVA 0xBE9B40):
  Save:  BSON(Newtonsoft) -> XOR encrypt -> gzip compress
  Load:  gunzip -> XOR decrypt -> BSON parse
Inner cipher is XOR with a 6-byte key, applied PER 0xA000-byte chunk with the
key index RESET to 0 at every chunk boundary (the stream is processed in
buffer.Length = 0xA000 reads, and the keystream restarts each read):
  decoded[i] = raw[i] XOR KEY[(i % 0xA000) % 6]
KEY = 70 DD 1F 2A 0B 4A  (recovered from known BSON plaintext "textureSize"
header; the .cctor allocates a 6-element byte[] and InitializeArray-copies it).

Decoded payload is a single Newtonsoft BSON document
(WalkerBlueprintContainerSerializableProxy): textureSize, textureRawData,
walker{...}, format, iconVersion, firstNameIndex, secondNameIndex,
creationTime, name, isBackup. Name parts are 0-based indices into the Unity
Localization tables WalkerFirstName / WalkerSecondName.

Subcommands:
  decode  <wbt>                       -> write decoded bytes
  snap    [--all | <wbt>...]          -> decode + save numbered snapshot per file
  diff    <before.dec> <after.dec>    -> human-readable diff
  check   <wbt>                       -> one-shot: snap + diff against latest prior snap
  watch   <wbt>                       -> poll the file forever; run check on every mtime change
"""
import sys, gzip, argparse, pathlib, re, time

KEY = bytes.fromhex('70dd1f2a0b4a')
CHUNK = 0xA000

WBT_DIR = pathlib.Path('/home/downloadpizza/sand_tools/Walkers')   # symlink -> live game saves
SNAP_DIR = pathlib.Path.home() / 'sand_tools' / 'snapshots'

# Known fields to suppress in diffs (per-file). Each entry: (offset, length, note)
KNOWN_NOISE = {
    # file_uuid_prefix : list of (offset, length, label)
    '*': [
        (0x000000, 2, 'leading 2-byte per-save field'),
    ],
    '215949d3': [
        (0x1040f3, 4, 'per-save nonce/timestamp'),
    ],
}

def decode_bytes(raw: bytes) -> bytes:
    out = bytearray(len(raw))
    for i in range(len(raw)):
        out[i] = raw[i] ^ KEY[(i % CHUNK) % 6]
    return bytes(out)

def encode_bytes(dec: bytes) -> bytes:
    # XOR is symmetric; same transform re-encrypts a (modified) decoded payload.
    return decode_bytes(dec)

def decode_file(path: pathlib.Path) -> bytes:
    return decode_bytes(gzip.decompress(path.read_bytes()))

def stem8(path: pathlib.Path) -> str:
    return path.stem[:8]

def next_snap_path(file_stem: str) -> pathlib.Path:
    SNAP_DIR.mkdir(parents=True, exist_ok=True)
    existing = sorted(SNAP_DIR.glob(f'{file_stem}_v*.dec'))
    n = 1
    if existing:
        nums = [int(re.search(r'_v(\d+)\.dec$', p.name).group(1)) for p in existing]
        n = max(nums) + 1
    return SNAP_DIR / f'{file_stem}_v{n:03d}.dec'

def latest_snap(file_stem: str) -> pathlib.Path | None:
    snaps = sorted(SNAP_DIR.glob(f'{file_stem}_v*.dec'))
    return snaps[-1] if snaps else None

def cmd_decode(args):
    data = decode_file(pathlib.Path(args.input))
    out = pathlib.Path(args.output) if args.output else None
    if out:
        out.write_bytes(data); print(f"wrote {out} ({len(data)} bytes)", file=sys.stderr)
    else:
        sys.stdout.buffer.write(data)

def cmd_snap(args):
    if args.all:
        files = sorted(WBT_DIR.glob('*.wbt'))
    else:
        files = [pathlib.Path(p) for p in args.files]
    for f in files:
        data = decode_file(f)
        out = next_snap_path(stem8(f))
        out.write_bytes(data)
        print(f"snap: {f.name[:8]}... -> {out.name} ({len(data)} bytes)")

def diff_regions(a: bytes, b: bytes, merge_gap=8):
    common = min(len(a), len(b))
    regions = []
    state = None; start = 0
    for i in range(common):
        eq = a[i] == b[i]
        cur = 'eq' if eq else 'df'
        if state is None: state = cur; start = i
        elif cur != state: regions.append((state, start, i)); state = cur; start = i
    regions.append((state, start, common))
    merged = []
    for r in regions:
        if merged and r[0]=='eq' and (r[2]-r[1])<merge_gap and merged[-1][0]=='df':
            merged[-1] = ('df', merged[-1][1], r[2])
        elif merged and merged[-1][0]==r[0]:
            merged[-1] = (r[0], merged[-1][1], r[2])
        else:
            merged.append(list(r))
    return [tuple(r) for r in merged if r[0]=='df'], len(a), len(b)

def filter_noise(diffs, file_stem):
    noise = KNOWN_NOISE.get('*', []) + KNOWN_NOISE.get(file_stem, [])
    def overlaps(s, e):
        return any(s < no+nn and e > no for no, nn, _ in noise)
    kept = [d for d in diffs if not overlaps(d[1], d[2])]
    suppressed = len(diffs) - len(kept)
    return kept, suppressed

def render_diffs(a, b, diffs, ctx=12):
    for t, s, e in diffs:
        n = e - s
        pre = a[max(0,s-ctx):s].hex()
        post = a[e:e+ctx].hex()
        print(f"@ {s:#010x} .. {e:#010x}  ({n} bytes)")
        print(f"  -before: {pre}  [{a[s:e].hex()}]  {post}")
        print(f"  +after : {pre}  [{b[s:e].hex()}]  {post}")
        if n <= 4:
            bv = int.from_bytes(a[s:e], 'little')
            av = int.from_bytes(b[s:e], 'little')
            print(f"  LE int : {bv} -> {av}  (Δ {av-bv:+d})")
        print()
    sys.stdout.flush()

def cmd_diff(args):
    a = pathlib.Path(args.before).read_bytes()
    b = pathlib.Path(args.after).read_bytes()
    diffs, la, lb = diff_regions(a, b)
    file_stem = re.match(r'([0-9a-f]{8})', pathlib.Path(args.before).name)
    stem = file_stem.group(1) if file_stem else ''
    if not args.no_filter:
        diffs, suppressed = filter_noise(diffs, stem)
    else:
        suppressed = 0
    print(f"sizes: before={la} after={lb} Δ={lb-la:+d}")
    print(f"{len(diffs)} differing region(s)" + (f"  ({suppressed} known-noise suppressed)" if suppressed else ""))
    print()
    render_diffs(a, b, diffs)

def cmd_check(args):
    path = pathlib.Path(args.input)
    stem = stem8(path)
    prev = latest_snap(stem)
    data = decode_file(path)
    out = next_snap_path(stem)
    out.write_bytes(data)
    print(f"snap: {path.name[:8]}... -> {out.name}")
    if not prev:
        print("(no prior snapshot to diff against)")
        return
    a = prev.read_bytes()
    b = data
    diffs, la, lb = diff_regions(a, b)
    if not args.no_filter:
        diffs, suppressed = filter_noise(diffs, stem)
    else:
        suppressed = 0
    print(f"diff vs {prev.name}: sizes {la}/{lb} Δ={lb-la:+d}")
    print(f"{len(diffs)} region(s)" + (f"  ({suppressed} known-noise suppressed)" if suppressed else ""))
    print()
    render_diffs(a, b, diffs)

def cmd_watch(args):
    path = pathlib.Path(args.input)
    print(f"watching {path.name} (poll every {args.interval}s, Ctrl-C to stop)", flush=True)
    last_mtime = path.stat().st_mtime
    while True:
        try:
            time.sleep(args.interval)
        except KeyboardInterrupt:
            print("stopped."); return
        try:
            cur = path.stat().st_mtime
        except FileNotFoundError:
            continue
        if cur != last_mtime:
            last_mtime = cur
            print(f"\n[{time.strftime('%H:%M:%S')}] change detected", flush=True)
            cmd_check(args)

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sp = ap.add_subparsers(dest='cmd', required=True)

    p = sp.add_parser('decode'); p.add_argument('input'); p.add_argument('-o', '--output'); p.set_defaults(fn=cmd_decode)
    p = sp.add_parser('snap'); g = p.add_mutually_exclusive_group(required=True)
    g.add_argument('--all', action='store_true'); g.add_argument('files', nargs='*', default=[])
    p.set_defaults(fn=cmd_snap)
    p = sp.add_parser('diff'); p.add_argument('before'); p.add_argument('after'); p.add_argument('--no-filter', action='store_true'); p.set_defaults(fn=cmd_diff)
    p = sp.add_parser('check'); p.add_argument('input'); p.add_argument('--no-filter', action='store_true'); p.set_defaults(fn=cmd_check)
    p = sp.add_parser('watch'); p.add_argument('input'); p.add_argument('--no-filter', action='store_true'); p.add_argument('--interval', type=float, default=1.0); p.set_defaults(fn=cmd_watch)
    args = ap.parse_args()
    args.fn(args)

if __name__ == '__main__':
    main()