#!/usr/bin/env python3
"""
gbad2_music2xm.py - Convert GBA Doom II music to FastTracker 2 .XM modules.

Usage:
    python3 gbad2_music2xm.py DoomII.gba [output_folder] [--format xm|mod|both] [--pan 0-100] [--sfx]

Writes song_00 ... song_18 as .xm and/or .mod (default: both). Opens in OpenMPT, MilkyTracker, FT2.
--format which module type to write (default both).
--pan   XM only: stereo separation in percent (default 100 = hard L/R like the GBA).
        .MOD always uses the tracker's own fixed Amiga panning.
--sfx   also dump the 98 sound effects as .wav files.

See GBA_Doom2_Music_Format.txt for the format this reads.
No third-party modules needed (plain Python 3).
"""
import os
import struct
import sys
import wave

# ---------------------------------------------------------------- ROM parsing

def find_musc(rom):
    """Locate the MUSC chunk. Prefer the Torus DPAK directory, fall back to a raw scan."""
    p = rom.find(b"DPAK")
    while p != -1:
        if rom[p + 6:p + 11] == b"Torus":
            e = p + 16
            while e + 16 <= len(rom):
                t, off, size, zero = struct.unpack_from("<4I", rom, e)
                if zero != 0 or t > 0x100:
                    break
                if t == 7 and rom[p + off:p + off + 4] == b"MUSC":
                    return p + off
                e += 16
        p = rom.find(b"DPAK", p + 1)
    p = rom.find(b"MUSC")
    while p != -1:
        song_off, bank_off, sfx_off = struct.unpack_from("<3I", rom, p + 4)
        if 0 < bank_off < song_off < sfx_off < len(rom) - p:
            return p
        p = rom.find(b"MUSC", p + 1)
    raise SystemExit("Could not find the MUSC music chunk in this ROM.")


class Sample:
    def __init__(self, index, data, finetune, volume, loop_start, loop_len, loop_count):
        self.index = index          # 1-based instrument number in the game
        self.data = data            # bytes, 8-bit signed PCM
        self.finetune = finetune    # 0..15 (MOD style)
        self.volume = volume        # 0..64
        self.loop_start = loop_start  # bytes
        self.loop_len = loop_len      # bytes (0 = no loop)
        self.loop_count = loop_count


def read_bank(rom, musc, bank_off, end_off):
    """Read a sample bank. The music bank's count field is unreliable, so the entry
    count is derived from where the next structure starts (end_off)."""
    base = musc + bank_off + 4
    count = (end_off - (bank_off + 4)) // 12
    samples = []
    for i in range(count):
        a = base + i * 12
        w0, ln, ft, vol, ls, ll = struct.unpack_from("<iHBBHH", rom, a)
        if ln == 0:
            samples.append(None)
            continue
        start = a + (w0 >> 8)
        data = rom[start:start + ln * 2]
        samples.append(Sample(i + 1, data, ft & 15, min(vol, 64), ls * 2, ll * 2, w0 & 0xFF))
    return samples


