#!/usr/bin/env python3
"""Extract SAND MonoBehaviour data (loot tables, crafting graph, …) using IL2CPP typetrees.

Builds a UnityPy TypeTreeGenerator from GameAssembly.dll + global-metadata.dat (Unity
6000.0.40f1), loads the relevant bundles into one environment so PPtr references resolve,
reads each MonoBehaviour against its generated typetree, and rewrites PPtrs as the target's
m_Name where known. Output: structured JSON in extracted/.
"""
import os, sys, json, UnityPy
from UnityPy.helpers.TypeTreeGenerator import TypeTreeGenerator

GAME = "/mnt/d/SteamLibrary/steamapps/common/Sand Playtest"
BD = os.path.join(GAME, "Sand_Data/StreamingAssets/aa/StandaloneWindows64")
META = os.path.join(GAME, "Sand_Data/il2cpp_data/Metadata/global-metadata.dat")
DLL = os.path.join(GAME, "GameAssembly.dll")
OUT = "/home/downloadpizza/sand_tools/extracted"
UNITY = "6000.0.40f1"

def build_generator():
    g = TypeTreeGenerator(UNITY)
    g.load_il2cpp(open(DLL, "rb").read(), open(META, "rb").read())
    return g

_node_cache = {}
def nodes_for(gen, script):
    key = (script.m_AssemblyName, script.m_Namespace, script.m_ClassName)
    if key not in _node_cache:
        full = (script.m_Namespace + "." if script.m_Namespace else "") + script.m_ClassName
        try:
            _node_cache[key] = gen.get_nodes(script.m_AssemblyName, full)
        except Exception as e:
            _node_cache[key] = None
    return _node_cache[key]

def load_env(*bundles):
    paths = [os.path.join(BD, b) for b in bundles]
    return UnityPy.load(*paths)

def build_name_index(env):
    """path_id -> m_Name for every object we can cheaply name (GameObjects, named MBs)."""
    idx = {}
    for o in env.objects:
        try:
            if o.type.name in ("GameObject",):
                idx[o.path_id] = o.read().m_Name
        except Exception:
            pass
    return idx

def read_mb(gen, obj, name_idx=None):
    """Read a MonoBehaviour into a plain dict; resolve PPtr dicts to {'->': name|pathid}."""
    d = obj.read()
    try:
        script = d.m_Script.read()
    except Exception:
        return None
    nodes = nodes_for(gen, script)
    if not nodes:
        return None
    tree = obj.read_typetree(nodes)
    return _clean(tree, name_idx, obj)

def _clean(v, name_idx, obj):
    if isinstance(v, dict):
        # PPtr shape
        if set(v.keys()) >= {"m_FileID", "m_PathID"} and len(v) <= 3:
            pid = v["m_PathID"]
            if pid == 0:
                return None
            nm = name_idx.get(pid) if name_idx else None
            return {"ref": nm or pid}
        return {k: _clean(x, name_idx, obj) for k, x in v.items()}
    if isinstance(v, list):
        return [_clean(x, name_idx, obj) for x in v]
    return v

if __name__ == "__main__":
    print("building generator…", flush=True)
    gen = build_generator()
    print("generator ready", flush=True)
    # quick validation on one loot table
    env = load_env("lootsets_assets_all.bundle", "sand_monoscripts.bundle",
                   "epb_assets_all.bundle", "configuration_assets_all.bundle")
    name_idx = build_name_index(env)
    print("name index:", len(name_idx), "entries", flush=True)
    o = next(x for x in env.objects if x.type.name == "MonoBehaviour"
             and x.read().m_Name == "POIShipMediumWeapons")
    import pprint
    pprint.pprint(read_mb(gen, o, name_idx))