#!/usr/bin/env python3
"""
subsystems.py - rebuild a mech's .subsystems source from its packed record.

    python3 subsystems.py <record.subsystems> <manifest.tsv> [-o out.subsystems]
    python3 subsystems.py --verify              # check against all known chassis

The packed record is `WORD span` followed by one CreateMessage per [Page]
(MWObject::CreateSubsystemStream, MWObject_Tool.cpp:1196). Message layout and
the derivation of every offset below is documented in ../../DECOMPILING.md.

Layout beyond the Entity header (Subsystem.hpp, Weapon.hpp, Armor.hpp):

    off  96  i32   subsystemIndex
    off 100  u8    locationID          -> InternalLocation
    off 104  i32   criticalHitsTaken   -> CriticalHitsTaken
  Armor (152):
    off 108  9xf32 armour points       -> tons, see ARMOR_POINTS_PER_TON
    off 144  i32   m_armorType         -> ArmorType
    off 148  i32   m_internalType      -> InternalType
  Engine (112):
    off 108  i32   m_engineUpgrades    -> EngineUpgrades
  SearchLight (236):
    off 108  char[128] siteName        -> Site
  Weapon (380, or 376 without the trailing field):
    off 108  char[128] siteName        -> Site
    off 236  char[128] ejectSiteName   -> EjectSite
    off 364  i32   groupIndex          -> GroupIndex
    off 368  i32   ammoCount           -> AmmoCount
    off 372  i32   initialAmmoCount
    off 376  i32   m_weaponFacing      -> WeaponFacing   (absent when len == 376)

`m_weaponFacing` is the "MSL 5.04 Rear Firing Weapons" field. V4H packed
champion/griffin/marauder without it (376 bytes) - matching their own release
note "Any new mech will not have rear facing weapons" - while dasher, jenner2c
and thunderbolt have it. Both forms are read; re-packing with our own exe always
emits the 380-byte form.

PAGE NAMES ARE NOT STORED. The packer only serialises page order, so names are
regenerated from the Model= reference with a per-kind counter. They are labels;
the runtime keys on order.
"""
import sys, os, re, glob, struct, argparse, collections

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

OUR_MECH_SOURCE = "/home/rich/Repositories/firestorm/Gameleap/mw4/Content/Mechs"
OUR_RECORDS = "/home/rich/Repositories/FS_Ours_extracted/core/mechs"
OUR_MANIFEST = "/home/rich/Repositories/FS_Ours_extracted/_manifest.tsv"
ARMOR_DATA = "/home/rich/Repositories/firestorm/Gameleap/mw4/Content/Subsystems/Armor.data"

# Derived empirically from 63 aligned chassis and cross-checked against the
# engine's own text tables.
EXEC_STATE = {1: "NeverExecuteState", 2: "AlwaysExecuteState", 6: "ActiveState"}
# InternalDamageObject enum, DamageObject.hpp:205
ZONE = {255: "NullZone", 0: "LeftLeg", 1: "RightLeg", 2: "LeftArm", 3: "RightArm",
        4: "RightTorso", 5: "LeftTorso", 6: "CenterTorso", 7: "Head",
        8: "Special1", 9: "Special2", 10: "VehicleHull", 11: "VehicleWeapon",
        12: "VehicleSpecial", 13: "DefaultZone"}
ARMOR_TYPE = {0: "Standard", 1: "FerroFiberus", 2: "Reactive", 3: "Reflective",
              4: "Solarian"}
# Armor.hpp:275 - a separate 2-value enum, not the armour table. (The engine's own
# InternalTypeAsciiToText returns the typo "Statndard"; TextToAscii wants "Standard".)
INTERNAL_TYPE = {0: "Standard", 1: "EndoSteel"}
ARMOR_KEYS = ["LeftLeg", "RightLeg", "LeftArm", "RightArm", "LeftFrontTorso",
              "RightFrontTorso", "CenterFrontTorso", "CenterRearTorso", "Head"]

