#!/usr/bin/env python
"""
replay-mr.py - parse and "replay" a MechWarrior4 / FireStorm mission-review file
(the mreplay2 ".mr" format, magic 0xab9a655f, written by CMRPFull::SaveMR).

A .mr is a full network replay: session info, the player roster (names, clans,
AI scripts, team, mech/tonnage, and each player's vehicle CreateMessage), then a
chain of per-frame "packs" of network messages (incl. mech-damage events).

A true *visual* replay needs the game engine (CMRPFull::DoReplay reconstructs the
match from the recorded CreateMessages + update messages). This script does the
standalone, scriptable part: it decodes the structure and replays it as a textual
play-by-play -- match info, roster, frame timeline, and a timestamped combat log
built from the MRP_MECHDAMAGETAKEN events.

Usage:
    python replay-mr.py FILE.mr [--frames] [--commands] [--core PATH]
"""
import struct, sys, argparse, os

MAGIC      = 0xab9a655f
MAGIC_OLDA = 0xab9a655e   # g_dwMR_MetaOldA (packs have no per-frame time)

# SMRP command ids (mreplay2.h)
MRP_ADD_PLAYER_NAME, MRP_REMOVE_PLAYER_NAME = 1000, 1001
MRP_STOP_FRAME, MRP_FULLCONFIRM, MRP_RESPAWN = 2000, 2001, 2002
MRP_MECHDAMAGETAKEN = 3000
CMD_NAME = {1000:"ADD_PLAYER",1001:"REMOVE_PLAYER",2000:"STOP_FRAME",
            2001:"FULLCONFIRM",2002:"RESPAWN",3000:"MECH_DAMAGE"}

class Reader:
    def __init__(self, buf): self.b=buf; self.p=0
    def bytes(self,n):
        v=self.b[self.p:self.p+n]; self.p+=n
        if len(v)!=n: raise EOFError("short read")
        return v
    def u8(self):  return self.bytes(1)[0]
    def s8(self):  return struct.unpack("<b",self.bytes(1))[0]
    def u16(self): return struct.unpack("<H",self.bytes(2))[0]
    def i32(self): return struct.unpack("<i",self.bytes(4))[0]
    def u32(self): return struct.unpack("<I",self.bytes(4))[0]
    def f32(self): return struct.unpack("<f",self.bytes(4))[0]
    def mstr(self):                       # MString: [u32 len]; if len>0: len+1 bytes (chars+null)
        n=self.i32()
        if n<0 or n>1_000_000: raise ValueError(f"bad MString len {n} @ {self.p-4}")
        return self.bytes(n+1)[:n].decode("latin-1") if n else ""
    def resid(self):                      # Adept::ResourceID {short fileID, WORD recordID}
        fid=struct.unpack("<h",self.bytes(2))[0]; rid=self.u16(); return (fid,rid)

def parse(path):
    r=Reader(open(path,"rb").read())
    meta=r.u32()
    if meta not in (MAGIC, MAGIC_OLDA):
        raise SystemExit(f"{path}: not an mreplay2 .mr (magic={meta:#x})")
    # DPSESSIONDESC2 (80 bytes): pull a few useful fields, skip the rest
    sess=r.bytes(80)
    dwFlags, = struct.unpack_from("<I",sess,4)
    maxPlayers,curPlayers = struct.unpack_from("<II",sess,40)
    game = {"meta":meta,"dwFlags":dwFlags,"maxPlayers":maxPlayers,"curPlayers":curPlayers}
    game["gameName"]=r.mstr()
    game["myDPID"]=r.u32()
    game["player"]=r.mstr()
    game["password"]=r.mstr()
    flags=r.u8()
    game.update(starting=bool(flags&1),started=bool(flags&2),stopping=bool(flags&4),
                stopped=bool(flags&8),teamGame=bool(flags&16))
    msz=r.i32(); game["missionParam"]=r.bytes(msz)
    nbot=r.i32(); nload=r.i32(); nbcast=r.i32()
    game.update(nBOT=nbot,nLoading=nload,nBroadcast=nbcast)
    players=[]
    for _ in range(nbot+nload+nbcast):
        bot=bool(r.u8()&1)               # WriteBit + WriteByteAlign => 1 byte
        pl={"bot":bot,"name":r.mstr(),"clan":r.mstr(),"script":r.mstr(),
            "conn":r.u8(),"team":r.i32(),"pilotDecals":r.i32(),"teamDecals":r.i32(),
            "mechChassisID":r.i32(),"pilotSkins":r.s8(),"tonnage":r.i32(),
            "aiModel":r.resid()}
        vcmlen=r.i32(); pl["vcm"]=r.bytes(vcmlen)
        players.append(pl)
    game["frameTime0"]=r.f32()
    npacks=r.i32(); packs=[]
    for _ in range(npacks):
        ncount=r.i32(); ftime=0.0 if meta==MAGIC_OLDA else r.f32()
        entries=[]
        for _ in range(ncount):
            cmd=r.i32(); typ=r.i32(); ln=r.i32(); payload=r.bytes(ln)
            entries.append((cmd,typ,payload))
        packs.append((ftime,entries))
    return game, players, packs, r.p, len(r.b)

def core_id2name(core_path):
    try: f=open(core_path,"rb")
    except OSError: return {}
    hdr=f.read(20)
    if hdr[:4]!=b"#VBD": return {}
    _,_,_,indexsize,nrec,_=struct.unpack_from("<4sIIIHH",hdr,0)
    idx=f.read(indexsize-20); off=0; m={}
    for _ in range(nrec):
        off+=8+12
        rid,nlen=struct.unpack_from("<HB",idx,off); off+=3
        m[rid]=idx[off:off+nlen].decode("latin-1"); off+=nlen
    return m

