#!/usr/bin/env python3
"""sf2extract.py -- extract the BattleTech AWE32 SoundFonts (AUDIO1/2.RES) into
loose WAV files + generate the allPresets[] C++ table for the WinTesla port.

The banks are SoundFont **v1.0** (ifil 1.0, EMU8000): the `shdr` sample headers are
16 bytes {dwStart,dwEnd,dwStartLoop,dwEndLoop} (NO name/rate field -- the v2.0
46-byte layout does not apply).  Samples are 16-bit mono PCM in `smpl`.  Preset ->
sample resolves through phdr->pbag->pgen(inst=41)->inst->ibag->igen(sampleID=53)->shdr.

Output:
  content/AUDIO/<PresetName>.wav   (one per preset; the L4AUDRES loader sf_opens these)
  game/reconstructed/audiopresets.cpp  (the allPresets[2][100] table, replaces the
                                        btstubs.cpp zero-init stub)

Sample rate: SF2 v1.0 shdr has no per-sample rate; these were authored at 22050 Hz
(SAMPLE_RATE below -- adjust if pitch is wrong).  Loop is taken from the shdr loop
points (a real loop region -> LoopAtWill; else ForceStatic one-shot).
"""
import struct, os, wave

HERE = os.path.dirname(__file__)
AUDIO = os.path.join(HERE, "..", "content", "AUDIO")
OUT_CPP = os.path.join(HERE, "..", "game", "reconstructed", "audiopresets.cpp")
SAMPLE_RATE = 22050

def parse_bank(path):
    d = open(path, "rb").read()
    assert d[:4]==b"RIFF" and d[8:12]==b"sfbk"
    def find(fc): p=d.find(fc); return p+8, struct.unpack_from("<I",d,p+4)[0]
    smpl_o, smpl_s = find(b"smpl"); smpl = d[smpl_o:smpl_o+smpl_s]
    pho,_=find(b"phdr"); pbo,_=find(b"pbag"); pgo,_=find(b"pgen")
    ino,_=find(b"inst"); ibo,_=find(b"ibag"); igo,_=find(b"igen")
    sho,shs=find(b"shdr")                       # 16-byte v1.0 records
    def u16(o): return struct.unpack_from("<H",d,o)[0]
    def genlk(base,g0,g1,oper):
        for g in range(g0,g1):
            op,amt=struct.unpack_from("<HH",d,base+g*4)
            if op==oper: return amt
        return None
    nphdr = 0
    while d[pho+nphdr*38:pho+nphdr*38+3] != b"EOP" and nphdr < 200:
        nphdr += 1
    out = []
    for i in range(nphdr):
        name=d[pho+i*38:pho+i*38+20].split(b"\x00")[0].decode("latin1").strip()
        preset,bank,pbag=struct.unpack_from("<HHH",d,pho+i*38+20)
        pbag2=struct.unpack_from("<H",d,pho+(i+1)*38+24)[0]
        instId=None
        for b in range(pbag,pbag2):
            instId=genlk(pgo,u16(pbo+b*4),u16(pbo+(b+1)*4),41)
            if instId is not None: break
        sampId=None; sampleModes=0
        if instId is not None:
            ib=u16(ino+instId*22+20); ib2=u16(ino+(instId+1)*22+20)
            for b in range(ib,ib2):
                g0=u16(ibo+b*4); g1=u16(ibo+(b+1)*4)
                sid=genlk(igo,g0,g1,53)                  # oper 53 = sampleID
                if sid is not None:
                    sampId=sid
                    sm=genlk(igo,g0,g1,54)               # oper 54 = sampleModes
                    if sm is not None: sampleModes=sm
                    break
        if sampId is None: continue
        st,en,ls,le=struct.unpack_from("<IIII", d, sho+sampId*16)
        # SF2 sampleModes: 0/2 = no loop (one-shot), 1 = loop, 3 = loop-until-release.
        # (Presence of shdr loop points does NOT mean loop -- almost every sample has
        # them; only sampleModes decides.  The old loop-points heuristic looped ~all
        # 241 samples, so weapon fire / explosions / steps played forever.)
        looping = (sampleModes == 1 or sampleModes == 3)
        out.append((preset, name, smpl[st*2:en*2], looping))
    return out

def write_wav(name, pcm):
    with wave.open(os.path.join(AUDIO, name+".wav"), "wb") as w:
        w.setnchannels(1); w.setsampwidth(2); w.setframerate(SAMPLE_RATE)
        w.writeframes(pcm)

def main():
    banks = {1:"AUDIO1.RES", 2:"AUDIO2.RES"}
    table = {1:{}, 2:{}}
    for bn, fn in banks.items():
        presets = parse_bank(os.path.join(AUDIO, fn))
        for (preset, name, pcm, looping) in presets:
            if not name or not pcm: continue
            write_wav(name, pcm)
            table[bn][preset] = (name, looping)
        print(f"bank {bn} ({fn}): {len(table[bn])} presets extracted to WAV")
    # emit allPresets[2][100]
    with open(OUT_CPP, "w", newline="\n") as f:
        f.write('// GENERATED by scratchpad/sf2extract.py -- the audio soundbank table.\n')
        f.write('// Replaces the zero-init allPresets stub (btstubs.cpp) with the real\n')
        f.write('// bank/preset -> sample map recovered from the AWE32 SoundFonts\n')
        f.write('// (AUDIO1/2.RES, SF2 v1.0).  Each preset -> a loose <Name>.wav the\n')
        f.write('// L4AUDRES loader sf_opens.  Regenerate: python scratchpad/sf2extract.py\n')
        f.write('#include <L4AUDLVL.hpp>\n\n')
        f.write('PRESETINFO allPresets[2][128] = {};\n\n')
        f.write('namespace {\nstruct AudioPresetInit {\n\tAudioPresetInit() {\n')
        for bn in (1,2):
            for preset, (name, looping) in sorted(table[bn].items()):
                loop = "SampleLoop::LoopAtWill" if looping else "SampleLoop::ForceStatic"
                f.write(f'\t\t{{ SAMPLEINFO &s = allPresets[{bn-1}][{preset}].samples[0];\n')
                f.write(f'\t\t  allPresets[{bn-1}][{preset}].sampleNum = 1;\n')
                f.write(f'\t\t  s.implemented = true; s.bufferIndex = -1; s.file = "{name}.wav";\n')
                f.write(f'\t\t  s.chan = SampleChannel::CHANNEL_CENTER; s.loop = {loop}; }}\n')
        f.write('\t}\n} s_audioPresetInit;\n}\n')
    total = len(table[1]) + len(table[2])
    print(f"wrote {OUT_CPP} ({total} presets)")

if __name__ == "__main__":
    main()