#!/usr/bin/env python3
"""
smallmodel.py - decompilers for a mech's `.torso` and `.engine` files.

Both are single-page `[GameData]` subsystem models stored as flat structs, so
they reuse `datamap.chain_layout()` with their own class chains:

    Entity__GameModel -> Subsystem__GameModel -> Torso__GameModel    1608 bytes
    Entity__GameModel -> Subsystem__GameModel -> Engine__GameModel     64 bytes

Both chains compute to exactly the record size with no anchoring, which is the
check that the member list and alignment are right.

Sources author the numbers as `$(SYMBOL)` macros from a `!include`d defines file
(`!NAME=value` syntax, not `#define`), and the record keeps only the resolved
float, so the symbol is restored by reverse lookup where one matches.

Two quirks worth keeping:

  * The five Torso angles are declared plain `Stuff::Scalar` but `Torso_Tool.cpp`
    multiplies them by `Radians_Per_Degree`, exactly like the Mech spring fields.
  * `TotalCritLocations` is read by **nothing** in the engine -- only the 3DS Max
    export plugin writes it, and the factory reads `TotalSlotsTaken`, which no
    source sets (the record holds the default 1 while every source says 2). It
    is emitted as the constant it always is rather than derived.

    python3 smallmodel.py torso|engine <chassis-record-dir> [-o out]
    python3 smallmodel.py --verify
"""
import argparse, collections, os, re, subprocess, sys

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

C = datamap.CODE
CONTENT = "/home/rich/Repositories/firestorm/Gameleap/mw4/Content"

SPECS = {
    "torso": {
        "chain": [("Entity__GameModel", C + "/mw4/Libraries/Adept/Entity.hpp"),
                  ("Subsystem__GameModel", C + "/mw4/Code/MW4/Subsystem.hpp"),
                  ("Torso__GameModel", C + "/mw4/Code/MW4/Torso.hpp")],
        "size": 1608,
        "factory": C + "/mw4/Code/MW4/Torso_Tool.cpp",
        "defines": CONTENT + "/Defines/MechTorso.defines",
        "include": r"Content\Defines\MechTorso.defines",
        "class": "MechWarrior4::Torso",
        "keys": [("TwistJointName", "twistJointName"),
                 ("PitchJointName", "pitchJointName"),
                 ("LeftArmJointName", "leftArmJointName"),
                 ("RightArmJointName", "rightArmJointName"),
                 ("EyeJointName", "eyeJointName"),
                 ("ArmRatioAngle", "armRatioAngle"),
                 ("TwistSpeed", "twistSpeed"),
                 ("PitchSpeed", "pitchSpeed"),
                 ("TwistRadius", "twistRadius"),
                 ("PitchRadius", "pitchRadius"),
                 ("CageJointName", "cageJointName"),
                 ("CageRatioAngle", "cageRatioAngle")],
    },
    "engine": {
        "chain": [("Entity__GameModel", C + "/mw4/Libraries/Adept/Entity.hpp"),
                  ("Subsystem__GameModel", C + "/mw4/Code/MW4/Subsystem.hpp"),
                  ("Engine__GameModel", C + "/mw4/Code/MW4/Engine.hpp")],
        "size": 64,
        "factory": C + "/mw4/Code/MW4/Engine_Tool.cpp",
        "defines": CONTENT + "/Subsystems/HeatSink.Defines",
        "include": r"Content\Subsystems\HeatSink.defines",
        "class": "Mechwarrior4::Engine",       # lowercase 'w' in every source
        "keys": [("NumHeatSinks", "m_numHeatSinks"),
                 ("TonsPerUpgrade", "m_tonsPerUpgrade"),
                 ("MPSPerUpgrade", "m_mpsPerUpgrade"),
                 ("HeatSinkEfficiency", "m_heatSinkEfficiency")],
    },
}

# Written only by the 3DS Max exporter; the engine never reads it. Uniformly 2.
TOTAL_CRIT_LOCATIONS = "2"


def defines(path):
    """-> {value: symbol} from a `!NAME=value` defines file."""
    out = {}
    if not os.path.exists(path):
        return out
    txt = open(path, encoding="latin-1", errors="replace").read()
    txt = re.sub(r'//[^\n]*', '', txt)
    for name, val in re.findall(r'^\s*!(\w+)\s*=\s*([-\d.]+)\s*$', txt, re.M):
        out.setdefault(round(float(val), 6), name)
    return out


def angle_members(factory):
    out = set()
    if not os.path.exists(factory):
        return out
    txt = open(factory, encoding="latin-1").read()
    for m in re.finditer(r'model->(\w+)\s*=\s*model->\w+\s*\*\s*Radians_Per_Degree', txt):
        out.add(m.group(1))
    return out


def record(mech_dir, kind):
    for fn in os.listdir(mech_dir):
        if fn.lower().endswith(f".{kind}{{gamemodel}}"):
            return open(os.path.join(mech_dir, fn), "rb").read()
    return None


def num(x):
    return f"{int(x)}" if float(x) == int(x) else f"{x:g}"


def decompile(mech_dir, kind):
    """-> [(key, value)] for the [GameData] page."""
    spec = SPECS[kind]
    blob = record(mech_dir, kind)
    if blob is None:
        raise SystemExit(f"no .{kind}{{GameModel}} record in {mech_dir}")
    layout, _ = datamap.chain_layout(spec["chain"], {})
    angles = angle_members(spec["factory"])
    symbols = defines(spec["defines"])

    kv = [("Class", spec["class"]), ("TotalCritLocations", TOTAL_CRIT_LOCATIONS)]
    for key, member in spec["keys"]:
        off, typ, size = layout[member]
        val = datamap.read(blob, off, typ, size, member in angles)
        if typ == "char":
            kv.append((key, val))
        elif typ == "int":
            kv.append((key, str(val)))
        else:
            sym = symbols.get(round(val, 6))
            kv.append((key, f"$({sym})" if sym else num(val)))
    return kv


def emit(kind, kv):
    lines = [f"!include = {SPECS[kind]['include']}", "", "[GameData]"]
    lines += [f"{k}={v}" for k, v in kv]
    return "\r\n".join(lines) + "\r\n"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("kind", nargs="?", choices=sorted(SPECS))
    ap.add_argument("mech_dir", nargs="?")
    ap.add_argument("-o", "--output")
    ap.add_argument("--verify", action="store_true")
    args = ap.parse_args()
    if args.verify:
        here = os.path.dirname(os.path.abspath(__file__))
        raise SystemExit(subprocess.call([sys.executable,
                                          os.path.join(here, "verify_smallmodel.py")]))
    if not args.kind or not args.mech_dir:
        ap.error("kind and mech_dir required")
    text = emit(args.kind, decompile(args.mech_dir, args.kind))
    if args.output:
        with open(args.output, "wb") as fh:
            fh.write(text.encode("latin-1"))
        print(f"wrote {args.output}")
    else:
        sys.stdout.write(text)


if __name__ == "__main__":
    main()