#!/usr/bin/env python3
"""rp_sf2extract.py -- extract Red Planet's AWE32 soundbanks (AUDIO1/2.RES) into
loose WAV files and regenerate RP_L4/WTPresets.cpp.

Adapted from the BattleTech tree's tools/sf2extract.py (same MUNGA engine, same
SoundFont v1.0 / SBK banks) for RP's own PRESETINFO layout.  See docs/SOUND.md.

The banks are SoundFont **v1.0** (ifil 1.0, EMU8000): `shdr` records are 16 bytes
{dwStart,dwEnd,dwStartLoop,dwEndLoop} with 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.

RP-specific facts (measured, not assumed -- see docs/SOUND.md):
  * 154 presets, 395 zones, max 4 zones in any one preset (fits samples[5]).
  * ZERO key-splits: every multi-zone preset is a pure LAYER stack whose zones
    all share one keyRange.  RP's authored content also predates
    NoteAudioControlID, so every source plays at note 60 -- key ranges are
    therefore not emitted, because nothing could ever select on them.
  * The previously shipped WAVs were all flat 44100 Hz with tuning discarded,
    and 93 presets were missing zones.

Baked into the PCM / WAV header (no engine change needed):
  samplePitch(55) / overridingRootKey(58) / coarseTune(51) / fineTune(52)
      -> the WAV's declared sample rate, so playback at note 60 reproduces the
         EMU8000 exactly:  rate = 44100 * 2^(((6000 - rootCents) + tune)/1200)
  initialAttenuation(48)  -> SBK INVERTED scale (127 = full volume), 0.375 dB
                             per step; layer balance, multiplied into the PCM
  initialFilterFc(8)/Q(9) -> the authored static resonant low-pass the EMU8000
                             applied in hardware, run over the PCM

Read from the bank and emitted to the table:
  sampleModes(54)  -> LoopAtWill / ForceStatic
  pan(17)          -> CHANNEL_LEFT / CENTER / RIGHT

NOT yet carried (needs SAMPLEINFO fields + engine work -- docs/SOUND.md F13):
  loop regions (shdr dwStartLoop/dwEndLoop) and releaseVolEnv(38) fades.

Usage:
    python tools/rp_sf2extract.py --stats     # report only, writes nothing
    python tools/rp_sf2extract.py             # extract WAVs + rewrite the table
"""
import struct, os, sys, wave, math, collections

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
AUDIO = os.path.join(ROOT, "assets", "RP411", "AUDIO")
OUT_CPP = os.path.join(ROOT, "RP_L4", "WTPresets.cpp")

BASE_RATE = 44100          # EMU8000 v1 base rate (awesfx sffile.c: v1 == 44100)
MAX_ZONES = 5              # PRESETINFO.samples[5]
MAX_PRESETS = 100          # allPresets[2][100]

G_FILTER_FC, G_FILTER_Q, G_PAN, G_RELEASE, G_INST, G_KEYRANGE = 8, 9, 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
SIGNED_GENS = {G_COARSE, G_FINE, G_RELEASE}

# zone 0 keeps the bare preset name; further layers take RP's a/b/c suffix, which
# is the convention the previously shipped assets already used.
ZONE_SUFFIX = ["", "a", "b", "c", "d"]


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", f"{path}: not an sfbk bank"
        self.d = d

        def find(fc):
            p = d.find(fc)
            assert p >= 0, f"{path}: missing {fc!r} chunk"
            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 < 400):
            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, [merged zone-generator dicts])"""
        off = self.pho + i * 38
        name = self.d[off:off + 20].split(b"\x00")[0].decode("latin1").strip()
        preset, bank, pbag = struct.unpack_from("<HHH", self.d, off + 20)
        pbag2 = self.u16(off + 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 = self.u16(self.ino + inst * 22 + 20)
            ib2 = 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)
                for pk in (pglobal, pz):           # preset layer: tunes ADD
                    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]

    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)))

    looping = (g.get(G_MODES, 0) & 3) in (1, 3)

    atten = g.get(G_ATTEN, 127)                    # SBK: 127 == full volume
    gain = 10.0 ** (-((127 - atten) * 0.375) / 20.0) if atten < 127 else 1.0

    pan = g.get(G_PAN, 64)
    chan = "CHANNEL_LEFT" if pan <= 42 else ("CHANNEL_RIGHT" if pan >= 85 else "CHANNEL_CENTER")

    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

    return dict(pcm=pcm, rate=rate, looping=looping, gain=gain, chan=chan,
                release=release, sid=g[G_SAMPLEID], nframes=en - st,
                fc=g.get(G_FILTER_FC, 127), q=g.get(G_FILTER_Q, 0),
                lstart=ls - st, lend=le - st)


def bake_lowpass(pcm, rate, fc127, q127):
    """Bake the zone's authored static resonant low-pass into the PCM.  Cutoff
    follows the AWE NRPN curve (100 + fc*7900/127 Hz); SBK Q 0..127 maps to
    0..+12 dB of biquad peak.  Designed at the zone's BAKED rate so note-60
    playback reproduces the hardware's absolute cutoff."""
    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)
    w0 = 2.0 * math.pi * cutoff_hz / rate
    alpha = math.sin(w0) / (2.0 * q_biquad)
    cw = math.cos(w0)
    b0 = b2 = (1.0 - cw) / 2.0
    b1 = 1.0 - cw
    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):
    """Whole bake pipeline in float with a single int16 conversion at the end:
    a resonance peak can overshoot full scale, so normalize the file down to fit
    rather than clipping each stage."""
    fvals = None
    if fc127 < 127:
        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:
            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)


