"""Extract every MeshRenderer a Unity prefab or scene places, with its world transform.

`unity_extract.py` walks the same hierarchy to find colliders; this walks it to find
*visible* geometry, so the level can be rebuilt from the original models instead of
from boxes. It reuses that module's transform maths, which is already validated
against known collider sizes.

Output per renderer: the source FBX and which mesh inside it, the Unity materials the
renderer overrides with, and the composed world TRS.
"""
import json
import os
import re
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from import_config import CFG as _CFG
_UNITY = os.path.expanduser(_CFG['unityRoot'])
# unity_extract resolves its project root from argv[1]; point it at the active tree.
if len(sys.argv) < 2:
    sys.argv = [sys.argv[0], os.path.join(_UNITY, 'Assets', 'Scenes')]

import unity_extract as ue  # noqa: E402

MAX_DEPTH = 10
# the project's TagManager: a layer some games render only into a visibility or shadow mask.
SHADOW_LAYER = 18
_LOD_SUFFIX = re.compile(r'_(LOD)?[Ll]?[123]$')


def _is_lod_child(name):
    """LODGroups list LOD0 first; the _L1/_L2 rungs are the same shape at lower detail."""
    return bool(re.search(r'_L[123]$', name)) or bool(re.search(r'_LOD[123]$', name))


def _under_lodgroup(docs, tbody):
    """Does this transform's parent GameObject carry a LODGroup (class 205)?

    A name alone is not evidence of a LOD rung: some asset packs ship one FBX per detail level,
    so a `..._L2` model placed by hand under a plain group is a whole object. The scene this was
    built on had 16 of those — the four side-middle wall pieces and all twelve corner
    staircases — and dropping them by name emptied the corners and the middle of every side.
    """
    f = ue.fid(tbody, 'm_Father')
    pt = docs.get(f)
    if not pt or pt[0] != '4':
        return False
    pg = docs.get(ue.fid(pt[1], 'm_GameObject'))
    if not pg:
        return False
    for m in re.finditer(r'- component: \{fileID: (\d+)\}', pg[1]):
        c = docs.get(m.group(1))
        if c and c[0] == '205':
            return True
    return False