def read_songs(rom, musc, song_off):
    st = musc + song_off
    n = struct.unpack_from("<I", rom, st)[0]
    offs = struct.unpack_from("<%dI" % n, rom, st + 4)
    songs = []
    for o in offs:
        a = st + o
        length, npat, nch = rom[a + 4], rom[a + 5], rom[a + 6]
        order_offs = struct.unpack_from("<%dH" % length, rom, a + 8)
        pat_base = 8 + 2 * length
        pat_size = nch * 64 * 4
        orders = [(x - pat_base) // pat_size for x in order_offs]
        # make sure every referenced pattern is included
        npat = max([npat] + [x + 1 for x in orders])
        patterns = []
        for p in range(npat):
            pa = a + pat_base + p * pat_size
            rows = []
            for r in range(64):
                rows.append([rom[pa + (r * nch + c) * 4: pa + (r * nch + c) * 4 + 4] for c in range(nch)])
            patterns.append(rows)
        songs.append(dict(channels=nch, orders=orders, patterns=patterns))
    return songs

# ------------------------------------------------------- effect translation

def convert_effect(fx, param):
    """GBA player effect -> XM effect, keeping only what the GBA player really does."""
    if fx == 0x0:
        return (0, param) if param else (0, 0)
    if fx in (0x1, 0x2, 0xA):
        # the GBA player has no effect memory for these; XM would recall memory on 00
        return (fx, param) if param else (0, 0)
    if fx in (0x3, 0x4, 0x7, 0x9, 0xB, 0xC):
        if fx == 0xC:
            param = min(param, 64)
        return (fx, param)
    if fx == 0x5:
        return (0x5, param) if param else (0x3, 0)
    if fx == 0x6:
        return (0x6, param) if param else (0x4, 0)
    if fx == 0xD:
        # GBA uses a plain hex row number; XM expects decimal digits
        row = param if param < 64 else 0
        return (0xD, (row // 10) << 4 | (row % 10))
    if fx == 0xE:
        sub = param >> 4
        if sub in (0x1, 0x2, 0x4, 0x6, 0x7, 0x9, 0xA, 0xB, 0xC, 0xE):
            return (0xE, param)
        return (0, 0)  # E0, E3, E5, E8, ED, EF are no-ops on GBA
    if fx == 0xF:
        if 1 <= param <= 0x1F:
            return (0xF, param)
        return (0, 0)  # tempo changes are ignored by the GBA player
    return (0, 0)      # 8xx panning: no-op on GBA

# ------------------------------------------------------------- XM writing

def pad(s, n):
    s = s.encode("ascii", "replace")[:n]
    return s + b"\0" * (n - len(s))


def pack_pattern(rows, nch, ins_map, pan_vals):
    out = bytearray()
    for r in rows:
        for c in range(nch):
            cell = r[c]
            fx_hi, ins, note, param = cell[0] >> 4, cell[1], cell[2], cell[3]
            xm_note = note + 24 if note else 0
            xm_ins = ins_map.get(ins, 0) if ins else 0
            # Instrument triggers reset panning in XM, so re-apply the channel pan.
            vol = (0xC0 | pan_vals[c]) if (xm_ins and pan_vals[c] is not None) else 0
            efx, eparam = convert_effect(fx_hi, param)
            fields = [xm_note, xm_ins, vol, efx, eparam]
            mask = 0
            for bit, v in enumerate(fields):
                if v:
                    mask |= 1 << bit
            if mask == 0x1F:
                out += bytes(fields)
            else:
                out.append(0x80 | mask)
                out += bytes(v for v in fields if v)
    return bytes(out)


def delta_encode(data):
    out = bytearray(len(data))
    prev = 0
    for i, b in enumerate(data):
        out[i] = (b - prev) & 0xFF
        prev = b
    return bytes(out)


# The GBA plays at ~3557600/period Hz (Amiga PAL rate, 0.3% sharp). XM note n+24 lands
# ~10.6 cents sharp of that; fold the difference into finetune (1/128 semitone units).
GLOBAL_FT = -14


def xm_instrument(sample):
    ft = sample.finetune - 16 if sample.finetune >= 8 else sample.finetune
    ft = ft * 16 + GLOBAL_FT
    rel = 0
    while ft < -128:
        ft += 128
        rel -= 1
    data = sample.data
    ls, ll = sample.loop_start, sample.loop_len
    looped = ll > 2 and ls < len(data)
    if looped:
        ll = min(ll, len(data) - ls)
    else:
        ls = ll = 0

    name = "GBA ins %03d" % sample.index
    hdr = struct.pack("<I", 263) + pad(name, 22) + bytes([0]) + struct.pack("<H", 1)
    hdr += struct.pack("<I", 40) + bytes(96) + bytes(48) + bytes(48)
    hdr += bytes(14)                 # env counts/points/types, vibrato
    hdr += struct.pack("<H", 0)      # fadeout
    hdr += bytes(22)                 # reserved
    assert len(hdr) == 263
    sh = struct.pack("<III", len(data), ls, ll)
    sh += struct.pack("<BbBBbB", sample.volume, ft, 1 if looped else 0, 0x80, rel, 0)
    sh += pad(name, 22)
    assert len(sh) == 40
    return hdr + sh + delta_encode(data)


def write_xm(path, title, song, samples, pan_sep):
    nch = song["channels"]
    used = sorted({c[1] for p in song["patterns"] for r in p for c in r if c[1]})
    used = [i for i in used if i <= len(samples) and samples[i - 1] is not None]
    ins_map = {ins: n + 1 for n, ins in enumerate(used)}

    if pan_sep > 0:
        off = round(7.5 * pan_sep / 100.0)
        pan_vals = [max(0, min(15, (8 - off) if c % 2 == 0 else (7 + off))) for c in range(nch)]
    else:
        pan_vals = [None] * nch

    xm_channels = nch + (nch & 1)  # XM wants an even channel count
    out = bytearray()
    out += b"Extended Module: " + pad(title, 20) + b"\x1a" + pad("gbad2_music2xm", 20)
    out += struct.pack("<H", 0x0104)
    orders = song["orders"]
    out += struct.pack("<IHHHHHHHH", 276, len(orders), 0, xm_channels,
                       len(song["patterns"]), len(used), 1, 6, 125)
    out += bytes(orders) + bytes(256 - len(orders))

    for rows in song["patterns"]:
        if xm_channels != nch:
            rows = [r + [bytes(4)] for r in rows]
        packed = pack_pattern(rows, xm_channels, ins_map, pan_vals + [None])
        out += struct.pack("<IBHH", 9, 0, 64, len(packed)) + packed

    for ins in used:
        out += xm_instrument(samples[ins - 1])

    with open(path, "wb") as f:
        f.write(out)
    return len(used)


def write_mod(path, title, song, samples):
    """ProTracker .MOD. The game's own period values and samples are stored unchanged."""
    nch = song["channels"]
    used = sorted({c[1] for p in song["patterns"] for r in p for c in r if c[1]})
    used = [i for i in used if i <= len(samples) and samples[i - 1] is not None]
    if len(used) > 31:
        raise ValueError("song uses more than 31 instruments")
    ins_map = {ins: n + 1 for n, ins in enumerate(used)}

    out = bytearray(pad(title, 20))
    datas = []
    for n in range(31):
        if n < len(used):
            smp = samples[used[n] - 1]
            data = smp.data
            ln = len(data) // 2
            ls, ll = smp.loop_start, smp.loop_len
            if ll > 2 and ls < len(data):
                ll = min(ll, len(data) - ls)
                ls, ll = ls // 2, ll // 2
            else:
                ls, ll = 0, 1
            out += pad("GBA ins %03d" % smp.index, 22)
            out += struct.pack(">HBBHH", ln, smp.finetune, smp.volume, ls, ll)
            datas.append(data)
        else:
            out += bytes(22) + struct.pack(">HBBHH", 0, 0, 0, 0, 1)
    orders = song["orders"]
    out += bytes([len(orders), 127]) + bytes(orders) + bytes(128 - len(orders))
    out += {4: b"M.K.", 5: b"5CHN", 6: b"6CHN", 8: b"8CHN"}[nch]
    for rows in song["patterns"][:max(orders) + 1]:
        for r in rows:
            for cell in r:
                fx, ins, note, param = cell[0] >> 4, cell[1], cell[2], cell[3]
                period = PERIODS[note - 1] if note else 0
                mi = ins_map.get(ins, 0) if ins else 0
                efx, eparam = convert_effect(fx, param)
                out += bytes([(mi & 0xF0) | (period >> 8), period & 0xFF,
                              ((mi & 0x0F) << 4) | efx, eparam])
    for d in datas:
        out += d
    with open(path, "wb") as f:
        f.write(out)
    return len(used)


def read_periods(rom):
    """The game's own 72-entry finetune-0 period table (1712 ... 28)."""
    seq = struct.pack("<4H", 1712, 1616, 1524, 1440)
    p = rom.find(seq)
    if p == -1 or p > 0x100000:
        raise SystemExit("period table not found")
    return list(struct.unpack_from("<72H", rom, p))


def write_wav(path, data, rate):
    with wave.open(path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(1)
        w.setframerate(rate)
        w.writeframes(bytes((b + 128) & 0xFF for b in data))  # signed -> unsigned


def main(argv):
    args, opts, it = [], {}, iter(argv[1:])
    for a in it:
        if a in ("--pan", "--format"):
            opts[a] = next(it, "")
        elif a.startswith("--"):
            opts[a] = True
        else:
            args.append(a)
    pan = int(opts.get("--pan", 100))
    fmt = opts.get("--format", "both")
    if not args or fmt not in ("xm", "mod", "both"):
        print(__doc__)
        return 1
    rom = open(args[0], "rb").read()
    outdir = args[1] if len(args) > 1 else "gbad2_music"
    os.makedirs(outdir, exist_ok=True)

    musc = find_musc(rom)
    song_off, bank_off, sfx_off = struct.unpack_from("<3I", rom, musc + 4)
    print("MUSC chunk at 0x%X" % musc)
    samples = read_bank(rom, musc, bank_off, song_off)
    songs = read_songs(rom, musc, song_off)
    global PERIODS
    PERIODS = read_periods(rom)
    for i, s in enumerate(songs):
        title = "GBA Doom II song %02d" % i
        if fmt in ("xm", "both"):
            path = os.path.join(outdir, "song_%02d.xm" % i)
            n = write_xm(path, title, s, samples, pan)
            print("  %s  %d ch, %d orders, %d patterns, %d instruments"
                  % (path, s["channels"], len(s["orders"]), len(s["patterns"]), n))
        if fmt in ("mod", "both"):
            path = os.path.join(outdir, "song_%02d.mod" % i)
            write_mod(path, title, s, samples)
            print("  %s" % path)

    if "--sfx" in opts:
        sfxdir = os.path.join(outdir, "sfx")
        os.makedirs(sfxdir, exist_ok=True)
        cnt = struct.unpack_from("<I", rom, musc + sfx_off)[0]
        sfx = read_bank(rom, musc, sfx_off, sfx_off + 4 + cnt * 12)
        for s in sfx:
            if s:
                # SFX are triggered at a fixed rate; 4156 Hz is the music C-2 rate,
                # the real SFX rate is unverified, adjust if they sound off.
                write_wav(os.path.join(sfxdir, "sfx_%02d.wav" % s.index), s.data, 8000)
        print("  wrote %d sound effects to %s" % (sum(1 for s in sfx if s), sfxdir))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))