#!/usr/bin/env python3
"""
psf2xm.py - convert Inspector Gadget: Gadget's Crazy Maze (PS1) ".PSF" music to FastTracker 2 .XM

The game's PSF files are NOT Portable Sound Format rips. They're the game's own sequence format,
played by a small driver in SLUS_012.67 (0x8003aa3c tick / 0x8003ae58 play):

  u32 nSongs, u32 nTracks, u32 trackSize[nTracks]
  nSongs x { s16 count, s16 trackIndex[count] }     # a song = N tracks played on N voices at once
  track data, concatenated

  Track = byte stream, run once per 60 Hz frame:
    0x80|n  wait n frames (n=0: end of track, loops if the song was started looping)
    else one command per frame, flag bits:
      0x01 u16 SPU pitch (0x1000 = 44.1 kHz)   0x02 u8 volume (0-63)
      0x04 u8 pan (0=left..15=right)            0x08 u8 sample index into SAMPLES.BIN
      0x20 key on (+0x10: u8 start offset, units of 73 ADPCM bytes)
      0x10 alone: key off

  SAMPLES.BIN: u32 n, u32 size[n], then n x (48-byte VAGp header + SPU-ADPCM data)

The music timing (rows every 7.2 frames) shows the original was an XM at speed 6 / 125 BPM, so this
rebuilds it on that grid: 1 voice -> 1 XM channel, notes recovered from pitch, volume ramps baked into
the volume column, pitch slides/vibrato approximated with E1x/E2x/X1x/X2x/1xx/2xx per row.

Usage:
  python psf2xm.py GEMZ.DAT out.xm               (reads MENUE.PSF + SAMPLES.BIN from the archive)
  python psf2xm.py --psf MENUE.PSF --bank SAMPLES.BIN out.xm
  options: --song N (default 0, the title/menu theme), --list
"""
import argparse, math, struct, sys
from collections import Counter

FRAMES_PER_ROW = 7.2          # 60 Hz driver frames per XM row at speed 6 / 125 BPM
SPEED, BPM = 6, 125
ROWS = 64

# ---------------------------------------------------------------- containers
def read_gemz(path, want):
    d = open(path, 'rb').read()
    if d[:16] != b'vision file-lib.':
        sys.exit('not a GEMZ.DAT archive')
    n = struct.unpack_from('<I', d, 16)[0]
    out = {}
    for i in range(n):
        o = 20 + i * 40
        name = d[o:o + 32].split(b'\0')[0].decode('latin1').upper()
        off, sz = struct.unpack_from('<II', d, o + 32)
        out[name] = d[off:off + sz]
    return [out[w] for w in want]

def load_psf(d):
    nsong, ntrk = struct.unpack_from('<II', d, 0)
    sizes = struct.unpack_from('<%dI' % ntrk, d, 8)
    o = 8 + 4 * ntrk
    songs = []
    for _ in range(nsong):
        n = struct.unpack_from('<h', d, o)[0]
        songs.append(list(struct.unpack_from('<%dh' % n, d, o + 2)))
        o += 2 + 2 * n
    tracks = []
    for s in sizes:
        tracks.append(d[o:o + s]); o += s
    return songs, tracks

def track_events(t):
    """[(frame, {...})] ; every non-wait command occupies one frame, like the driver."""
    i = tick = 0; out = []
    while i < len(t):
        c = t[i]; i += 1
        if c & 0x80:
            if c & 0x7f == 0:
                break
            tick += c & 0x7f; continue
        e = {}
        if c & 1: e['pitch'] = t[i] | t[i + 1] << 8; i += 2
        if c & 2: e['vol'] = t[i]; i += 1
        if c & 4: e['pan'] = t[i]; i += 1
        if c & 8: e['smp'] = t[i]; i += 1
        if c & 0x20:
            e['on'] = 1
            if c & 0x10: e['offs'] = t[i]; i += 1
        elif c & 0x10:
            e['off'] = 1
        out.append((tick, e)); tick += 1
    return out, tick

# ---------------------------------------------------------------- samples
F0 = [0, 60, 115, 98, 122]; F1 = [0, 0, -52, -55, -60]