def mech_of(player, core):
    """The captured VCM stub is [BYTE][WORD count][WORD modelRecordID]; the model id
    indexes a core.mw4 record (mechs\\X\\X.instance). Returns the chassis folder or None."""
    vcm=player["vcm"]
    if len(vcm)>=5:
        rid=struct.unpack_from("<H",vcm,3)[0]
        n=core.get(rid,"")
        if n.startswith("mechs\\"):
            return n.split("\\")[1]
    return None

def fmt_time(t): return f"{int(t)//60:d}:{t%60:05.2f}"

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("file"); ap.add_argument("--frames",action="store_true")
    ap.add_argument("--commands",action="store_true")
    ap.add_argument("--core",default=r"Gameleap/mw4/resource/core.mw4")
    a=ap.parse_args()
    game,players,packs,consumed,total=parse(a.file)
    core=core_id2name(a.core)
    for p in players: p["mech"]=mech_of(p,core)

    # ---- collect damage events keyed by full ReplicatorID (connectionID, localID) ----
    # id1 = victim (took damage), id2 = attacker (inflicted)
    dmg_dealt={}; dmg_taken={}; events=[]; cmd_hist={}; ids=set()
    for ftime,entries in packs:
        for cmd,typ,payload in entries:
            cmd_hist[cmd]=cmd_hist.get(cmd,0)+1
            if cmd==MRP_MECHDAMAGETAKEN and len(payload)>=12:
                dmg,con1,wid,lid1,con2,dtype,lid2=struct.unpack_from("<fBBHBBH",payload,0)
                vic=(con1,lid1); att=(con2,lid2)
                events.append((ftime,att,vic,dmg,wid,dtype)); ids|={vic,att}
                dmg_dealt[att]=dmg_dealt.get(att,0.0)+dmg
                dmg_taken[vic]=dmg_taken.get(vic,0.0)+dmg

    # ---- resolve ReplicatorID -> name ----
    # Humans pilot entity (conn, localID 2) on their own connection (reliable).
    # Bots are all replicated on the host connection -> resolved by spawn order (inferred '~').
    name_of={}
    humans=[p for p in players if not p["bot"] and p["mech"]]
    for p in humans: name_of[(p["conn"],2)] = p["name"]
    bot_ids=sorted(i for i in ids if i not in name_of)
    bots=[p for p in players if p["bot"]]              # roster (spawn) order
    inferred=set()
    for i,bid in enumerate(bot_ids):
        if i<len(bots): name_of[bid]=bots[i]["name"]+"~"; inferred.add(bid)
        else: name_of[bid]=f"AI[{bid[0]}:{bid[1]}]"
    def who(rid): return name_of.get(rid, f"{rid[0]}:{rid[1]}")

    print(f"=========== MISSION REVIEW: {os.path.basename(a.file)} ===========")
    print(f"game name : {game['gameName']}")
    print(f"host      : {game['player']}   team game: {game['teamGame']}   maxPlayers: {game['maxPlayers']}")
    print(f"players   : {len(players)}  (bots={game['nBOT']} loading={game['nLoading']} bcast={game['nBroadcast']})")
    print(f"frames    : {len(packs)}   end time: {fmt_time(packs[-1][0]) if packs else '0'}")
    print(f"parsed {consumed}/{total} bytes" + ("" if consumed==total else "  <-- WARNING: trailing bytes"))

    print("\n--- ROSTER ---")
    print(f"{'conn':>4} {'BOT':>3} {'team':>4} {'ton':>4}  {'mech':<10} {'clan':<14} name")
    for p in sorted(players,key=lambda x:(x["bot"],x["conn"])):
        mech=p["mech"] or ("(spectator)" if not p["bot"] else f"id{p['mechChassisID']}")
        print(f"{p['conn']:>4} {'Y' if p['bot'] else '-':>3} {p['team']:>4} {p['tonnage']:>4}  "
              f"{mech:<10} {p['clan'][:14]:<14} {p['name']}"
              + (f"   [AI:{os.path.basename(p['script'])}]" if p['bot'] and p['script'] else ""))

    print("\n--- DAMAGE TOTALS ---")
    print(f"{'dealt':>10} {'taken':>10}  combatant")
    for rid in sorted(ids, key=lambda r:-dmg_dealt.get(r,0)):
        print(f"{dmg_dealt.get(rid,0):>10.1f} {dmg_taken.get(rid,0):>10.1f}  {who(rid)}")

    print(f"\n--- COMBAT LOG ({len(events)} hits) ---")
    show = events if a.frames else events[:50]
    for ftime,att,vic,dmg,wid,dtype in show:
        print(f"  [{fmt_time(ftime)}] {who(att):>16} -> {who(vic):<16}  {dmg:6.1f} dmg (weaponID {wid}, type {dtype})")
    if not a.frames and len(events)>len(show):
        print(f"  ... {len(events)-len(show)} more (use --frames for all)")
    if inferred:
        print("\n  (~ = bot identity inferred from spawn order; bots replicate on the host"
              "\n   connection so their exact name<->entity link isn't stored in the file.)")

    if a.commands:
        print("\n--- COMMAND HISTOGRAM ---")
        for cmd,cnt in sorted(cmd_hist.items()):
            print(f"  {CMD_NAME.get(cmd, f'net_msg_type_{cmd}'):<18} {cnt}")

if __name__=="__main__":
    main()