#!/usr/bin/env python3
"""
damage.py - decompiler for a mech `<chassis>.damage` file.

The record is a bare concatenation of variable-length objects with no index and
no length prefixes: parsing means walking forward, reading a classID, and using
it to decide what follows. Written by `MWObject::CreateDamageStream`
(MWObject_Tool.cpp:1149), which iterates the source pages in order and dispatches
on whether a page carries a `DamageZone` entry.

Armour page -- DamageObject::ConstructDamageObjectStream (DamageObject.cpp:157):

    classID, baseArmorValue, currentArmorValue, scaleSplashDamage,
    damageObjectName (MString), internalDamageZoneID, armorZone, damageLevel,
    armorType, maxArmorValue, attachedToZone

Internal page -- InternalDamageObject::ConstructInternalDamageObjectStream
(DamageObject.cpp:931) plus the MW4 subclass (MWDamageObject.cpp:88):

    classID, baseInternalDamage, currentInternalDamage,
    parentEntityName (MString), damageMode, damageZone, damagePropagationZone,
    internalType, attachedTo, damageEffects[], missileSlots, projectileSlots,
    beamSlots, omniSlots

MString is an int length NOT counting the terminator, the characters, then a NUL
-- the same encoding as the {FootSteps} stream.

Two traps:

  * The armour page stores its own name, but the internal page does NOT. Internal
    page names are reconstructed as `<Zone>Internal`, which every one of the 89
    mech .damage files follows, using the damageZone that IS stored.
  * ArmorZone and InternalZone are DIFFERENT enums. ArmorZone has
    CenterRearTorso at 7 and Head at 8; InternalZone has Head at 7 and no rear
    torso entry. Conflating them silently mislabels head and torso zones.

    python3 damage.py <chassis-record-dir> [-o out.damage]
    python3 damage.py --verify
"""
import argparse, collections, os, re, struct, subprocess, sys

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

MANIFEST = "/home/rich/Repositories/FS_Ours_extracted/_manifest.tsv"

ARMOR_CLASS = 468            # Adept::DamageObject
INTERNAL_CLASS = 1162        # MechWarrior4::MWInternalDamageObject

# DamageObject.hpp:363 -- note CenterRearTorso, absent from the internal enum.
ARMOR_ZONE = {
    -1: "NullZone", 0: "LeftLeg", 1: "RightLeg", 2: "LeftArm", 3: "RightArm",
    4: "RightTorso", 5: "LeftTorso", 6: "CenterTorso", 7: "CenterRearTorso",
    8: "Head", 9: "Special1", 10: "Special2", 11: "DefaultZone",
}
# DamageObject.hpp:204
INTERNAL_ZONE = {
    -1: "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",
}
# DamageObject.hpp:185
DAMAGE_MODE = {
    0: "GeneralDamageMode", 1: "GimpLeftDamageMode", 2: "GimpRightDamageMode",
    3: "DestructionDamageMode", 4: "DetachableDamageMode", 5: "EngineDamageMode",
    6: "NextDamageMode", 7: "HeadShotDamageMode", 8: "GyroHitDamageMode",
    9: "TorsoLeftDamageMode", 10: "TorsoRightDamageMode",
    11: "ArmLeftDamageMode", 12: "ArmRightDamageMode",
}


class Reader:
    def __init__(self, blob):
        self.b, self.o = blob, 0

    def i32(self):
        v = struct.unpack_from("<i", self.b, self.o)[0]
        self.o += 4
        return v

    def u32(self):
        v = struct.unpack_from("<I", self.b, self.o)[0]
        self.o += 4
        return v

    def f32(self):
        v = struct.unpack_from("<f", self.b, self.o)[0]
        self.o += 4
        return v

    def mstring(self):
        n = self.u32()
        s = self.b[self.o:self.o + n].decode("latin-1")
        self.o += n + 1                       # length excludes the terminator
        return s

    def done(self):
        return self.o >= len(self.b)


