#!/usr/bin/env python3
"""Convert Resident Evil (PS1) SOUND/SEPxx.HSB (VAB header + SEP + VB) sequences to Impulse Tracker (.it).

Usage: sep2it.py FILE.HSB OUTDIR [--seq N] [--name PREFIX]
"""
import struct, sys, os, math, argparse, bisect

# ---------------------------------------------------------------- HSB / VAB / SEP parsing
def vlq(d, i):
    v = 0
    while True:
        b = d[i]; i += 1; v = (v << 7) | (b & 0x7f)
        if not b & 0x80: return v, i

def load_hsb(path):
    d = open(path, 'rb').read()
    magic, ver, vid, fsize, _, ps, ts, vs = struct.unpack_from('<4sIIIHHHH', d, 0)
    assert magic == b'pBAV', 'not a VAB-headed HSB'
    vh = 32 + 2048 + ps * 512 + 512
    vbsize = fsize - vh
    # VB body sits 4 bytes later than fsize implies (checked on every SEP file: flag bytes align at +4)
    vb = d[len(d) - vbsize + 4:] + bytes(4)
    # vag sizes
    vo = 32 + 2048 + ps * 512
    sizes = [struct.unpack_from('<H', d, vo + 2 * i)[0] * 8 for i in range(256)]
    vags = {}; pos = 0
    for i in range(1, vs + 1):
        vags[i] = vb[pos:pos + sizes[i]]; pos += sizes[i]
    programs = {}
    bi = 0
    for p in range(128):
        po = 32 + 16 * p
        nt = d[po]
        if not nt: continue
        prog = dict(vol=d[po + 1], pan=d[po + 4], tones=[])
        for t in range(nt):
            o = 32 + 2048 + (bi * 16 + t) * 32
            prog['tones'].append(dict(
                prior=d[o], mode=d[o + 1], vol=d[o + 2], pan=d[o + 3], center=d[o + 4], shift=d[o + 5],
                min=d[o + 6], max=d[o + 7], pbmin=d[o + 12], pbmax=d[o + 13],
                adsr1=struct.unpack_from('<H', d, o + 16)[0], adsr2=struct.unpack_from('<H', d, o + 18)[0],
                vag=struct.unpack_from('<H', d, o + 22)[0], prog=p, idx=t))
        programs[p] = prog
        bi += 1
    # SEP
    o = vh
    assert d[o:o + 4] == b'pQES'
    seqend = len(d) - vbsize
    i = o + 6; seqs = []
    while i + 13 <= seqend:
        sid, res = struct.unpack_from('>HH', d, i)
        tempo = int.from_bytes(d[i + 4:i + 7], 'big')
        size = struct.unpack_from('>I', d, i + 9)[0]
        if size == 0: break
        seqs.append(dict(id=sid, res=res, tempo=tempo, data=d[i + 13:i + 13 + size]))
        i += 13 + size
    return programs, vags, seqs

def parse_events(data):
    """-> list of (tick, kind, ch, a, b)"""
    ev = []; t = 0; st = 0; i = 0
    while i < len(data):
        dt, i = vlq(data, i); t += dt
        b = data[i]
        if b & 0x80: st = b; i += 1
        hi, ch = st & 0xf0, st & 0x0f
        if st == 0xff:
            mt = data[i]; i += 1
            if mt == 0x2f: ev.append((t, 'end', 0, 0, 0)); break
            if mt == 0x51:
                ev.append((t, 'tempo', 0, int.from_bytes(data[i:i + 3], 'big'), 0)); i += 3; continue
            raise ValueError('meta %x' % mt)
        if hi == 0x90:
            n, v = data[i], data[i + 1]; i += 2
            ev.append((t, 'on' if v else 'off', ch, n, v))
        elif hi == 0x80: ev.append((t, 'off', ch, data[i], 0)); i += 2
        elif hi == 0xa0: i += 2
        elif hi == 0xb0: ev.append((t, 'cc', ch, data[i], data[i + 1])); i += 2
        elif hi == 0xc0: ev.append((t, 'prog', ch, data[i], 0)); i += 1
        elif hi == 0xd0: i += 1
        elif hi == 0xe0: ev.append((t, 'bend', ch, (min(data[i + 1], 127) << 7) | (data[i] & 0x7f), 0)); i += 2
        else: raise ValueError('status %x' % st)
    else:
        ev.append((t, 'end', 0, 0, 0))
    return ev

