#!/usr/bin/env python3
"""
assembly.py - show how a mech's parts bolt together.

There is no `.erf` viewer on Linux, but the assembly itself is not hidden: the
`.armature` gives every joint's parent, offset and rotation, and the `.contents`
plus `armaturevideo/` say which geometry hangs off each joint. Printing that as
a tree answers "what are the parts and how do they fit" without a renderer.

For actual 3D you need MW4Ed2 on the Windows box (Gameleap/mw4/run-editor.bat) --
its Game View renders a loaded mech through DDrawCompat.

    python3 assembly.py <mech-source-dir> [--geometry-only]
"""
import argparse, collections, os, re, sys


def pages(path):
    """-> OrderedDict pageName -> [(key, value)]"""
    txt = re.sub(r'//[^\n]*', '', open(path, "rb").read().decode("latin-1"))
    out, cur = collections.OrderedDict(), None
    for line in txt.splitlines():
        line = line.strip()
        if not line or line.startswith("!"):
            continue
        m = re.match(r'^\[([^\]]+)\]$', line)
        if m:
            cur = m.group(1)
            out.setdefault(cur, [])
        elif "=" in line and cur is not None:
            k, v = line.split("=", 1)
            out[cur].append((k.strip(), v.strip()))
    return out


def find(mech_dir, ext):
    for f in sorted(os.listdir(mech_dir)):
        if f.lower().endswith(ext) and "}" not in f:
            return os.path.join(mech_dir, f)
    return None


def geometry_for(mech_dir, part):
    """The .erf a joint draws, via armaturevideo/<part>.video.

    Takes the [lod] block specifically -- a torso's file starts with the
    running-lights block, whose geometry is not the part.
    """
    vid = os.path.join(mech_dir, "armaturevideo", part + ".video")
    if not os.path.exists(vid):
        return None
    t = open(vid, "rb").read().decode("latin-1")
    m = re.search(r'^\[lod\][^\[]*?Geometry=([^\r\n]+)', t, re.M | re.I | re.S)
    if not m:
        m = re.search(r'Geometry=([^\r\n]+)', t)
    return m.group(1).strip() if m else None


def build(mech_dir):
    arm = find(mech_dir, ".armature")
    con = find(mech_dir, ".contents")
    if not arm:
        raise SystemExit(f"no .armature in {mech_dir}")

    apages = pages(arm)
    children = collections.defaultdict(list)
    info = {}
    for name, kv in apages.items():
        d = dict(kv)
        info[name] = d
        for k, v in kv:
            if k.lower() == "child":
                children[name].append(v.strip())

    models = {}
    if con:
        for name, kv in pages(con).items():
            for k, v in kv:
                if k.lower() == "model":
                    models[name.lower()] = v.strip()

    parented = {c for cs in children.values() for c in cs}
    roots = [n for n in apages if n not in parented]
    return apages, children, info, models, roots


def render(mech_dir, geometry_only=False):
    apages, children, info, models, roots = build(mech_dir)
    sizes = {f.lower(): os.path.getsize(os.path.join(mech_dir, f))
             for f in os.listdir(mech_dir) if f.lower().endswith(".erf")}
    lines = []

    def walk(name, depth, last, prefix):
        d = info.get(name, {})
        model = models.get(name.lower(), "")
        part = None
        m = re.match(r'armaturedata[\\/](.+)\.data', model, re.I)
        if m:
            part = m.group(1)
        geo = geometry_for(mech_dir, part) if part else None
        size = sizes.get((geo or "").lower())
        tag = ""
        if geo:
            tag = f"  <- {geo}" + (f" ({size:,} B)" if size else "")
        elif model and model.lower() != "basic.data":
            tag = f"  <- {model}"
        trans = d.get("Translation", d.get("translation", ""))
        rot = d.get("Rotation", d.get("rotation", ""))
        pos = f"  @[{trans}]" if trans and trans.strip() not in ("0.0 0.0 0.0", "0 0 0") else ""
        if geometry_only and not geo:
            pass
        else:
            branch = "" if depth == 0 else ("`-- " if last else "|-- ")
            lines.append(f"{prefix}{branch}{name}{tag}{pos}")
        kids = children.get(name, [])
        newprefix = prefix + ("" if depth == 0 else ("    " if last else "|   "))
        for i, k in enumerate(kids):
            walk(k, depth + 1, i == len(kids) - 1, newprefix)

    for r in roots:
        walk(r, 0, True, "")
    return lines, sizes


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("mech_dir")
    ap.add_argument("--geometry-only", action="store_true",
                    help="hide joints and sites that draw nothing")
    args = ap.parse_args()
    lines, sizes = render(args.mech_dir, args.geometry_only)
    name = os.path.basename(args.mech_dir.rstrip("/"))
    print(f"=== {name}: {len(lines)} nodes, {len(sizes)} .erf, "
          f"{sum(sizes.values()):,} B of geometry ===")
    print("\n".join(lines))


if __name__ == "__main__":
    main()