#!/usr/bin/env python3
"""Harvest per-part hashes from SAND walker saves (100% offline, no injection/proxy).

Every compartment placed in the in-game editor writes its DefinitionHash (and
CompartmentHash) into the .wbt save. This scans all saves, decodes them, and
merges every (EpbId -> {DefinitionHash, CompartmentHash}) into a JSON table,
deduping and keeping a history of any older hash values seen per part.

Re-run it any time after building/saving new walkers to grow the table.

Usage:
  harvest_hashes.py            # scan the live Walkers dir + snapshots, update table
  harvest_hashes.py <dir> ...  # scan extra dirs of .wbt too
"""
import sys, gzip, json, glob, os, hashlib

KEY = bytes.fromhex('70dd1f2a0b4a'); CHUNK = 0xA000
WBT_DIR = '/home/downloadpizza/sand_tools/Walkers'   # symlink -> live game saves (*.wbt)
EXTRACTED = '/home/downloadpizza/sand_tools/extracted'
OUT = os.path.join(EXTRACTED, 'definition_hashes_known.json')
COMPDB = os.path.join(EXTRACTED, 'CompartmentsDatabase.json')

try:
    from bson import decode
except ImportError:
    sys.exit("need pymongo: ~/sand_tools/venv/bin/pip install pymongo")

def decode_wbt(path):
    raw = gzip.decompress(open(path, 'rb').read())
    dec = bytes(raw[i] ^ KEY[(i % CHUNK) % 6] for i in range(len(raw)))
    return decode(dec)

def main():
    dirs = [WBT_DIR, os.path.join(os.path.expanduser('~'), 'sand_tools', 'snapshots')] + sys.argv[1:]
    files = []
    for d in dirs:
        files += glob.glob(os.path.join(d, '*.wbt'))
    # snapshots are decoded already (.dec) -> handle separately
    dec_files = []
    for d in dirs:
        dec_files += glob.glob(os.path.join(d, '*.dec'))

    # rebuild fresh from saves+snapshots (the authoritative source); snapshots persist
    # history even if live .wbt files are later deleted.
    table = {}

    def note(epb, comp, deff):
        e = table.setdefault(epb, {'DefinitionHash': None, 'DefinitionHash_history': [],
                                   'CompartmentHash': None, 'CompartmentHash_history': []})
        for field, hist, val in [('DefinitionHash', 'DefinitionHash_history', deff),
                                 ('CompartmentHash', 'CompartmentHash_history', comp)]:
            if not val:
                continue
            if e[field] and e[field] != val and e[field] not in e[hist]:
                e[hist].append(e[field])
            e[field] = val

    scanned = 0
    def walkers_from(doc):
        w = doc['walker']
        return [w['Chassis']] + w['Compartments']

    for f in sorted(set(files)):
        try:
            doc = decode_wbt(f); scanned += 1
        except Exception as ex:
            print(f"skip {os.path.basename(f)}: {ex}"); continue
        for c in walkers_from(doc):
            note(c['EpbId'], c.get('CompartmentHash'), c.get('DefinitionHash'))
    for f in sorted(set(dec_files)):
        try:
            doc = decode(open(f, 'rb').read()); scanned += 1
        except Exception:
            continue
        try:
            for c in walkers_from(doc):
                note(c['EpbId'], c.get('CompartmentHash'), c.get('DefinitionHash'))
        except Exception:
            pass

    # cross-check CompartmentHash against CompartmentsDatabase (authoritative current value)
    total_parts = None
    if os.path.exists(COMPDB):
        db = json.load(open(COMPDB))
        total_parts = len(db['compartments'])
        calc = {c['entityId']: hashlib.md5(json.dumps(c, separators=(',', ':')).encode()).hexdigest().upper()
                for c in db['compartments']}
        for epb, e in table.items():
            if epb in calc:
                e['CompartmentHash_computed'] = calc[epb]

    out = {
        '_note': ("Per-part hashes harvested from walker saves. CompartmentHash is also computed "
                  "offline from CompartmentsDatabase (CompartmentHash_computed, authoritative). "
                  "DefinitionHash is server-sourced and only known for parts that appear in a save — "
                  "build a part in-game and re-run this to capture it."),
        'definition_known': sum(1 for e in table.values() if e.get('DefinitionHash')),
        'total_parts_in_db': total_parts,
        'parts': dict(sorted(table.items())),
    }
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    json.dump(out, open(OUT, 'w'), indent=2)
    print(f"scanned {scanned} saves; {out['definition_known']} parts with DefinitionHash"
          + (f" / {total_parts} total" if total_parts else "") + f"\nwrote {OUT}")

if __name__ == '__main__':
    main()