#!/usr/bin/env python3
"""Build / edit SAND .wbt walker saves offline.

Pipeline (both directions verified byte-exact on all 5 local walkers):
  load:  gunzip -> XOR -> Newtonsoft-BSON decode
  save:  BSON encode -> XOR -> gzip
pymongo's bson.encode reproduces Newtonsoft.Bson byte-for-byte for these docs,
so a decoded->encoded round-trip is identity.

What "arbitrary" means here:
  * You can freely EDIT a loaded walker: rename, move/rotate/add/remove parts and
    connections, swap texture, etc. pack() recomputes the two OFFLINE hashes
    (CompartmentsHash, ConnectionsHash) so the file stays self-consistent.
  * DefinitionsHash (aggregate) and per-part DefinitionHash are SERVER-SOURCED.
    pack() copies DefinitionsHash from the source doc and reuses each part's
    DefinitionHash. That is correct as long as the *set of part definitions*
    doesn't change. If you introduce a part whose EpbId we've never harvested a
    DefinitionHash for, pack() raises -- we can't invent that value offline.
    Known DefinitionHashes live in extracted/definition_hashes_known.json
    (currently 18/126 parts; build a part in-game once to harvest more).

CLI:
  build_wbt.py repack  <wbt>                 sanity: load->pack->reload == identity
  build_wbt.py rename  <wbt> <first> <second> [-o out.wbt]   set name indices (0-31)
  build_wbt.py pack    <wbt> -o out.wbt      recompute hashes, write a fresh .wbt
  build_wbt.py get-icon <wbt> [-o out.png]   extract the walker icon to a PNG (upright)
  build_wbt.py set-icon <wbt> <png> [-o out] paint a PNG as the icon (auto-resize+flip)
"""
import sys, gzip, json, os, argparse, pathlib
from bson import decode, encode
import walker_hashes as wh

KEY = bytes.fromhex('70dd1f2a0b4a'); CHUNK = 0xA000
KNOWN = pathlib.Path('/home/downloadpizza/sand_tools'
                     '/extracted/definition_hashes_known.json')

def _xor(b: bytes) -> bytes:
    return bytes(b[i] ^ KEY[(i % CHUNK) % 6] for i in range(len(b)))

def load(path) -> dict:
    """gunzip -> XOR -> BSON -> dict (full container, not just doc['walker'])."""
    return decode(_xor(gzip.decompress(pathlib.Path(path).read_bytes())))

def _known_defhashes() -> dict:
    if not KNOWN.exists():
        return {}
    d = json.load(open(KNOWN))
    out = {}
    for epb, e in d['parts'].items():
        if e.get('DefinitionHash'):
            out[epb] = e['DefinitionHash']
    return out

def pack(doc: dict, *, recompute=True, strict=True) -> bytes:
    """Recompute offline hashes on doc['walker'] and return loadable .wbt bytes.

    recompute: refresh CompartmentsHash/ConnectionsHash from the parts/connections.
    strict:    error if any placed part lacks a known DefinitionHash (i.e. its
               server-sourced value isn't recorded), since the file would carry a
               DefinitionHash the server can't validate.
    """
    w = doc['walker']
    if strict:
        known = _known_defhashes()
        missing = []
        for c in [w['Chassis']] + w['Compartments']:
            epb = c['EpbId']
            have = c.get('DefinitionHash')
            if not have and epb not in known:
                missing.append(epb)
        if missing:
            raise ValueError(
                "no known DefinitionHash for EpbId(s): " + ", ".join(sorted(set(missing)))
                + "\n  -> build that part in-game once and re-run harvest_hashes.py, "
                  "or pass strict=False to bypass.")
        # backfill any part missing its per-part DefinitionHash from the known table
        for c in [w['Chassis']] + w['Compartments']:
            if not c.get('DefinitionHash') and c['EpbId'] in known:
                c['DefinitionHash'] = known[c['EpbId']]
    if recompute:
        w['CompartmentsHash'] = wh.compartments_hash(w['Compartments'])
        w['ConnectionsHash'] = wh.connections_hash(w['Connections'])
        # DefinitionsHash (aggregate over CompartmentDefinitionDto) is server-sourced;
        # left as-is -- valid while the definition set is unchanged.
    return gzip.compress(_xor(encode(doc)))

def save(doc: dict, path, **kw):
    pathlib.Path(path).write_bytes(pack(doc, **kw))

# ---- CLI ----
def cmd_repack(a):
    doc = load(a.input)
    blob = pack(doc, strict=False)
    back = load_bytes(blob)
    same = back == doc
    print(f"load->pack->reload identity: {'OK' if same else 'MISMATCH'}")
    if not same:
        for k in doc:
            if doc[k] != back.get(k):
                print("  differs:", k)

def load_bytes(blob: bytes) -> dict:
    return decode(_xor(gzip.decompress(blob)))

def cmd_rename(a):
    doc = load(a.input)
    doc['firstNameIndex'] = a.first
    doc['secondNameIndex'] = a.second
    out = a.output or a.input
    save(doc, out, strict=False)
    print(f"set firstNameIndex={a.first} secondNameIndex={a.second} -> {out}")

def cmd_pack(a):
    doc = load(a.input)
    save(doc, a.output, strict=not a.no_strict)
    print(f"packed -> {a.output}  "
          f"(CompartmentsHash={doc['walker']['CompartmentsHash'][:8]} "
          f"ConnectionsHash={doc['walker']['ConnectionsHash'][:8]})")

# textureRawData is RGBA8888 stored bottom-up (Unity convention); flip on the way
# in/out so PNGs are upright both in files and in-engine. Icon isn't hashed.
def cmd_get_icon(a):
    from PIL import Image
    doc = load(a.input)
    sz = doc['textureSize']
    img = Image.frombytes('RGBA', (sz, sz), bytes(doc['textureRawData']))
    img = img.transpose(Image.FLIP_TOP_BOTTOM)
    out = a.output or (pathlib.Path(a.input).stem + '.png')
    img.save(out)
    print(f"wrote {out}  ({sz}x{sz} RGBA)")

def cmd_set_icon(a):
    from PIL import Image
    doc = load(a.input)
    sz = doc['textureSize']
    img = Image.open(a.png).convert('RGBA')
    if img.size != (sz, sz):
        img = img.resize((sz, sz), Image.LANCZOS)
    img = img.transpose(Image.FLIP_TOP_BOTTOM)        # -> bottom-up for in-engine upright
    doc['textureRawData'] = img.tobytes()
    out = a.output or a.input
    save(doc, out, strict=False)                      # icon doesn't affect any hash
    print(f"set icon {a.png} -> {out}")

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sp = ap.add_subparsers(dest='cmd', required=True)
    p = sp.add_parser('repack'); p.add_argument('input'); p.set_defaults(fn=cmd_repack)
    p = sp.add_parser('rename'); p.add_argument('input'); p.add_argument('first', type=int)
    p.add_argument('second', type=int); p.add_argument('-o', '--output'); p.set_defaults(fn=cmd_rename)
    p = sp.add_parser('pack'); p.add_argument('input'); p.add_argument('-o', '--output', required=True)
    p.add_argument('--no-strict', action='store_true'); p.set_defaults(fn=cmd_pack)
    p = sp.add_parser('get-icon'); p.add_argument('input'); p.add_argument('-o', '--output')
    p.set_defaults(fn=cmd_get_icon)
    p = sp.add_parser('set-icon'); p.add_argument('input'); p.add_argument('png')
    p.add_argument('-o', '--output'); p.set_defaults(fn=cmd_set_icon)
    a = ap.parse_args(); a.fn(a)

if __name__ == '__main__':
    main()