#!/usr/bin/env python3
"""For every EntityBlueprint carrying a CraftingWorkbenchDataComponent, list the
CraftingRecipeBundle assets it references (recipeBundles). Resolves the Odin
ext_index refs against the blueprint's serializationData.ReferencedUnityObjects
(PPtrs) -> CraftingRecipeBundle MonoBehaviour -> m_Name.
"""
import os, sys, json, UnityPy
from UnityPy.helpers.TypeTreeGenerator import TypeTreeGenerator
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import odin_read

GAME = "/mnt/d/SteamLibrary/steamapps/common/Sand Playtest"
META = os.path.join(GAME, "Sand_Data/il2cpp_data/Metadata/global-metadata.dat")
DLL = os.path.join(GAME, "GameAssembly.dll")
BUNDLES = "/home/downloadpizza/sand_tools/bundles"
UNITY = "6000.0.40f1"
WB = "CraftingWorkbenchDataComponent"


def find_recipe_indices(node, out):
    """Find a node of $type CraftingWorkbenchDataComponent and collect ext_index refs."""
    if isinstance(node, dict):
        if node.get("$type", "").endswith("CraftingWorkbenchDataComponent"):
            def grab(n):
                if isinstance(n, odin_read.Ref):
                    if n.kind == "ext_index":
                        out.append(n.val)
                elif isinstance(n, dict):
                    for v in n.values():
                        grab(v)
                elif isinstance(n, list):
                    for v in n:
                        grab(v)
            grab(node)
        for v in node.values():
            find_recipe_indices(v, out)
    elif isinstance(node, list):
        for v in node:
            find_recipe_indices(v, out)


def main():
    gen = TypeTreeGenerator(UNITY)
    gen.load_il2cpp(open(DLL, "rb").read(), open(META, "rb").read())
    # load the workbench prefabs (epb), the recipe-bundle assets they reference
    # (craftingrecipes), and monoscripts so m_Script resolves
    paths = [os.path.join(BUNDLES, b) for b in
             ("epb_assets_all.bundle", "craftingrecipes_assets_all.bundle",
              "sand_monoscripts.bundle")]
    env = UnityPy.load(*paths)
    nodecache = {}

    # path_id -> CraftingRecipeBundle name, to resolve the workbench's ext refs
    bundle_names = {}
    for o in env.objects:
        if o.type.name == "MonoBehaviour":
            try:
                d = o.read()
                if d.m_Script.read().m_ClassName == "CraftingRecipeBundle":
                    bundle_names[o.path_id] = d.m_Name
            except Exception:
                pass

    for path, obj in env.container.items():
        if obj.type.name != "GameObject":
            continue
        go = obj.read()
        eb = None
        for c in go.m_Components:
            co = c.read(); r = co.object_reader
            if r.type.name == "MonoBehaviour" and co.m_Script.read().m_ClassName == "EntityBlueprint":
                eb = r; break
        if eb is None:
            continue
        sc = eb.read().m_Script.read()
        full = (sc.m_Namespace + "." if sc.m_Namespace else "") + sc.m_ClassName
        if full not in nodecache:
            nodecache[full] = json.loads(gen.get_nodes_as_json(sc.m_AssemblyName, full))
        try:
            tt = eb.read_typetree(nodecache[full])
        except Exception:
            continue
        sd = tt.get("serializationData", {})
        sb = sd.get("SerializedBytes")
        if not sb:
            continue
        try:
            parsed = odin_read.parse(bytes(sb))
        except Exception:
            continue
        idxs = []
        for root in ("roots", "items"):
            find_recipe_indices(parsed.get(root), idxs)
        if not idxs:
            continue
        refs = sd.get("ReferencedUnityObjects", [])
        print("\n%s" % go.m_Name)
        for i in idxs:
            pid = refs[i].get("m_PathID") if i < len(refs) else None
            print("   recipeBundle[%d] -> %s" % (i, bundle_names.get(pid, "<unresolved pid %s>" % pid)))


if __name__ == "__main__":
    main()