#!/usr/bin/env python3
"""Census of ECS components across every EntityBlueprint in epb_assets_all.bundle.

Decodes all 1446 blueprints' Odin payloads and tallies which $type components
appear and how often. Use it to answer "is component X authored in data at all?"
firsthand (Odin stores type names as UTF-16, so an ascii grep gives false
negatives — this decodes properly). Optional filter substring narrows output.

    python bundle/component_census.py            # combat-ish components
    python bundle/component_census.py Damage     # only names containing 'Damage'
    python bundle/component_census.py ''          # everything
"""
import os, sys, json, UnityPy
from collections import Counter
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"
COMBAT = ("Damage", "Health", "Weapon", "Attack", "Melee", "Shoot", "Projectile",
          "Penetra", "Overheat", "Hit", "AoE", "Armor", "Resist")


def walk(n):
    if isinstance(n, dict):
        t = n.get("$type")
        if t:
            yield t
        for k, v in n.items():
            if k != "$type":
                yield from walk(v)
    elif isinstance(n, list):
        for v in n:
            yield from walk(v)


def main():
    filt = sys.argv[1] if len(sys.argv) > 1 else None  # None -> combat preset
    gen = TypeTreeGenerator(UNITY)
    gen.load_il2cpp(open(DLL, "rb").read(), open(META, "rb").read())
    env = UnityPy.load(os.path.join(BUNDLES, "epb_assets_all.bundle"),
                       os.path.join(BUNDLES, "sand_monoscripts.bundle"))
    comp = Counter(); ndone = 0; nodecache = {}
    for path, obj in env.container.items():
        if obj.type.name != "GameObject":
            continue
        eb = None
        for c in obj.read().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:
            sb = eb.read_typetree(nodecache[full]).get("serializationData", {}).get("SerializedBytes")
        except Exception:
            continue
        if not sb:
            continue
        try:
            p = odin_read.parse(bytes(sb))
        except Exception:
            continue
        seen = set()
        for root in ("roots", "items"):
            for t in walk(p.get(root)):
                seen.add(t)
        for t in seen:
            comp[t] += 1
        ndone += 1
    print("blueprints parsed:", ndone, " distinct components:", len(comp))
    for t, n in comp.most_common():
        short = t.split(".")[-1]
        if filt is None:
            if any(w in short for w in COMBAT):
                print("%4d  %s" % (n, t))
        elif filt == "" or filt in t:
            print("%4d  %s" % (n, t))


if __name__ == "__main__":
    main()