#!/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.

FULL-ZONE extraction (AUDIO_FIDELITY.md F1/F2/F13/F14).  The banks are SoundFont
**v1.0** (ifil 1.0, EMU8000): `shdr` records are 16 bytes {dwStart,dwEnd,
dwStartLoop,dwEndLoop} (no name/rate/pitch field).  Samples are 16-bit mono PCM
in `smpl`.  Preset -> instrument resolves phdr->pbag->pgen(41); every sample-
bearing ibag zone of that instrument becomes a SAMPLEINFO slot.

Per-zone recovered metadata:
  keyRange(43)            -> keyLo/keyHi (SetupPatch selects zones by the note)
  sampleModes(54)         -> loop flag (1/3 = loop)
  samplePitch(55, SBK)    -> root*100+cents, present on every zone
  overridingRootKey(58)   -> replaces the gen-55 root when present
  coarseTune(51)/fineTune(52) -> added cents
  initialAttenuation(48)  -> SBK INVERTED scale (127=full); 0.375 dB/step [T3],
                             baked into the WAV PCM (layer balance)
  releaseVolEnv(38)       -> timecents -> seconds, for the Stop-path gain fade
  pan(17)                 -> 0..127 -> CHANNEL_LEFT/CENTER/RIGHT
  shdr loop points        -> loopStart/loopEnd frames (AL_SOFT_loop_points)

PITCH (F2, algebraically exact): EMU8000 v1 base rate is 44100 Hz; a zone's true
data rate is baked into the WAV header as
    rate = round(44100 * 2^(((6000 - effRootCents) + tuneCents) / 1200))
so the engine's note-60 pitch model (BTNotePitchFactor) reproduces original
playback at every note.  effRootCents = gen58*100 if present else gen55.

Output:
  content/AUDIO/<PresetName>[_zK].wav
  game/reconstructed/audiopresets.cpp   (allPresets[2][128], all zones)