# ---------------------------------------------------------------- PS-ADPCM
F0 = [0, 60, 115, 98, 122]; F1 = [0, 0, -52, -55, -60]
def decode_vag(b):
    out = []; s1 = s2 = 0; loop_start = None; loop_end = None; looped = False
    for blk in range(len(b) // 16):
        o = blk * 16
        sh = b[o] & 0x0f; flt = min((b[o] >> 4) & 0x0f, 4); fl = b[o + 1]
        if fl & 4 and loop_start is None: loop_start = blk * 28
        for k in range(14):
            byte = b[o + 2 + k]
            for nib in (byte & 0x0f, byte >> 4):
                s = (nib << 12) & 0xffff
                if s & 0x8000: s -= 0x10000
                s = (s >> sh) + ((s1 * F0[flt] + s2 * F1[flt] + 32) >> 6)
                s = max(-32768, min(32767, s))
                out.append(s); s2, s1 = s1, s
        if fl & 1:
            loop_end = (blk + 1) * 28
            looped = bool(fl & 2)
            break
    if not looped: loop_start = None
    elif loop_start is None: loop_start = 0
    return out, loop_start, loop_end if looped else None

# ---------------------------------------------------------------- SPU ADSR simulation
def _step(rate, dec, exp, level):
    sh = (rate >> 2) - 11
    if dec: base = -8 + (rate & 3)
    else: base = 7 - (rate & 3)
    if sh < 0: cycles, step = 1, base << (-sh)
    else: cycles, step = 1 << sh, base
    if exp and not dec and level > 0x6000: cycles *= 4
    if exp and dec: step = (step * level) >> 15
    return cycles, step

def sim_phase(level, rate, dec, exp, target, tmax, pts, t0):
    t = t0
    if rate >= 0x7f and not (dec is True and target is not None and False):
        pass
    while t < tmax:
        if target is not None:
            if (not dec and level >= target) or (dec and level <= target): break
        cycles, step = _step(rate, dec, exp, level)
        if step == 0 or cycles > 44100 * 60: break
        level = max(0, min(0x7fff, level + step)); t += cycles
        pts.append((t, level))
        if level in (0, 0x7fff) and target is None: break
    return level, t

def adsr_curve(a1, a2, tmax_s=20.0):
    sr = 44100; pts = [(0, 0)]; tmax = int(tmax_s * sr)
    ar = (a1 >> 8) & 0x7f; aexp = a1 >> 15
    dr = ((a1 >> 4) & 0xf) * 4; sl = min(0x7fff, ((a1 & 0xf) + 1) * 0x800)
    sr_ = (a2 >> 6) & 0x7f; sdec = (a2 >> 14) & 1; sexp = a2 >> 15
    lvl, t = sim_phase(0, ar, False, aexp, 0x7fff, tmax, pts, 0)
    lvl, t = sim_phase(lvl, dr, True, True, sl, tmax, pts, t)
    sus_start = t
    if sr_ < 0x7f:
        lvl, t = sim_phase(lvl, sr_, bool(sdec), sexp, None, tmax, pts, t)
    return pts, lvl

def release_time(a2, start_level):
    rr = (a2 & 0x1f) * 4; rexp = (a2 >> 5) & 1
    pts = [(0, start_level)]
    lvl, t = sim_phase(start_level, rr, True, rexp, max(1, start_level // 30) if rexp else 0, 44100 * 20, pts, 0)
    return t / 44100.0

def curve_at(pts, t):
    ts = [p[0] for p in pts]
    i = bisect.bisect_right(ts, t) - 1
    if i < 0: return pts[0][1]
    if i >= len(pts) - 1: return pts[-1][1]
    (t0, l0), (t1, l1) = pts[i], pts[i + 1]
    return l0 + (l1 - l0) * (t - t0) / max(1, t1 - t0)

def rdp(points, eps):
    if len(points) < 3: return points
    (x0, y0), (x1, y1) = points[0], points[-1]
    dmax, idx = -1, 0
    for i in range(1, len(points) - 1):
        x, y = points[i]
        yy = y0 + (y1 - y0) * (x - x0) / max(1e-9, x1 - x0)
        dd = abs(y - yy)
        if dd > dmax: dmax, idx = dd, i
    if dmax > eps:
        return rdp(points[:idx + 1], eps)[:-1] + rdp(points[idx:], eps)
    return [points[0], points[-1]]

def build_envelope(a1, a2, tick_s):
    pts, end_lvl = adsr_curve(a1, a2)
    tend = pts[-1][0] / 44100.0
    nt = min(9000, max(1, int(math.ceil(tend / tick_s))))
    samp = []
    for k in range(nt + 1):
        samp.append((k, round(64 * curve_at(pts, k * tick_s * 44100) / 0x7fff)))
    eps = 0.5
    nodes = rdp(samp, eps)
    while len(nodes) > 25:
        eps *= 1.5; nodes = rdp(samp, eps)
    # de-duplicate ticks
    out = []
    for x, y in nodes:
        if out and out[-1][0] == x: out[-1] = (x, y)
        else: out.append((x, max(0, min(64, y))))
    if len(out) == 1: out.append((out[0][0] + 1, out[0][1]))
    return out, end_lvl

# ---------------------------------------------------------------- IT writer
def it_instrument(name, sample_no, env, fade, gvol):
    b = bytearray(554)
    b[0:4] = b'IMPI'
    b[0x11] = 0  # NNA cut
    struct.pack_into('<H', b, 0x14, fade)
    b[0x18] = max(0, min(128, gvol))
    b[0x19] = 0x80 | 32
    struct.pack_into('<H', b, 0x1C, 0x0214)
    b[0x1E] = 1
    nm = name.encode('ascii', 'replace')[:25]; b[0x20:0x20 + len(nm)] = nm
    for n in range(120):
        b[0x40 + 2 * n] = n; b[0x41 + 2 * n] = sample_no
    e = 0x130
    b[e] = 1 | 4  # on + sustain loop
    b[e + 1] = len(env)
    b[e + 4] = b[e + 5] = len(env) - 1
    for k, (x, y) in enumerate(env):
        b[e + 6 + 3 * k] = y; struct.pack_into('<H', b, e + 7 + 3 * k, x)
    # pan/pitch envelopes off but with 2 default nodes
    for e in (0x182, 0x1D4):
        b[e + 1] = 2; struct.pack_into('<H', b, e + 7 + 3, 100)
    return b

def it_sample_header(name, length, loop, c5, ptr):
    b = bytearray(80)
    b[0:4] = b'IMPS'
    b[0x11] = 64
    fl = 1 | 2
    if loop: fl |= 0x10
    b[0x12] = fl; b[0x13] = 64
    nm = name.encode('ascii', 'replace')[:25]; b[0x14:0x14 + len(nm)] = nm
    b[0x2E] = 1; b[0x2F] = 32
    struct.pack_into('<III', b, 0x30, length, loop[0] if loop else 0, loop[1] if loop else 0)
    struct.pack_into('<I', b, 0x3C, max(1, min(9999999, int(round(c5)))))
    struct.pack_into('<I', b, 0x48, ptr)
    return b

def pack_pattern(rows, nrows):
    """rows: dict row -> dict ch -> [note,ins,vol,cmd,param] (None = absent)"""
    out = bytearray()
    for r in range(nrows):
        for ch in sorted(rows.get(r, {})):
            c = rows[r][ch]
            mask = 0
            if c[0] is not None: mask |= 1
            if c[1] is not None: mask |= 2
            if c[2] is not None: mask |= 4
            if c[3] is not None: mask |= 8
            if not mask: continue
            out += bytes([(ch + 1) | 0x80, mask])
            if mask & 1: out.append(c[0])
            if mask & 2: out.append(c[1])
            if mask & 4: out.append(c[2])
            if mask & 8: out += bytes([c[3], c[4]])
        out.append(0)
    return struct.pack('<HH4x', len(out), nrows) + out

def write_it(path, title, nchan, init_tempo, orders, patterns, instruments, samples, mixvol=48):
    """patterns: list of bytes (packed); instruments: list of bytes; samples: list of (name, pcm16 list, loop, c5)"""
    hdr = bytearray(0xC0)
    hdr[0:4] = b'IMPM'
    t = title.encode('ascii', 'replace')[:25]; hdr[4:4 + len(t)] = t
    hdr[0x1E] = 4; hdr[0x1F] = 16
    ords = list(orders) + [255]
    struct.pack_into('<HHHHHHHH', hdr, 0x20, len(ords), len(instruments), len(samples), len(patterns), 0x0214, 0x0214, 1 | 4 | 8, 0)
    hdr[0x30] = 128; hdr[0x31] = mixvol; hdr[0x32] = 1; hdr[0x33] = init_tempo; hdr[0x34] = 128
    for c in range(64):
        hdr[0x40 + c] = 32 if c < nchan else 0x80 | 32
        hdr[0x80 + c] = 64
    head = bytes(hdr) + bytes(ords)
    off = len(head) + 4 * (len(instruments) + len(samples) + len(patterns))
    ins_off = []; blob = bytearray()
    for ins in instruments:
        ins_off.append(off + len(blob)); blob += ins
    smp_hdr_pos = []
    for s in samples:
        smp_hdr_pos.append(off + len(blob)); blob += bytes(80)
    pat_off = []
    for p in patterns:
        pat_off.append(off + len(blob)); blob += p
    # sample data
    for k, (name, pcm, loop, c5) in enumerate(samples):
        ptr = off + len(blob)
        blob += struct.pack('<%dh' % len(pcm), *pcm)
        h = it_sample_header(name, len(pcm), loop, c5, ptr)
        p = smp_hdr_pos[k] - off
        blob[p:p + 80] = h
    table = struct.pack('<%dI' % (len(ins_off) + len(smp_hdr_pos) + len(pat_off)), *(ins_off + smp_hdr_pos + pat_off))
    with open(path, 'wb') as f:
        f.write(head + table + blob)

# ---------------------------------------------------------------- conversion
FADE_COUNT = 1024  # IT fadeout counter (fadeout value subtracted per tick)

def convert(programs, vags, seq, title, max_voices=24, tail_s=2.0):
    res = seq['res']
    ev = parse_events(seq['data'])
    bpms = [60e6 / seq['tempo']] + [60e6 / e[3] for e in ev if e[1] == 'tempo' and e[3]]
    k = 1
    while k < 16 and 2 * max(bpms) * res / (48 * k) > 255: k *= 2
    def it_tempo(bpm): return max(32, min(255, int(round(bpm * res / (k * 24)))))
    tick_s = 2.5 / it_tempo(bpms[0])  # seconds per row (speed 1)
    rows_per_beat = res / k
    def row_of(t): return int(round(t / k))

    # loop markers
    loop_start = loop_end = None
    for (t, kind, ch, a, b) in ev:
        if kind == 'cc' and a == 99:
            if b == 20 and loop_start is None: loop_start = t
            elif b == 30 and loop_start is not None and loop_end is None and t > loop_start: loop_end = t
    end_tick = max(e[0] for e in ev)
    if loop_end is not None:
        ev = [e for e in ev if e[0] < loop_end]
        last_row = row_of(loop_end) - 1
    else:
        loop_start = None

    # instruments: one per (prog, tone)
    ins_map = {}; instruments = []; samples = []; smp_map = {}; ins_info = {}
    def get_ins(tone, prog):
        key = (tone['prog'], tone['idx'])
        if key in ins_map: return ins_map[key]
        skey = (tone['vag'], tone['center'], tone['shift'])
        if skey not in smp_map:
            pcm, ls, le = decode_vag(vags.get(tone['vag'], b''))
            if not pcm: pcm = [0] * 28
            c5 = 44100 * 2 ** ((60 - tone['center']) / 12 + tone['shift'] / 1536)
            samples.append(('vag%02d c%d' % (tone['vag'], tone['center']), pcm, (ls, le) if le else None, c5))
            smp_map[skey] = len(samples)
        env, sus_lvl = build_envelope(tone['adsr1'], tone['adsr2'], tick_s)
        rel = release_time(tone['adsr2'], max(sus_lvl, 0x800))
        rel_rows = max(1, rel / tick_s)
        fade = max(1, min(1024, int(round(FADE_COUNT / rel_rows))))
        gvol = int(round(128 * tone['vol'] / 127 * prog['vol'] / 127))
        instruments.append(it_instrument('p%02d t%02d' % key, smp_map[skey], env, fade, gvol))
        ins_map[key] = len(instruments)
        ins_info[ins_map[key]] = int(math.ceil(FADE_COUNT / fade)) + 1
        return ins_map[key]

    cells = {}  # (row, ch) -> [note, ins, vol, cmd, param]
    def cell(r, ch):
        c = cells.get((r, ch))
        if c is None: c = cells[(r, ch)] = [None, None, None, None, None]
        return c
    CTRL = None  # assigned after we know voice count
    ctrl_events = []  # (row, cmd, param)

    st = {ch: dict(prog=0, vol=127, expr=127, pan=64, bend=8192) for ch in range(16)}
    voices = []  # dict(ch, note, tone, on_row, off_row, free_row, pan, pitch_cur, pitch_tgt, pan_pending, vel)
    nvoices_used = 0
    def bend_units(v):
        b = (st[v['ch']]['bend'] - 8192) / 8192.0
        rng = v['tone']['pbmax'] if b >= 0 else v['tone']['pbmin']
        return int(round(b * rng * 64))
    def volcol(v):
        s = st[v['ch']]
        return max(0, min(64, int(round(64 * v['vel'] / 127 * s['vol'] / 127 * s['expr'] / 127))))
    def panval(v):
        p = st[v['ch']]['pan'] + (v['tone']['pan'] - 64) + (programs[v['tone']['prog']]['pan'] - 64)
        p = max(0, min(127, p))
        return min(255, int(round(p * 255 / 127)))
    chan_pan = {}
    def alloc(row, ch):
        nonlocal nvoices_used
        free = [i for i, v in enumerate(voices) if v is None or (v['off_row'] is not None and v['free_row'] <= row)]
        if free:
            same = [i for i in free if voices[i] is not None and voices[i]['ch'] == ch]
            return (same or free)[0]
        if len(voices) < max_voices:
            voices.append(None); return len(voices) - 1
        rel = [i for i, v in enumerate(voices) if v['off_row'] is not None]
        if rel: return min(rel, key=lambda i: voices[i]['off_row'])
        return min(range(len(voices)), key=lambda i: voices[i]['on_row'])

    # bucket events by row
    byrow = {}
    for e in ev:
        byrow.setdefault(row_of(e[0]), []).append(e)
    if loop_end is None:
        last_row = row_of(end_tick) + int(tail_s / tick_s)
    rows_total = last_row + 1
    cur_tempo = it_tempo(bpms[0])

    for r in range(rows_total):
        for (t, kind, ch, a, b) in byrow.get(r, []):
            s = st[ch]
            if kind == 'tempo' and a:
                nt = it_tempo(60e6 / a)
                if nt != cur_tempo: ctrl_events.append((r, 20, nt)); cur_tempo = nt
            elif kind == 'prog': s['prog'] = a
            elif kind == 'bend':
                s['bend'] = a
                for v in voices:
                    if v and v['ch'] == ch and v['off_row'] is None: v['pitch_tgt'] = bend_units(v)
            elif kind == 'cc':
                if a in (7, 11):
                    if a == 7: s['vol'] = b
                    else: s['expr'] = b
                    for i, v in enumerate(voices):
                        if v and v['ch'] == ch and v['off_row'] is None:
                            c = cell(r, i)
                            if c[0] is None: c[2] = volcol(v)
                elif a == 10:
                    s['pan'] = b
                    for v in voices:
                        if v and v['ch'] == ch and v['off_row'] is None: v['pan_pending'] = True
            elif kind == 'off':
                for i, v in enumerate(voices):
                    if v and v['ch'] == ch and v['note'] == a and v['off_row'] is None:
                        orow = max(r, v['on_row'] + 1)
                        v['off_row'] = orow; v['free_row'] = orow + v['rel_rows']
                        c = cell(orow, i)
                        if c[0] is None: c[0] = 246  # note fade
            elif kind == 'on':
                prog = programs.get(s['prog'])
                if not prog: continue
                for tone in prog['tones']:
                    if not (tone['min'] <= a <= tone['max']): continue
                    if tone['vag'] not in vags: continue
                    ins = get_ins(tone, prog)
                    i = alloc(r, ch)
                    v = dict(ch=ch, note=a, tone=tone, on_row=r, off_row=None, free_row=10 ** 9, vel=b,
                             pitch_cur=0, pan_pending=True, rel_rows=ins_info[ins])
                    v['pitch_tgt'] = bend_units(v)
                    voices[i] = v
                    c = cell(r, i)
                    c[0] = max(0, min(119, a)); c[1] = ins; c[2] = volcol(v); c[3] = c[4] = None
                    if chan_pan.get(i) == panval(v): v['pan_pending'] = False
        # per-row effects for voices
        for i, v in enumerate(voices):
            if not v or (v['off_row'] is not None and v['free_row'] <= r): continue
            c = cell(r, i) if (r, i) in cells else None
            if c is not None and c[3] is not None: continue
            diff = v['pitch_tgt'] - v['pitch_cur']
            if diff:
                c = cell(r, i)
                ad = abs(diff)
                if ad >= 4:
                    x = min(15, ad // 4); c[4] = 0xF0 | x; step = 4 * x
                else:
                    x = ad; c[4] = 0xE0 | x; step = x
                c[3] = 6 if diff > 0 else 5  # F up / E down
                v['pitch_cur'] += step if diff > 0 else -step
            elif v['pan_pending']:
                c = cell(r, i)
                pv = panval(v)
                c[3] = 24; c[4] = pv; chan_pan[i] = pv; v['pan_pending'] = False

    nchan = len(voices) + 1
    CTRL = len(voices)
    for (r, cmd, p) in ctrl_events:
        c = cell(r, CTRL); c[3] = cmd; c[4] = p

    # pattern segmentation
    P = int(rows_per_beat * max(1, int(192 // rows_per_beat))) if rows_per_beat <= 192 else 192
    P = max(16, min(200, P))
    bounds = []
    ls_row = row_of(loop_start) if loop_start is not None else None
    def segs(a, b):
        x = a
        while x < b:
            bounds.append((x, min(b, x + P))); x += P
    if ls_row is not None and 0 < ls_row < rows_total:
        segs(0, ls_row); loop_order = len(bounds); segs(ls_row, rows_total)
    else:
        loop_order = 0; segs(0, rows_total)
    if loop_end is not None:
        c = cell(last_row, CTRL)
        if c[3] is not None and c[3] != 2:
            # move tempo to the voice-free slot is unlikely needed; overwrite
            pass
        c[3] = 2; c[4] = loop_order
    byr = {}
    for (r, ch), c in cells.items():
        byr.setdefault(r, {})[ch] = c
    patterns = []
    for (a, b) in bounds:
        rows = {r - a: byr[r] for r in range(a, b) if r in byr}
        patterns.append(pack_pattern(rows, b - a))
    info = dict(rows=rows_total, chans=nchan, k=k, tempo=it_tempo(bpms[0]), ins=len(instruments), smp=len(samples),
                loop=loop_end is not None, seconds=rows_total * tick_s)
    return dict(nchan=nchan, tempo=it_tempo(bpms[0]), orders=list(range(len(patterns))), patterns=patterns,
                instruments=instruments, samples=samples, info=info)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('hsb'); ap.add_argument('outdir')
    ap.add_argument('--seq', type=int, default=None)
    ap.add_argument('--name', default=None)
    a = ap.parse_args()
    programs, vags, seqs = load_hsb(a.hsb)
    base = a.name or os.path.splitext(os.path.basename(a.hsb))[0]
    os.makedirs(a.outdir, exist_ok=True)
    for k, s in enumerate(seqs):
        if a.seq is not None and k != a.seq: continue
        ev = parse_events(s['data'])
        if not any(e[1] == 'on' for e in ev): continue
        r = convert(programs, vags, s, '%s seq%d' % (base, k + 1))
        out = os.path.join(a.outdir, '%s_seq%d.it' % (base, k + 1))
        write_it(out, '%s seq%d' % (base, k + 1), r['nchan'], r['tempo'], r['orders'], r['patterns'], r['instruments'], r['samples'])
        print(out, r['info'])

if __name__ == '__main__':
    main()