def decode_vag(adpcm):
    pcm = []; s1 = s2 = 0; loop_start = None; loop_end = None
    for b in range(0, len(adpcm) - 15, 16):
        blk = adpcm[b:b + 16]
        sh = blk[0] & 15; f = min(blk[0] >> 4, 4); flags = blk[1]
        if flags & 4 and loop_start is None:
            loop_start = len(pcm)
        for j in range(28):
            nib = (blk[2 + j // 2] >> ((j & 1) * 4)) & 15
            if nib >= 8: nib -= 16
            v = ((nib << 12) >> sh) + ((s1 * F0[f] + s2 * F1[f] + 32) >> 6)
            v = max(-32768, min(32767, v))
            pcm.append(v); s2 = s1; s1 = v
        if flags & 1:
            if flags & 2:
                loop_end = len(pcm)
                if loop_start is None: loop_start = 0
            break
    if loop_end is None: loop_start = None
    return pcm, loop_start, loop_end

def load_bank(d):
    n = struct.unpack_from('<I', d, 0)[0]
    sizes = struct.unpack_from('<%dI' % n, d, 4)
    o = 4 + 4 * n; out = []
    for s in sizes:
        h = d[o:o + 48]
        rate = struct.unpack_from('>I', h, 16)[0]
        name = h[32:48].split(b'\0')[0].decode('latin1')
        pcm, ls, le = decode_vag(d[o + 48:o + s])
        out.append(dict(rate=rate or 22050, name=name, pcm=pcm, ls=ls, le=le))
        o += s
    return out

# ---------------------------------------------------------------- conversion
def row_start(r):
    return max(0, math.ceil(FRAMES_PER_ROW * r - 1e-9) - 1)

def semis(p):  # SPU pitch -> semitones relative to 44.1 kHz playback
    return 12 * math.log2(max(p, 1) / 4096)

def convert(song_tracks, tracks, bank):
    evs = []; length = 0
    for ti in song_tracks:
        ev, n = track_events(tracks[ti]); evs.append(ev); length = max(length, n)
    nrows = 0
    while row_start(nrows) < length: nrows += 1
    nrows = -(-nrows // ROWS) * ROWS

    # per-sample tuning: XM note 48 (C-5) = header rate shifted by the common fractional residue
    resid = {}
    for ev in evs:
        cur = 4096; smp = 0
        for _, e in ev:
            cur = e.get('pitch', cur); smp = e.get('smp', smp)
            if 'on' in e:
                x = semis(cur) - 12 * math.log2(bank[smp]['rate'] / 44100)
                resid.setdefault(smp, []).append(x - round(x))
    tune = {}
    for s, r in resid.items():
        r.sort(); tune[s] = r[len(r) // 2]
    def note_float(p, smp):
        return 48 + semis(p) - 12 * math.log2(bank[smp]['rate'] / 44100) - tune.get(smp, 0)

    chans = []
    for ev in evs:
        cells = [[0, 0, 0, 0, 0] for _ in range(nrows)]   # note, inst, volcol, eff, param
        st = dict(pitch=4096, vol=0, pan=8, smp=0)
        xm_inst = None; xm_vol = None; xm_pan = None; xm_note = None  # what the XM player has
        playing = False; pend_pan = False
        k = 0
        for r in range(nrows):
            a, b = row_start(r), row_start(r + 1)
            on = off = False; offs = None
            # events in the first 3 frames of the row define the row's starting state
            while k < len(ev) and ev[k][0] < b:
                t, e = ev[k]
                if t > a + 2 and not ('on' in e or 'off' in e):
                    break
                if 'on' in e and (on or off):
                    break   # second trigger in one row -> next row
                st.update({x: e[x] for x in ('pitch', 'vol', 'pan', 'smp') if x in e})
                if 'on' in e: on = True; offs = e.get('offs')
                if 'off' in e: off = True
                k += 1
            cell = cells[r]; effs = []
            if on:
                nf = note_float(st['pitch'], st['smp'])
                n = max(0, min(95, round(nf)))
                cell[0] = n + 1
                if xm_inst != st['smp'] + 1:
                    cell[1] = st['smp'] + 1; xm_inst = st['smp'] + 1
                    xm_pan = None
                xm_note = float(n); playing = True
                if offs:
                    effs.append((0x09, min(255, round(offs * 73 / 16 * 28 / 256))))
            elif off:
                cell[0] = 97; playing = False
            if st['vol'] != xm_vol or (on and cell[1]):
                cell[2] = 0x10 + min(64, st['vol']); xm_vol = st['vol']
            if st['pan'] != xm_pan:
                pend_pan = True
            # rest of the row: apply remaining events to state, track pitch at row end
            while k < len(ev) and ev[k][0] < b and not ('on' in ev[k][1] or 'off' in ev[k][1]):
                st.update({x: ev[k][1][x] for x in ('pitch', 'vol', 'pan', 'smp') if x in ev[k][1]})
                k += 1
            if playing and xm_note is not None:
                want = note_float(st['pitch'], st['smp'])
                d = round((want - xm_note) * 64)
                if d:
                    ad = abs(d); up = d > 0
                    if ad <= 15:
                        effs.append((33, (0x10 if up else 0x20) | ad)); step = ad
                    elif ad <= 60 and ad % 4 < 2:
                        x = round(ad / 4); effs.append((14, (0x10 if up else 0x20) | x)); step = 4 * x
                    else:
                        x = max(1, min(255, round(ad / (4 * (SPEED - 1)))))
                        effs.append((1 if up else 2, x)); step = 4 * x * (SPEED - 1)
                    xm_note += (step if up else -step) / 64
            if pend_pan and not effs:
                effs.append((8, min(255, st['pan'] * 17))); xm_pan = st['pan']; pend_pan = False
            if effs:
                cell[3], cell[4] = effs[0]
        chans.append(cells)
    return chans, nrows, tune

# ---------------------------------------------------------------- XM writer
def pack_pattern(chans, r0):
    out = bytearray()
    for r in range(r0, r0 + ROWS):
        for c in chans:
            n, i, v, e, p = c[r]
            if n and i and v and (e or p):
                out += bytes([n, i, v, e, p]); continue
            m = 0x80; body = bytearray()
            for bit, val in enumerate((n, i, v, e, p)):
                if val: m |= 1 << bit; body.append(val)
            out.append(m); out += body
    return bytes(out)

def write_xm(path, title, chans, nrows, bank, tune):
    pats = []; order = []
    for r0 in range(0, nrows, ROWS):
        p = pack_pattern(chans, r0)
        if p in pats: order.append(pats.index(p))
        else: order.append(len(pats)); pats.append(p)
    nch = len(chans) + (len(chans) & 1)
    if nch != len(chans):  # XM likes even channel counts; pad an empty one
        pats = []; order = []
        empty = [[0, 0, 0, 0, 0] for _ in range(nrows)]
        chans = chans + [empty]
        for r0 in range(0, nrows, ROWS):
            p = pack_pattern(chans, r0)
            if p in pats: order.append(pats.index(p))
            else: order.append(len(pats)); pats.append(p)
    f = bytearray()
    f += b'Extended Module: ' + title.encode()[:20].ljust(20, b' ') + b'\x1a'
    f += b'psf2xm'.ljust(20, b' ') + struct.pack('<H', 0x0104)
    f += struct.pack('<IHHHHHHHH', 276, len(order), 0, nch, len(pats), len(bank), 1, SPEED, BPM)
    f += bytes(order).ljust(256, b'\0')
    for p in pats:
        f += struct.pack('<IBHH', 9, 0, ROWS, len(p)) + p
    for idx, s in enumerate(bank):
        rate = s['rate'] * 2 ** (tune.get(idx, 0) / 12)
        tot = 12 * math.log2(rate / 8363)
        rel = round(tot); fine = max(-128, min(127, round((tot - rel) * 128)))
        name = ('%02d %s' % (idx, s['name']))[:22].encode().ljust(22, b'\0')
        f += struct.pack('<I', 263) + name + struct.pack('<BHI', 0, 1, 40)
        f += bytes(96) + bytes(48) + bytes(48) + bytes(14)
        f += struct.pack('<H', 0) + bytes(22)
        pcm = s['pcm']
        if s['ls'] is not None:
            typ = 0x11; ls = s['ls']; ll = s['le'] - s['ls']
        else:
            typ = 0x10; ls = ll = 0
        f += struct.pack('<IIIBbBBbB', len(pcm) * 2, ls * 2, ll * 2, 64, fine, typ, 128, rel, 0)
        f += name
        prev = 0; data = bytearray()
        for v in pcm:
            data += struct.pack('<h', ((v - prev + 32768) & 0xffff) - 32768); prev = v
        f += data
    open(path, 'wb').write(f)
    return len(pats), order

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('gemz', nargs='?', help='GEMZ.DAT (reads MENUE.PSF + SAMPLES.BIN from it)')
    ap.add_argument('out', nargs='?', default='menue.xm')
    ap.add_argument('--psf'); ap.add_argument('--bank')
    ap.add_argument('--song', type=int, default=0)
    ap.add_argument('--list', action='store_true')
    a = ap.parse_args()
    if a.psf and a.bank:
        if a.gemz and a.out == 'menue.xm': a.out = a.gemz
        psf = open(a.psf, 'rb').read(); bankd = open(a.bank, 'rb').read()
    elif a.gemz:
        psf, bankd = read_gemz(a.gemz, ['\\GAME\\SFX\\MENUE.PSF', '\\GAME\\SFX\\SAMPLES.BIN'])
    else:
        ap.error('give GEMZ.DAT, or --psf and --bank')
    songs, tracks = load_psf(psf)
    if a.list:
        for i, s in enumerate(songs):
            print('song %d: %d voices, %d frames' % (i, len(s), max(track_events(tracks[t])[1] for t in s)))
        return
    bank = load_bank(bankd)
    chans, nrows, tune = convert(songs[a.song], tracks, bank)
    npat, order = write_xm(a.out, 'Gadget Crazy Maze %d' % a.song, chans, nrows, bank, tune)
    print('%s: %d channels, %d rows, %d unique patterns, order %s' % (a.out, len(chans), nrows, npat, order))

if __name__ == '__main__':
    main()