CLASS_ENGINE, CLASS_LAMS = 1073, 1183

# Model= basename -> page-name stem. Anything unmatched falls back to the
# weapon-category table below.
PAGE_STEM = {
    "heatsinksubsystem.data": "HeatSink", "armor.data": "Armor",
    "advancedgyrosubsystem.data": "AdvancedGyro", "sensorsubsystem.data": "Sensor",
    "searchlightsubsystem.data": "SearchLight", "jumpjetsubsystem.data": "JumpJet",
    "ecmsubsystem.data": "ECM", "beaglesubsystem.data": "Beagle",
    "lams.data": "LAMS", "narcbeacon.data": "Narc",
}
WEAPON_CATEGORY = [
    ("laserweaponsubsystem", "Beam"), ("pulselaserweaponsubsystem", "Beam"),
    ("ppcweaponsubsystem", "Beam"), ("flamerweaponsubsystem", "Beam"),
    ("lrmweaponsubsystem", "Missile"), ("srmweaponsubsystem", "Missile"),
    ("ssrmweaponsubsystem", "Missile"), ("smrmweaponsubsystem", "Missile"),
    ("missileweaponsubsystem", "Missile"), ("narcbeaconweaponsubsystem", "Narc"),
    ("machinegunweaponsubsystem", "Ballistic"), ("ultraacweaponsubsystem", "Ballistic"),
    ("acweaponsubsystem", "Ballistic"), ("gaussweaponsubsystem", "Ballistic"),
    ("lbxweaponsubsystem", "Ballistic"), ("rtxweaponsubsystem", "Ballistic"),
]


def armor_points_per_ton(path=ARMOR_DATA):
    txt = open(path, "rb").read().decode("latin-1")

    def get(key, default):
        m = re.search(rf"^{key}=(\d+)", txt, re.M | re.I)
        return int(m.group(1)) if m else default
    return {0: get("PointsPerStandardTon", 32), 1: get("PointsPerFerroTon", 38),
            2: get("PointsPerReactiveTon", 30), 3: get("PointsPerReflectiveTon", 30),
            4: get("PointsPerSolarianTon", 60)}


def load_manifest(path, package="core.mw4"):
    """record id -> entry name, for resolving dataListID back to a Model= path."""
    out = {}
    with open(path, encoding="latin-1") as fh:
        for line in fh:
            parts = line.rstrip("\n").split("\t")
            if len(parts) > 2 and parts[0].lower() == package.lower():
                out[int(parts[1])] = parts[2]
    return out


def cstr(msg, off, size=128):
    return msg[off:off + size].split(b"\0")[0].decode("latin-1").strip()


def group_flags_to_list(flags):
    """groupIndex is a BITMASK, not an index (Weapon_Tool.cpp ~76).

    A page may carry several `GroupIndex=` lines; the factory ORs
    `1 << (n-1)` for each. Returns the group numbers, so the caller can emit one
    line per group.
    """
    return [n for n in range(1, 7) if flags & (1 << (n - 1))]


def page_name(model, counters):
    base = os.path.basename(model.replace("\\", "/")).lower()
    stem = PAGE_STEM.get(base)
    if stem is None:
        low = model.lower().replace("\\", "/")
        stem = next((s for key, s in WEAPON_CATEGORY if key in low), None)
    if stem is None:
        stem = "Torso" if base.endswith(".torso") else \
               "Engine" if base.endswith(".engine") else "Subsystem"
    counters[stem] += 1
    # Singletons keep a bare name, matching how the sources are written.
    if stem in ("Armor", "AdvancedGyro", "Sensor", "SearchLight", "Torso",
                "Engine", "ECM", "Beagle", "LAMS"):
        return stem if counters[stem] == 1 else f"{stem}{counters[stem]}"
    return f"{stem}{counters[stem]}"