HEADER = """#include "..\\munga_l4\\L4AUDLVL.h"

#pragma warning ( disable : 4482)

//
// GENERATED by tools/rp_sf2extract.py from assets/RP411/AUDIO/AUDIO{1,2}.RES --
// Red Planet's original AWE32 soundbanks.  Do not hand-edit; regenerate instead.
//
// Every preset now carries ALL of its instrument zones.  RP's banks are pure
// LAYER stacks (no key-splits anywhere), so each zone is a simultaneous voice.
// Per-zone tuning, layer attenuation and the authored static low-pass are baked
// into the WAVs; see docs/SOUND.md and the tool's docstring.
//
// The trailing 0 on each preset is the unused is3d field.
//

PRESETINFO allPresets[2][100] =
{
"""


def emit(table):
    lines = [HEADER]
    for bn in (1, 2):
        lines.append(f"\t//BANK {bn}\n\t{{\n")
        for preset in range(MAX_PRESETS):
            entry = table[bn].get(preset)
            if entry is None:
                lines.append("\t\t//(unused)\n\t\t{\n\t\t\t0,\n\t\t\t{\n")
                for _ in range(MAX_ZONES):
                    lines.append('\t\t\t\t{-1, 0, "", SampleChannel::CHANNEL_CENTER,'
                                 ' SampleLoop::LoopAtWill},\n')
                lines.append("\t\t\t},\n\t\t\t0\n\t\t},\n\n")
                continue
            name, zs = entry
            lines.append(f"\t\t//{name}\n\t\t//BANK {bn} PRESET {preset}\n")
            lines.append("\t\t{\n\t\t\t%d,\n\t\t\t{\n" % len(zs))
            for z in zs:
                loop = "ForceStatic" if not z["looping"] else "LoopAtWill"
                lines.append(f'\t\t\t\t{{-1, 1, "{z["file"]}", SampleChannel::{z["chan"]},'
                             f' SampleLoop::{loop}}},\n')
            for _ in range(MAX_ZONES - len(zs)):
                lines.append('\t\t\t\t{-1, 0, "", SampleChannel::CHANNEL_CENTER,'
                             ' SampleLoop::LoopAtWill},\n')
            lines.append("\t\t\t},\n\t\t\t0\n\t\t},\n\n")
        lines.append("\t},\n")
    lines.append("};\n")
    with open(OUT_CPP, "w", newline="\n") as f:
        f.write("".join(lines))


def main():
    stats = "--stats" in sys.argv
    table = {1: {}, 2: {}}
    total_zones = capped = 0
    written = []
    loopcount = filtered = attenuated = 0

    for bn, fn in ((1, "AUDIO1.RES"), (2, "AUDIO2.RES")):
        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
            if preset >= MAX_PRESETS:
                print(f"WARNING: {name} preset {preset} >= {MAX_PRESETS}, skipped")
                continue
            if len(zones) > MAX_ZONES:
                print(f"WARNING: {name} has {len(zones)} zones, capped at {MAX_ZONES}")
                capped += 1
            zs = []
            for k, g in enumerate(zones[:MAX_ZONES]):
                z = zone_fields(bank, g)
                if not z["pcm"]:
                    continue
                z["file"] = f"{name}{ZONE_SUFFIX[len(zs)]}.wav"
                zs.append(z)
            if not zs:
                continue
            for z in zs:
                loopcount += z["looping"]
                filtered += z["fc"] < 127
                attenuated += z["gain"] < 0.999
                if not stats:
                    write_wav(z["file"], z["pcm"], z["rate"], z["gain"], z["fc"], z["q"])
                written.append(z["file"])
            table[bn][preset] = (name, zs)
            total_zones += len(zs)
            multi += len(zs) > 1
        print(f"bank {bn} ({fn}): {len(table[bn])} presets, {multi} multi-zone")

    print(f"total zones: {total_zones}   looping: {loopcount}   "
          f"authored low-pass: {filtered}   attenuated layers: {attenuated}")
    if capped:
        print(f"presets capped at {MAX_ZONES} zones: {capped}")

    rates = collections.Counter(
        z["rate"] for bn in (1, 2) for _, zs in table[bn].values() for z in zs)
    print(f"distinct baked sample rates: {len(rates)}  "
          f"(most common: {', '.join(f'{r}Hz x{n}' for r, n in rates.most_common(5))})")

    if stats:
        print("\n--stats: nothing written")
        return

    emit(table)
    print(f"wrote {OUT_CPP}")

    live = set(written)
    orphans = sorted(f for f in os.listdir(AUDIO)
                     if f.lower().endswith(".wav") and f not in live)
    print(f"wrote {len(live)} wav files; {len(orphans)} pre-existing wav(s) now unreferenced")
    for o in orphans:
        print(f"   orphan: {o}")


if __name__ == "__main__":
    main()