Regenerate: python tools/sf2extract.py   (add --stats for the audit cross-check)
"""
import struct, os, sys, wave, math

HERE = os.path.dirname(__file__)
AUDIO = os.path.join(HERE, "..", "content", "AUDIO")
OUT_CPP = os.path.join(HERE, "..", "game", "reconstructed", "audiopresets.cpp")
BASE_RATE = 44100          # EMU8000 v1 (awesfx sffile.c hardcodes v1 samplerate 44100)
MAX_ZONES = 25             # PRESETINFO samples[] capacity (AllExplosion = 25 layers)

G_KEYRANGE, G_PAN, G_RELEASE, G_INST, G_KEYRANGE2 = 43, 17, 38, 41, 43
G_ATTEN, G_COARSE, G_FINE, G_SAMPLEID, G_MODES, G_PITCH, G_ROOT = 48, 51, 52, 53, 54, 55, 58
G_FILTER_FC, G_FILTER_Q = 8, 9
SIGNED_GENS = {G_COARSE, G_FINE, G_RELEASE}


def s16(v):
    return v - 0x10000 if v >= 0x8000 else v


class Bank:
    def __init__(self, path):
        d = open(path, "rb").read()
        assert d[:4] == b"RIFF" and d[8:12] == b"sfbk"
        self.d = d

        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")
        self.smpl = d[smpl_o:smpl_o + smpl_s]
        self.pho, _ = find(b"phdr"); self.pbo, _ = find(b"pbag"); self.pgo, _ = find(b"pgen")
        self.ino, _ = find(b"inst"); self.ibo, _ = find(b"ibag"); self.igo, _ = find(b"igen")
        self.sho, _ = find(b"shdr")
        self.nphdr = 0
        while d[self.pho + self.nphdr * 38:self.pho + self.nphdr * 38 + 3] != b"EOP" and self.nphdr < 200:
            self.nphdr += 1

    def u16(self, o):
        return struct.unpack_from("<H", self.d, o)[0]

    def gens(self, base, g0, g1):
        out = {}
        for g in range(g0, g1):
            op, amt = struct.unpack_from("<HH", self.d, base + g * 4)
            out[op] = s16(amt) if op in SIGNED_GENS else amt
        return out

    def preset_zones(self, i):
        """-> (name, presetNum, [zone-dict])  zone gens merged: inst-global <- inst-zone <- preset adds"""
        name = self.d[self.pho + i * 38:self.pho + i * 38 + 20].split(b"\x00")[0].decode("latin1").strip()
        preset, bank, pbag = struct.unpack_from("<HHH", self.d, self.pho + i * 38 + 20)
        pbag2 = self.u16(self.pho + (i + 1) * 38 + 24)
        zones = []
        pglobal = {}
        for b in range(pbag, pbag2):
            pz = self.gens(self.pgo, self.u16(self.pbo + b * 4), self.u16(self.pbo + (b + 1) * 4))
            if G_INST not in pz:
                pglobal = pz                      # preset-level global zone
                continue
            inst = pz[G_INST]
            ib, ib2 = self.u16(self.ino + inst * 22 + 20), self.u16(self.ino + (inst + 1) * 22 + 20)
            iglobal = {}
            for zb in range(ib, ib2):
                iz = self.gens(self.igo, self.u16(self.ibo + zb * 4), self.u16(self.ibo + (zb + 1) * 4))
                if G_SAMPLEID not in iz:
                    iglobal = iz                  # instrument-level global zone
                    continue
                g = dict(iglobal); g.update(iz)
                # preset-level: keyRange INTERSECTS, tunes/release ADD (rare in these banks)
                for pk in (pglobal, pz):
                    if G_KEYRANGE in pk and G_KEYRANGE in g:
                        alo, ahi = g[G_KEYRANGE] & 0xFF, g[G_KEYRANGE] >> 8
                        blo, bhi = pk[G_KEYRANGE] & 0xFF, pk[G_KEYRANGE] >> 8
                        g[G_KEYRANGE] = max(alo, blo) | (min(ahi, bhi) << 8)
                    elif G_KEYRANGE in pk:
                        g[G_KEYRANGE] = pk[G_KEYRANGE]
                    for add in (G_COARSE, G_FINE, G_RELEASE):
                        if add in pk:
                            g[add] = g.get(add, 0) + pk[add]
                zones.append(g)
        return name, preset, zones


def zone_fields(bank, g):
    st, en, ls, le = struct.unpack_from("<IIII", bank.d, bank.sho + g[G_SAMPLEID] * 16)
    pcm = bank.smpl[st * 2:en * 2]
    # pitch: gen55 = root*100+cents on every zone; gen58 root override
    root_cents = g.get(G_PITCH, 6000)
    if G_ROOT in g and 0 <= g[G_ROOT] <= 127:
        root_cents = g[G_ROOT] * 100
    tune = g.get(G_COARSE, 0) * 100 + g.get(G_FINE, 0)
    rate = int(round(BASE_RATE * 2.0 ** (((6000 - root_cents) + tune) / 1200.0)))
    kr = g.get(G_KEYRANGE, 0x7F00)
    keylo, keyhi = kr & 0xFF, kr >> 8
    sm = g.get(G_MODES, 0) & 3
    looping = sm in (1, 3)
    lstart = max(0, min(ls - st, en - st)) if looping else 0
    lend = max(0, min(le - st, en - st)) if looping else 0
    if looping and lend <= lstart:               # degenerate loop -> whole buffer
        lstart, lend = 0, en - st
    # SBK attenuation: INVERTED (127 = full volume), 0.375 dB/step [T3 exact curve]
    atten = g.get(G_ATTEN, 127)
    gain = 10.0 ** (-((127 - atten) * 0.375) / 20.0) if atten < 127 else 1.0
    rel_tc = g.get(G_RELEASE, None)
    release = min(8.0, 2.0 ** (rel_tc / 1200.0)) if rel_tc is not None else 0.0
    pan = g.get(G_PAN, 64)
    chan = "CHANNEL_LEFT" if pan <= 42 else ("CHANNEL_RIGHT" if pan >= 85 else "CHANNEL_CENTER")
    # (AUDIO_FIDELITY F14) the authored static resonant low-pass the EMU8000
    # applied in hardware: initialFilterFc(8) / initialFilterQ(9), SBK 0..127
    # scale.  Absent fc = filter fully open.
    filter_fc = g.get(G_FILTER_FC, 127)
    filter_q  = g.get(G_FILTER_Q, 0)
    return dict(pcm=pcm, rate=rate, keylo=keylo, keyhi=keyhi, looping=looping,
                lstart=lstart, lend=lend, gain=gain, release=release, chan=chan,
                sid=g[G_SAMPLEID], root_cents=root_cents, tune=tune,
                fc=filter_fc, q=filter_q)


def bake_lowpass(pcm, rate, fc127, q127):
    """(AUDIO_FIDELITY F14) bake the zone's authored static resonant low-pass
    into the PCM.  Cutoff uses the AWE NRPN curve (100 + fc*7900/127 Hz);
    resonance maps SBK Q 0..127 -> 0..+12 dB biquad peak [T3 -- curve shape
    approximate, off/open endpoints exact].  The filter is designed at the
    zone's BAKED rate, so note-60 playback (1 file second == 1 real second)
    reproduces the hardware's absolute cutoff exactly; pitch-shifted notes
    carry the filter with them (same limitation as the rate-baked tuning)."""
    cutoff_hz = 100.0 + fc127 * 7900.0 / 127.0
    if cutoff_hz >= 0.45 * rate:                  # at/above Nyquist headroom: no-op
        return None
    q_biquad = 0.7071 * 10.0 ** ((q127 / 127.0 * 12.0) / 20.0)
    # RBJ cookbook 2-pole low-pass, run in FLOAT (the caller converts once)
    w0 = 2.0 * math.pi * cutoff_hz / rate
    alpha = math.sin(w0) / (2.0 * q_biquad)
    cw = math.cos(w0)
    b0 = (1.0 - cw) / 2.0
    b1 = 1.0 - cw
    b2 = (1.0 - cw) / 2.0
    a0 = 1.0 + alpha
    a1 = -2.0 * cw
    a2 = 1.0 - alpha
    b0 /= a0; b1 /= a0; b2 /= a0; a1 /= a0; a2 /= a0
    n = len(pcm) // 2
    vals = struct.unpack("<%dh" % n, pcm[:n * 2])
    out = [0.0] * n
    x1 = x2 = y1 = y2 = 0.0
    for i, x0 in enumerate(vals):
        y0 = b0 * x0 + b1 * x1 + b2 * x2 - a1 * y1 - a2 * y2
        x2 = x1; x1 = x0
        y2 = y1; y1 = y0
        out[i] = y0
    return out


def write_wav(fname, pcm, rate, gain, fc127=127, q127=0):
    # Full bake pipeline in FLOAT, one int16 conversion at the end: the
    # resonance peak can overshoot full scale (165 of 232 filtered zones
    # hard-clipped in an int-per-stage draft) -- normalize the whole file
    # down to fit instead of clipping (level error <= the overshoot, wave
    # shape preserved; the EMU's filter stage had internal headroom).
    fvals = None
    if fc127 < 127:                               # authored static low-pass
        fvals = bake_lowpass(pcm, rate, fc127, q127)
    if fvals is None and gain < 0.999:
        n = len(pcm) // 2
        fvals = [float(v) for v in struct.unpack("<%dh" % n, pcm[:n * 2])]
    if fvals is not None:
        if gain < 0.999:                          # layer-balance attenuation
            fvals = [v * gain for v in fvals]
        peak = max(1.0, max(abs(v) for v in fvals))
        norm = (32767.0 / peak) if peak > 32767.0 else 1.0
        pcm = struct.pack("<%dh" % len(fvals),
                          *[int(v * norm) for v in fvals])
    with wave.open(os.path.join(AUDIO, fname), "wb") as w:
        w.setnchannels(1); w.setsampwidth(2); w.setframerate(rate)
        w.writeframes(pcm)


def main():
    stats = "--stats" in sys.argv
    banks = {1: "AUDIO1.RES", 2: "AUDIO2.RES"}
    table = {1: {}, 2: {}}
    grid_outliers = []
    total_zones = 0
    for bn, fn in banks.items():
        bank = Bank(os.path.join(AUDIO, fn))
        multi = 0
        for i in range(bank.nphdr):
            name, preset, zones = bank.preset_zones(i)
            if not name or not zones:
                continue
            zs = []
            for g in zones[:MAX_ZONES]:
                z = zone_fields(bank, g)
                if not z["pcm"]:
                    continue
                zs.append(z)
            if len(zones) > MAX_ZONES:
                print(f"WARNING: {name} has {len(zones)} zones, capped at {MAX_ZONES}")
            if not zs:
                continue
            for k, z in enumerate(zs):
                z["file"] = f"{name}.wav" if len(zs) == 1 else f"{name}_z{k}.wav"
                write_wav(z["file"], z["pcm"], z["rate"], z["gain"], z["fc"], z["q"])
                cents = 1200.0 * math.log2(z["rate"] / 44100.0)
                if abs(cents - round(cents / 100.0) * 100) > 3 and not z["tune"] % 100:
                    grid_outliers.append((name, k, z["rate"], round(cents, 1)))
            table[bn][preset] = (name, zs)
            total_zones += len(zs)
            if len(zs) > 1:
                multi += 1
        print(f"bank {bn} ({fn}): {len(table[bn])} presets, {multi} multi-zone")
    print(f"total zones: {total_zones}")
    if grid_outliers:
        print(f"rate-grid outliers (expect few): {len(grid_outliers)}")
        for o in grid_outliers[:10]:
            print("   ", o)

    if stats:
        for bn in (1, 2):
            for preset, (name, zs) in sorted(table[bn].items()):
                if len(zs) > 1:
                    print(f"b{bn} p{preset} {name}: " + " ".join(
                        f"[{z['keylo']}-{z['keyhi']} r{z['rate']} {'L' if z['looping'] else '1'}"
                        f"{' rel%.1f' % z['release'] if z['release'] else ''}]" for z in zs))
        return

    with open(OUT_CPP, "w", newline="\n") as f:
        f.write('// GENERATED by tools/sf2extract.py -- the audio soundbank table.\n')
        f.write('// FULL-ZONE extraction (AUDIO_FIDELITY.md F1/F2/F13/F14): every preset\n')
        f.write('// carries ALL its instrument zones (key-splits, layers, stereo pairs)\n')
        f.write('// with per-zone key range, loop region (frames), release seconds and\n')
        f.write('// channel; per-zone tuning + attenuation are baked into the WAVs.\n')
        f.write('// Regenerate: python tools/sf2extract.py\n')
        f.write('#include <L4AUDLVL.hpp>\n\n')
        f.write('PRESETINFO allPresets[2][128] = {};\n\n')
        f.write('namespace {\n')
        f.write('void Z(int b, int p, const char *file, int keyLo, int keyHi, SampleLoop loop,\n')
        f.write('       int loopStart, int loopEnd, float releaseSec, SampleChannel chan)\n')
        f.write('{\n')
        f.write('\tPRESETINFO &pi = allPresets[b][p];\n')
        f.write('\tif (pi.sampleNum >= MAX_PRESET_SAMPLES) return;\n')
        f.write('\tSAMPLEINFO &s = pi.samples[pi.sampleNum++];\n')
        f.write('\ts.implemented = true; s.bufferIndex = -1; s.file = file;\n')
        f.write('\ts.keyLo = keyLo; s.keyHi = keyHi; s.loop = loop;\n')
        f.write('\ts.loopStart = loopStart; s.loopEnd = loopEnd;\n')
        f.write('\ts.releaseSec = releaseSec; s.chan = chan;\n')
        f.write('}\n\n')
        f.write('struct AudioPresetInit {\n\tAudioPresetInit() {\n')
        for bn in (1, 2):
            for preset, (name, zs) in sorted(table[bn].items()):
                for z in zs:
                    loop = "LoopAtWill" if z["looping"] else "ForceStatic"
                    f.write(f'\t\tZ({bn-1},{preset},"{z["file"]}",{z["keylo"]},{z["keyhi"]},'
                            f'{loop},{z["lstart"]},{z["lend"]},{z["release"]:.3f}f,{z["chan"]});\n')
        f.write('\t}\n} s_audioPresetInit;\n}\n')
    npresets = len(table[1]) + len(table[2])
    print(f"wrote {OUT_CPP} ({npresets} presets, {total_zones} zones)")


if __name__ == "__main__":
    main()