def parse(blob):
    """-> [(kind, {field: value})] in stream order."""
    r = Reader(blob)
    out = []
    while not r.done():
        class_id = r.i32()
        if class_id == ARMOR_CLASS:
            rec = {
                "baseArmorValue": r.f32(),
                "currentArmorValue": r.f32(),
                "scaleSplashDamage": r.f32(),
                "name": r.mstring(),
                "internalDamageZone": r.i32(),
                "armorZone": r.i32(),
                "damageLevel": r.i32(),
                "armorType": r.i32(),
                "maxArmorValue": r.f32(),
                "attachedToZone": r.i32(),
            }
            out.append(("armor", rec))
        elif class_id == INTERNAL_CLASS:
            rec = {
                "baseInternalDamage": r.f32(),
                "currentInternalDamage": r.f32(),
                "parentEntityName": r.mstring(),
                "damageMode": r.i32(),
                "damageZone": r.i32(),
                "damagePropagationZone": r.i32(),
                "internalType": r.i32(),
                "attachedTo": r.i32(),
            }
            count = r.i32()
            rec["effects"] = [(r.u32(), r.f32()) for _ in range(count)]
            rec["missileSlots"] = r.i32()
            rec["projectileSlots"] = r.i32()
            rec["beamSlots"] = r.i32()
            rec["omniSlots"] = r.i32()
            out.append(("internal", rec))
        else:
            raise ValueError(f"unknown classID {class_id} at offset {r.o - 4}")
    return out


def fmt(x):
    """Match the authored style: 1.0, 80.0, 0.14, .99"""
    if x == int(x):
        return f"{x:.1f}"
    return f"{x:g}"


def decompile(mech_dir, manifest=None):
    """-> [(pageName, [(key, value)])] in stream order."""
    manifest = manifest if manifest is not None else subsystems.load_manifest(MANIFEST)
    blob = record(mech_dir)
    pages = []
    for kind, r in parse(blob):
        if kind == "armor":
            kv = [
                ("BaseArmorValue", fmt(r["baseArmorValue"])),
                ("MaxArmorValue", fmt(r["maxArmorValue"])),
                ("ScaleSplashDamage", f"{r['scaleSplashDamage']:g}"),
                ("InternalDamageZone", INTERNAL_ZONE.get(r["internalDamageZone"])),
                ("ArmorZone", ARMOR_ZONE.get(r["armorZone"])),
            ]
            if r["attachedToZone"] != -1:
                kv.append(("SpecialAttachedToZone", ARMOR_ZONE.get(r["attachedToZone"])))
            pages.append((r["name"], kv))
        else:
            zone = INTERNAL_ZONE.get(r["damageZone"], "Null")
            kv = [
                ("DamageZone", zone),
                ("BaseInternalDamage", fmt(r["baseInternalDamage"])),
            ]
            for rid, pct in r["effects"]:
                path = manifest.get(rid >> 16, f"<unresolved:{rid}>")
                kv.append(("DamageEffect", f"{path},{pct:g}"))
            # no source writes GeneralDamageMode explicitly, so 0 means "omitted"
            if r["damageMode"]:
                kv.append(("DamageMode", DAMAGE_MODE.get(r["damageMode"])))
            kv.append(("ParentEntityName", r["parentEntityName"]))
            if r["damagePropagationZone"] != -1:
                kv.append(("DamagePropagationZone",
                           INTERNAL_ZONE.get(r["damagePropagationZone"])))
            if r["attachedTo"] != -1:
                kv.append(("SpecialAttachedToZone", ARMOR_ZONE.get(r["attachedTo"])))
            for key, field in (("MissileSlots", "missileSlots"),
                               ("ProjectileSlots", "projectileSlots"),
                               ("BeamSlots", "beamSlots"),
                               ("OmniSlots", "omniSlots")):
                if r[field]:
                    kv.append((key, str(r[field])))
            pages.append((zone + "Internal", kv))
    return pages


def record(mech_dir):
    ch = os.path.basename(mech_dir.rstrip("/")).lower()
    for fn in os.listdir(mech_dir):
        if fn.lower().endswith(".damage"):
            return open(os.path.join(mech_dir, fn), "rb").read()
    raise SystemExit(f"no .damage record in {mech_dir}")


def emit(pages):
    lines = []
    for name, kv in pages:
        lines.append(f"[{name}]")
        lines.extend(f"{k}={v}" for k, v in kv)
        lines.append("")
    return "\r\n".join(lines)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("mech_dir", nargs="?")
    ap.add_argument("-o", "--output")
    ap.add_argument("-m", "--manifest", default=MANIFEST)
    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_damage.py")]))
    if not args.mech_dir:
        ap.error("mech_dir required")
    text = emit(decompile(args.mech_dir, subsystems.load_manifest(args.manifest)))
    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()