#!/usr/bin/env python3
"""
mw4msg.py - reader for the GameOS "CreateMessage" streams found inside .mw4
records (.subsystems, [joint_*]{armature}, .contents, .instance, ...).

Layout, derived from the engine source rather than guessed:

  MWObject::CreateSubsystemStream / CreateArmatureStream (MWObject_Tool.cpp)
      WORD span                      number of replicator IDs consumed
      N x CreateMessage              one per [Page] that was serialised

  Each message is a plain C struct, /Zp4, built by a per-class factory in one of
  the 51 mw4/Code/MW4/*_Tool.cpp files. The common prefix is:

      off  0  u32   messageLength        Connection__Message
      off  4  u32   messageID
      off  8  u32   priority
      off 12  u32   messageFlags
      off 16  u32   classID              Replicator__CreateMessage
      off 20  u32   replicatorFlags
      off 24  u32   replicatorID
      off 28  12f   localToParent        Entity__CreateMessage (LinearMatrix4D)
      off 76  u32   executionState
      off 80  f32   initialAge
      off 84  u32   dataListID           (ResourceID; record id is the HIGH word)
      off 88  u32   alignment
      off 92  u32   nameID

  Mover__CreateMessage then adds two 24-byte Motion3D fields (96, 120), and
  MWMover__CreateMessage adds:

      off 144 u32   siteStreamResourceID
      off 148 u32   armatureStreamResourceID
      off 152 char  jointName[128]

  Verified: a joint message is exactly 280 bytes, which is 152 + 128.

localToParent is 3 rows of 4 floats; the rotation is columns 0..2 and the
translation is column 3 of each row. Checked against every joint of all 65
chassis: 1453/1453 translations matched the source .armature exactly.
"""
import struct

HDR_LEN = 0
HDR_CLASSID = 16
HDR_REPLICATORID = 24
ENT_MATRIX = 28
ENT_EXECSTATE = 76
ENT_INITIALAGE = 80
ENT_DATALISTID = 84
ENT_ALIGNMENT = 88
ENT_NAMEID = 92
MWMOVER_SITEID = 144
MWMOVER_ARMID = 148
MWMOVER_JOINTNAME = 152


def walk(data, start=2):
    """Yield (offset, messageBytes) for each message in a span-prefixed stream."""
    off = start
    while off + 4 <= len(data):
        length, = struct.unpack_from("<I", data, off)
        if length < 16 or off + length > len(data):
            raise ValueError(f"bad messageLength {length} at offset {off}")
        yield off, data[off:off + length]
        off += length
    if off != len(data):
        raise ValueError(f"trailing {len(data) - off} bytes")


def span(data):
    return struct.unpack_from("<H", data, 0)[0]


def u32(msg, off):
    return struct.unpack_from("<I", msg, off)[0]


def class_id(msg):
    return u32(msg, HDR_CLASSID)


def record_id(msg, off=ENT_DATALISTID):
    """ResourceID packs the package record id in its high word."""
    return u32(msg, off) >> 16


def matrix(msg):
    return struct.unpack_from("<12f", msg, ENT_MATRIX)


def translation(msg):
    m = matrix(msg)
    return (m[3], m[7], m[11])


def rotation3x3(msg):
    m = matrix(msg)
    return (m[0], m[1], m[2], m[4], m[5], m[6], m[8], m[9], m[10])


def joint_name(msg):
    if len(msg) <= MWMOVER_JOINTNAME:
        return None
    return msg[MWMOVER_JOINTNAME:].split(b"\0")[0].decode("latin-1")


def read_sites(data):
    """{sites} record: repeated [YawPitchRoll 3f][Point3D 3f][u32 len][name][NUL].

    Written by MWMover__CreateMessage::ConstructCreateMessage (MWMover_Tool.cpp
    ~145): `site_stream << rotation; << translation; << site_name;`
    Angles are radians.
    """
    out, off = [], 0
    while off + 28 <= len(data):
        rx, ry, rz, tx, ty, tz = struct.unpack_from("<6f", data, off)
        off += 24
        n, = struct.unpack_from("<I", data, off)
        off += 4
        name = data[off:off + n].decode("latin-1")
        off += n + 1                       # names are length-prefixed AND NUL-terminated
        out.append((name, (rx, ry, rz), (tx, ty, tz)))
    return out