class RendererWalk:
    def __init__(self):
        self.out = []

    def walk(self, path, docs, tid, world, depth, mods_for_file, tag):
        if depth > MAX_DEPTH:
            return
        entry = docs.get(tid)
        if not entry:
            return
        cls, body = entry
        if cls == '1001':
            self.prefab(path, docs, tid, body, world, depth, tag, mods_for_file)
            return
        if 'm_LocalPosition' not in body:
            return  # a stripped stub standing in for a nested prefab; the 1001 sweep places it

        trs = ue.local_trs(body)
        mod = mods_for_file.get(tid)
        if mod:
            trs = ue.apply_mods(trs, mod)
        goid = ue.fid(body, 'm_GameObject')
        name = ue.go_name(docs, goid)
        if depth == 0 and path.endswith('.unity'):
            dp = ue.patch_move_for(name, trs[0])          # level patch: slide a scene object
            if any(dp):
                trs = ((trs[0][0] + dp[0], trs[0][1] + dp[1], trs[0][2] + dp[2]), trs[1], trs[2])
        w = ue.compose(world, trs)

        if not ue.go_active(docs, goid, mods_for_file):
            return
        layer = ue.go_layer(docs, goid, mods_for_file)

        gd = docs.get(goid)
        if gd and not (_is_lod_child(name) and _under_lodgroup(docs, body)):
            self.components(path, docs, gd[1], name, w, tag, layer, mods_for_file)

        seg = ''
        if 'm_Children:' in body and 'm_Father:' in body:
            seg = body[body.find('m_Children:'):body.find('m_Father:')]
        for m in re.finditer(r'- \{fileID: (\d+)\}', seg):
            self.walk(path, docs, m.group(1), w, depth + 1, mods_for_file, tag)

    def components(self, path, docs, gbody, name, world, tag, layer, mods=None):
        """Pair the MeshFilter and MeshRenderer hanging off one GameObject."""
        mesh = None
        mats = None
        enabled = True
        casts = True
        for m in re.finditer(r'- component: \{fileID: (\d+)\}', gbody):
            cd = docs.get(m.group(1))
            if not cd:
                continue
            ccls, cbody = cd
            if ccls == '33':  # MeshFilter
                mm = re.search(r'm_Mesh: \{fileID: (-?\d+), guid: ([0-9a-f]{32})', cbody)
                if mm:
                    mesh = (mm.group(2), mm.group(1))
            elif ccls == '23':  # MeshRenderer
                if not ue.comp_enabled(cbody, m.group(1), mods):
                    enabled = False
                cs = re.search(r'm_CastShadows: (\d)', cbody)
                if cs and cs.group(1) == '0':
                    casts = False
                seg = cbody[cbody.find('m_Materials:'):]
                seg = seg[:seg.find('\n  m_', 1)] if '\n  m_' in seg[1:] else seg
                mats = re.findall(r'guid: ([0-9a-f]{32})', seg)
        if not mesh or not enabled:
            return
        mesh_path = ue.guid2path.get(mesh[0])
        if not mesh_path or not os.path.exists(mesh_path):
            return
        if not mesh_path.lower().endswith(('.fbx', '.obj')):
            return  # a built-in primitive; the collider pass already covers those

        (p, q, s) = world
        self.out.append({
            'name': name,
            'tag': tag,
            'layer': layer,
            'fbx': mesh_path,
            'meshFileId': mesh[1],
            'materials': [ue.guid2path.get(g) for g in (mats or []) if ue.guid2path.get(g)],
            'castShadow': casts,
            'pos': [round(v, 5) for v in p],
            'quat': [round(v, 6) for v in q],
            'scale': [round(v, 5) for v in s],
        })

    def prefab(self, path, docs, tid, body, world, depth, tag, outer_mods=None):
        m = re.search(r'm_SourcePrefab: \{fileID: \d+, guid: ([0-9a-f]{32})', body)
        if not m:
            return
        src = ue.guid2path.get(m.group(1))
        if not src or not os.path.exists(src):
            return
        pdocs = ue.parse_file(src)
        if not pdocs:
            return
        modmap = ue.instance_mods(body, docs, tid, pdocs, outer_mods)
        proot = next((pid for pid, (c, b) in pdocs.items() if c == '4' and ue.fid(b, 'm_Father') == '0'), None)
        if proot is not None and not ue.go_active(pdocs, ue.fid(pdocs[proot][1], 'm_GameObject'), modmap):
            return   # the instance switches the whole prefab off (a scene-level m_IsActive override)

        newtag = tag or os.path.basename(src).replace('.prefab', '')
        if depth == 0:
            # scene-level instance: honour level patches (tools/level_patches.json), the same way
            # unity_colliders2d does, so what is drawn and what is hit stay identical.
            root = next((pid for pid, (c, b) in pdocs.items() if c == '4' and ue.fid(b, 'm_Father') == '0'), None)
            if root is not None:
                rp = ue.apply_mods(ue.local_trs(pdocs[root][1]), modmap.get(root, {}))
                if ue.patch_skips_instance(newtag, ue.compose(world, rp)[0]):
                    return
        self.expand(src, pdocs, world, depth + 1, modmap, newtag)

    def expand(self, path, docs, world, depth, mods, tag):
        """Everything one file places: its root transforms, then its own nested prefab instances.

        The second half is not optional. A nested instance appears in its parent's m_Children only
        as a stripped stub the transform walk cannot follow, so a prefab built out of other prefabs
        — which most of a kit-built level is — yields only the handful of objects stored directly in its
        own file unless the 1001 records are swept as well. Sweeping only the top-level file, as
        this used to, meant one placed prefab kept 16 of its 75 pieces and another
        arrived empty.
        """
        if depth > MAX_DEPTH:
            return
        for tid, (cls, body) in docs.items():
            if cls == '4' and ue.fid(body, 'm_Father') == '0' and 'm_LocalPosition' in body:
                self.walk(path, docs, tid, world, depth, mods, tag)

        for tid, (cls, body) in docs.items():
            if cls != '1001':
                continue
            chain = []
            cur = ue.fid(body, 'm_TransformParent')
            guard = 0
            while cur != '0' and cur in docs and guard < 40:
                ccls, cbody = docs[cur]
                if ccls != '4' or 'm_LocalPosition' not in cbody:
                    break
                chain.append((cur, cbody))
                cur = ue.fid(cbody, 'm_Father')
                guard += 1
            w = world
            for cid, cbody in reversed(chain):
                trs = ue.local_trs(cbody)
                mod = mods.get(cid)
                if mod:
                    trs = ue.apply_mods(trs, mod)
                w = ue.compose(w, trs)
            self.prefab(path, docs, tid, body, w, depth, tag, mods)


