#!/usr/bin/env python
"""
variant-to-subsystems.py - translate a MechLab-saved variant .mw4 back into the
human-readable .subsystems text format.

A variant .mw4 stores the loadout as compiled binary replication CreateMessages
(Mech::SaveMakeMessage). This decodes the {Subsystem} record:
  - each subsystem = a length-prefixed message; base message = 108 bytes
  - dataListID (Adept::ResourceID {short fileID, WORD recordID}) @ off 84
        recordID -> record name in core.mw4 (the weapon/component .data path)
  - locationID (BYTE) @ off 100  -> InternalDamageObject zone enum
  - Weapon messages append: siteName[128]@108, ejectSiteName[128]@236,
        groupIndex@364, ammoCount@368, initialAmmoCount@372, weaponFacing@376
  - SearchLight appends siteName[128]@108 ; Engine appends EngineUpgrades int @108

Usage: python variant-to-subsystems.py "Arctic Wolf 1.mw4" [...]  [--core PATH] [-o DIR]
"""
import struct, sys, os, re, argparse, glob

LOC = {0:"LeftLeg",1:"RightLeg",2:"LeftArm",3:"RightArm",4:"RightTorso",
       5:"LeftTorso",6:"CenterTorso",7:"Head",8:"Special1",9:"Special2",
       255:"Null"}
ARMOR_TYPE = {0:"Standard",1:"FerroFiberus",2:"Reactive",3:"Reflective",4:"Solarian"}
INTERNAL_TYPE = {0:"Standard",1:"EndoSteel"}
# Points-per-ton from Content\Subsystems\Armor.data; the compiled message stores
# absolute points = (.subsystems tonnage value) * pointsPerTon, so we divide to recover tonnage.
POINTS_PER_TON = {0:32.0, 1:38.0, 2:30.0, 3:30.0, 4:60.0}

def fmtf(v):
    s=f"{v:.4f}".rstrip("0").rstrip(".")
    return s if "." in s else s+".0"

# ---- LZW (gos_LZDecompress port) ----
def lz_decompress(src):
    src = bytes(src) + b"\x00\x00\x00\x00"
    bitpos = 0
    def getcode(width, mask):
        nonlocal bitpos
        i = bitpos >> 3
        v = src[i] | (src[i+1]<<8) | (src[i+2]<<16) | (src[i+3]<<24)
        v = (v >> (bitpos & 7)) & mask; bitpos += width; return v
    chain=[0]*8192; suffix=[0]*8192
    width,mask,maxindex,free = 9,511,512,258
    oldchain=oldsuffix=0; out=bytearray()
    while True:
        code=getcode(width,mask)
        if code==257: break
        if code==256:
            width,mask,maxindex,free=9,511,512,258
            code=getcode(width,mask); out.append(code&0xff)
            oldchain,oldsuffix=code,code&0xff; continue
        stack=[]
        if code>=free:
            stack.append(oldsuffix); suffix[code]=oldsuffix; chain[code]=oldchain; walk=oldchain
        else: walk=code
        while walk>=256: stack.append(suffix[walk]); walk=chain[walk]
        stack.append(walk); oldsuffix=walk&0xff; out+=bytes(reversed(stack))
        suffix[free]=oldsuffix; chain[free]=oldchain; free+=1; oldchain=code
        if free==maxindex and width!=12:
            width+=1; maxindex<<=1; mask=(mask<<1)|1
    return bytes(out)

# ---- #VBD container ----
def parse_db(path):
    data=open(path,"rb").read()
    tag,fmt,content,indexsize,nrec,nextid=struct.unpack_from("<4sIIIHH",data,0)
    assert tag==b"#VBD", f"{path}: not #VBD"
    recs={}; ids={}; off=20
    for _ in range(nrec):
        off+=8
        dlen,rlen,doff=struct.unpack_from("<III",data,off); off+=12
        rid,nlen=struct.unpack_from("<HB",data,off); off+=3
        name=data[off:off+nlen].decode("latin-1"); off+=nlen
        blob=b""
        if dlen:
            raw=data[indexsize+doff:indexsize+doff+rlen]
            blob=lz_decompress(raw) if rlen<dlen else raw
        recs[name]=blob; ids[rid]=name
    return content, recs, ids

def core_id2name(core_path):
    f=open(core_path,"rb"); hdr=f.read(20)
    _,_,_,indexsize,nrec,_=struct.unpack_from("<4sIIIHH",hdr,0)
    idx=f.read(indexsize-20); off=0; m={}
    for _ in range(nrec):
        off+=8; off+=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 propercase_map(content_root):
    """lowercased relative path -> actual-case relative path, for Model= fidelity."""
    m={}
    for dp,_,fns in os.walk(content_root):
        for fn in fns:
            full=os.path.join(dp,fn)
            rel=os.path.relpath(full,content_root)
            m[rel.lower().replace("/","\\")]=rel.replace("/","\\")
    return m

def cstr(b,o):
    e=b.find(b"\x00",o); return b[o:e if e>=0 else len(b)].decode("latin-1")

def weapon_section(path):
    p=path.lower()
    if any(x in p for x in ("laserweaponsubsystem","pulselaser","ppcweapon","flamerweapon","firehose","bombast")):
        return "Beam"
    if any(x in p for x in ("acweaponsubsystem","gaussweapon","machinegun","lbxac","ultraac","rtxac")):
        return "Ballistic"
    return "Missile"