def decompile(record_path, manifest, ppt=None):
    ppt = ppt or armor_points_per_ton()
    data = open(record_path, "rb").read()
    # Model= is written relative to the mech's own directory in the source, but the
    # package stores the full entry path. Use the parent folder, not the file name:
    # Black Hawk's chassis files are nova.* inside mechs/blackhawk/.
    own_dir = "mechs\\" + os.path.basename(os.path.dirname(record_path)).lower() + "\\"
    counters = collections.Counter()
    pages = []
    for _off, msg in mw4msg.walk(data):
        cid, n = mw4msg.class_id(msg), len(msg)
        model = manifest.get(mw4msg.record_id(msg), "?")
        if model.lower().startswith(own_dir):
            model = model[len(own_dir):]
        kv = collections.OrderedDict()
        kv["Model"] = model
        kv["ExecutionState"] = EXEC_STATE.get(struct.unpack_from("<i", msg, 76)[0], "?")
        kv["InternalLocation"] = ZONE.get(msg[100], f"?{msg[100]}")
        crit = struct.unpack_from("<i", msg, 104)[0]

        if n == 152:                                            # Armor
            vals = struct.unpack_from("<9f", msg, 108)
            atype, itype = struct.unpack_from("<2i", msg, 144)
            kv["ArmorType"] = ARMOR_TYPE.get(atype, str(atype))
            kv["InternalType"] = INTERNAL_TYPE.get(itype, str(itype))
            mult = ppt.get(atype, 32)
            for key, v in zip(ARMOR_KEYS, vals):
                kv[key] = f"{v / mult:g}"
        elif n == 112 and cid == CLASS_ENGINE:
            kv["EngineUpgrades"] = str(struct.unpack_from("<i", msg, 108)[0])
        elif n == 112 and cid == CLASS_LAMS:
            ammo = struct.unpack_from("<i", msg, 108)[0]
            if ammo >= 0:
                kv["AmmoCount"] = str(ammo)
        elif n == 236:                                          # SearchLight
            kv["Site"] = cstr(msg, 108)
        elif n in (376, 380):                                   # Weapon
            kv["Site"] = cstr(msg, 108)
            eject = cstr(msg, 236)
            grp, ammo, _init = struct.unpack_from("<3i", msg, 364)
            kv["GroupIndex"] = group_flags_to_list(grp)
            if ammo >= 0:
                kv["AmmoCount"] = str(ammo)
            if eject:
                kv["EjectSite"] = eject
            if n == 380:
                facing = struct.unpack_from("<i", msg, 376)[0]
                if facing:
                    kv["WeaponFacing"] = str(facing)
        if crit:
            kv["CriticalHitsTaken"] = str(crit)
        pages.append((page_name(model, counters), kv))
    return pages


def emit(pages):
    out = []
    for name, kv in pages:
        out.append(f"[{name}]")
        for k, v in kv.items():
            if isinstance(v, list):                     # GroupIndex: one line per group
                out.extend(f"{k}={n}" for n in v)
            else:
                out.append(f"{k}={v}")
        out.append("")
    return "\r\n".join(out) + "\r\n"


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("record", nargs="?")
    ap.add_argument("manifest", nargs="?", default=OUR_MANIFEST)
    ap.add_argument("-o", "--out")
    ap.add_argument("--verify", action="store_true")
    args = ap.parse_args()

    if args.verify:
        here = os.path.dirname(os.path.abspath(__file__))
        os.execvp("python3", ["python3", os.path.join(here, "verify_subsystems.py")])
    if not args.record:
        ap.error("give a packed .subsystems record, or --verify")

    pages = decompile(args.record, load_manifest(args.manifest))
    text = emit(pages)
    if args.out:
        os.makedirs(os.path.dirname(args.out), exist_ok=True)
        open(args.out, "wb").write(text.encode("latin-1"))
        print(f"{os.path.basename(args.record)}: {len(pages)} pages -> {args.out}")
    else:
        sys.stdout.write(text)


if __name__ == "__main__":
    main()