IDENT = ((0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 1.0), (1.0, 1.0, 1.0))


def _add_patched_instances(rw, docs):
    """Extra placements a level patch adds, expanded from the scene's own copy of the prefab."""
    for tid, (cls, body) in docs.items():
        if cls != '1001': continue
        m = re.search(r'm_SourcePrefab: \{fileID: \d+, guid: ([0-9a-f]{32})', body)
        src = ue.guid2path.get(m.group(1)) if m else None
        if not src: continue
        name = os.path.basename(src).replace('.prefab', '')
        adds = ue.patch_added_instances(name)
        if not adds: continue
        pdocs = ue.parse_file(src)
        root = next((pid for pid, (c, b) in pdocs.items() if c == '4' and ue.fid(b, 'm_Father') == '0'), None)
        if root is None: continue
        for (pos, quat) in adds:
            modmap = {root: {'m_LocalPosition.x': pos[0], 'm_LocalPosition.y': pos[1], 'm_LocalPosition.z': pos[2],
                             'm_LocalRotation.x': quat[0], 'm_LocalRotation.y': quat[1], 'm_LocalRotation.z': quat[2], 'm_LocalRotation.w': quat[3]}}
            rw.expand(src, pdocs, ((0.0, 0.0, 0.75), (0.0, 0.0, 0.0, 1.0), (1.0, 1.0, 1.0)), 1, modmap, name)
        break


def extract(path):
    """Every renderer a scene or prefab places, in that file's own space."""
    docs = ue.parse_file(path)
    rw = RendererWalk()
    rw.expand(path, docs, IDENT, 0, {}, None)
    if path.endswith('.unity'):
        _add_patched_instances(rw, docs)
    return _dedupe(rw.out)


def _rank(item):
    """How much a copy deserves to be the one kept. Higher wins.

    Coincident copies arise two ways, and they are not the same thing:

    * Unity keeps a second copy of most level geometry on the SHADOWS layer — a stand-in the
      original rendered only into its line-of-sight mask, never to the screen. It sits at exactly
      the same transform as the piece you can see, so it must never win: dropping SHADOWS after
      deduping deletes the visible piece along with it, which is what left whole structures as
      scattered fragments.
    * An object inside a nested prefab is also seen through the outer file, which is where Unity
      stores the override putting it on a gameplay layer. That copy is the authoritative one.
    """
    if item['layer'] == SHADOW_LAYER:
        return 0
    return 2 if item['layer'] != 0 else 1


def _dedupe(items):
    """One renderer reached twice is still one renderer."""
    best = {}
    order = []
    for it in items:
        key = (it['fbx'], it['meshFileId'], tuple(it['pos']), tuple(it['quat']), tuple(it['scale']))
        if key not in best:
            best[key] = it
            order.append(key)
        elif _rank(it) > _rank(best[key]):
            best[key] = it
    return [best[k] for k in order]


if __name__ == '__main__':
    print(json.dumps(extract(sys.argv[1] if len(sys.argv) > 1 else ue.scene_path_from_config()), indent=1))