def translate(path, core_map, case_map):
    content, recs, ids = parse_db(path)
    sub_key=[k for k in recs if k.endswith("{Subsystem}")]
    if not sub_key: raise SystemExit(f"{path}: no {{Subsystem}} record")
    blob=recs[sub_key[0]]
    chassis=sub_key[0][:-len("{Subsystem}")]
    n,=struct.unpack_from("<H",blob,0); pos=2; msgs=[]
    while pos+4<=len(blob):
        mlen,=struct.unpack_from("<I",blob,pos)
        if mlen<=0 or pos+mlen>len(blob): break
        msgs.append(blob[pos:pos+mlen]); pos+=mlen

    def model_of(b):
        fid,rid=struct.unpack_from("<hH",b,84)
        nm=core_map.get(rid,f"<rec {rid}>")
        return case_map.get(nm, nm)   # restore proper case if known
    def loc_of(b):
        return LOC.get(b[100], f"loc{b[100]}")
    def crit_of(b):
        return struct.unpack_from("<i",b,104)[0]

    out=[]; counters={"Beam":0,"Ballistic":0,"Missile":0,"HeatSink":0}
    for b in msgs:
        model=model_of(b); loc=loc_of(b); ml=model.lower(); L=len(b)
        base=os.path.basename(ml)
        if ml.endswith(".engine"):
            up,=struct.unpack_from("<i",b,108) if L>=112 else (0,)
            out.append(("Engine","NeverExecuteState",
                [("Model",os.path.basename(model)),("InternalLocation",loc),("EngineUpgrades",up)]))
        elif base=="armor.data":
            ll,rl,la,ra,lf,rf,cf,cr,hd,atype,itype=struct.unpack_from("<9f2i",b,108)
            ppt=POINTS_PER_TON.get(atype,1.0)   # convert absolute points back to tonnage
            out.append(("Armor","NeverExecuteState",[
                ("Model","Subsystems\\Armor.data"),("InternalLocation",loc),
                ("ArmorType",ARMOR_TYPE.get(atype,atype)),("InternalType",INTERNAL_TYPE.get(itype,itype)),
                ("LeftArm",fmtf(la/ppt)),("RightArm",fmtf(ra/ppt)),("LeftLeg",fmtf(ll/ppt)),("RightLeg",fmtf(rl/ppt)),
                ("LeftFrontTorso",fmtf(lf/ppt)),("RightFrontTorso",fmtf(rf/ppt)),("CenterFrontTorso",fmtf(cf/ppt)),
                ("CenterRearTorso",fmtf(cr/ppt)),("Head",fmtf(hd/ppt))]))
        elif "advancedgyro" in base:
            out.append(("AdvancedGyro","NeverExecuteState",
                [("Model","Subsystems\\AdvancedGyroSubsystem.data"),("InternalLocation",loc)]))
        elif "sensorsubsystem" in base:
            out.append(("Sensor","ActiveState",
                [("Model","Subsystems\\SensorSubsystem.data"),("InternalLocation",loc)]))
        elif ml.endswith(".torso"):
            out.append(("Torso","AlwaysExecuteState",
                [("Model",os.path.basename(model)),("InternalLocation",loc)]))
        elif "searchlight" in base:
            out.append(("SearchLight","AlwaysExecuteState",
                [("Model","Subsystems\\SearchLightSubsystem.data"),("CriticalHitsTaken",crit_of(b)),
                 ("Site",cstr(b,108)),("InternalLocation",loc)]))
        elif "heatsink" in base:
            counters["HeatSink"]+=1
            out.append((f"HeatSink{counters['HeatSink']}","NeverExecuteState",
                [("Model","Subsystems\\HeatSinkSubsystem.data"),("InternalLocation",loc)]))
        elif "jumpjet" in base:
            out.append(("JumpJet","NeverExecuteState",
                [("Model","Subsystems\\JumpJetSubsystem.data"),("CriticalHitsTaken",crit_of(b)),("InternalLocation",loc)]))
        elif "weaponsubsystems" in ml:
            site=cstr(b,108)
            grp,ammo,iammo,facing=struct.unpack_from("<iiii",b,364) if L>=380 else (0,0,0,0)
            sec=weapon_section(ml); counters[sec]+=1
            fields=[("Model",model),("CriticalHitsTaken",crit_of(b)),("Site",site),
                    ("GroupIndex",grp),("InternalLocation",loc)]
            if ammo > 0: fields.append(("AmmoCount",ammo))   # energy weapons report -1 (no ammo)
            if facing: fields.append(("WeaponFacing",facing))
            out.append((f"{sec} {counters[sec]}","NeverExecuteState",fields))
        else:
            out.append((f"Subsystem({base})","NeverExecuteState",[("Model",model),("InternalLocation",loc)]))
    return chassis, out

def render(chassis, sections):
    lines=[f"// translated from variant .mw4 (chassis: {chassis})",""]
    for name,exec_state,fields in sections:
        lines.append(f"[{name}]")
        lines.append(f"ExecutionState={exec_state}")
        for k,v in fields:
            lines.append(f"{k}={v}")
        lines.append("")
    return "\n".join(lines)

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("files",nargs="+")
    ap.add_argument("--core",default=r"Gameleap/mw4/resource/core.mw4")
    ap.add_argument("--content",default=r"Gameleap/mw4/Content")
    ap.add_argument("-o","--outdir")
    a=ap.parse_args()
    core_map=core_id2name(a.core)
    case_map=propercase_map(a.content) if os.path.isdir(a.content) else {}
    for f in a.files:
        chassis,secs=translate(f,core_map,case_map)
        text=render(chassis,secs)
        if a.outdir:
            os.makedirs(a.outdir,exist_ok=True)
            fn=os.path.join(a.outdir,os.path.splitext(os.path.basename(f))[0]+".subsystems")
            open(fn,"w").write(text); print("wrote",fn)
        else:
            print(f"\n######## {os.path.basename(f)} ########")
            print(text)

if __name__=="__main__":
    main()