#!/usr/bin/env python

# Converts an Impulse Tracker module into a text file of music macros as
# used in the disassembly of the Zelda Oracles games.

# This is free and unencumbered software released into the public domain.

from sys import argv, stderr, version_info

def die(msg):
    if isinstance(msg, BaseException):
        msg = str(msg)
    stderr.write(str(msg) +'\n')
    exit(1)

if version_info.major < 3:
    die('python3 only!')

from collections import namedtuple
from math import floor, log2
from struct import unpack_from

if len(argv) == 2:
    inpath = argv[1]
else:
    die('Usage: {} MODULE'.format(argv[0]))

ENVELOPE_ON = 1
ENVELOPE_LOOP_ON = 2
ENVELOPE_SUSTAIN_LOOP_ON = 4

Module = namedtuple('Module', ('speed', 'tempo', 'orders', 'instruments', 'samples', 'patterns'))
Instrument = namedtuple('Instrument', ('name', 'table', 'volume_env'))
Envelope = namedtuple('Envelope', ('flags', 'points'))
Sample = namedtuple('Sample', ('name'))
Cell = namedtuple('Cell', ('note', 'ins', 'vol', 'cmd', 'cmdval'))

def read_orders(data):
    ordnum = unpack_from('H', data, 0x20)[0]
    return unpack_from('B' * ordnum, data, 0xC0)

def instrument_offsets(data):
    ordnum, insnum = unpack_from('HH', data, 0x20)
    offset = 0xC0 + ordnum
    return unpack_from('I' * insnum, data, offset)

def read_envelope(data, offset):
    flags, n = unpack_from('BB', data, offset)
    points = []
    for i in range(n):
        points.append(unpack_from('=bH', data, offset + 6 + i * 3))
    return Envelope(flags, points)

def read_instruments(data):
    offsets = instrument_offsets(data)
    instruments = []
    for offset in offsets:
        name = unpack_from('26s', data, offset + 0x20)[0].decode('ascii').strip(' \x00')
        table = []
        for i in range(120):
            table.append(unpack_from('BB', data, offset + 0x40 + i * 2))
        volume_env = read_envelope(data, offset + 0x130)
        instruments.append(Instrument(name, table, volume_env))
    return tuple(instruments)

def sample_offsets(data):
    ordnum, insnum, smpnum = unpack_from('HHH', data, 0x20)
    offset = 0xC0 + ordnum + insnum * 4
    return unpack_from('I' * smpnum, data, offset)

def read_samples(data):
    offsets = sample_offsets(data)
    samples = []
    for offset in offsets:
        name = unpack_from('26s', data, offset + 0x14)[0].decode('ascii').strip(' \x00')
        samples.append(Sample(name))
    return tuple(samples)

def pattern_offsets(data):
    ordnum, insnum, smpnum, patnum = unpack_from('HHHH', data, 0x20)
    offset = 0xC0 + ordnum + insnum * 4 + smpnum * 4
    return unpack_from('I' * patnum, data, offset)

def read_pattern(data, offset):
    _, rows = unpack_from('HH', data, offset)
    offset += 8

    prev_maskvar, prev_note, prev_ins = ([0] * 64 for i in range(3))
    prev_vol, prev_cmd, prev_cmdval = ([0] * 64 for i in range(3))
    cells = [[None for y in range(rows)] for x in range(4)]

    for row in range(rows):
        while True:
            channelvariable = unpack_from('B', data, offset)[0]
            offset += 1
            if channelvariable == 0:
                break # end of row
            channel = (channelvariable - 1) & 63
            if channelvariable & 128:
                maskvar = unpack_from('B', data, offset)[0]
                offset += 1
            else:
                maskvar = prev_maskvar[channel]
            prev_maskvar[channel] = maskvar

            if maskvar & 1:
                note = unpack_from('B', data, offset)[0]
                prev_note[channel] = note
                offset += 1
            else:
                note = None

            if maskvar & 2:
                ins = unpack_from('B', data, offset)[0]
                prev_ins[channel] = ins
                offset += 1
            else:
                ins = None

            if maskvar & 4:
                vol = unpack_from('B', data, offset)[0]
                prev_vol[channel] = vol
                offset += 1
            else:
                vol = None

            if maskvar & 8:
                cmd, cmdval = unpack_from('BB', data, offset)
                prev_cmd[channel], prev_cmdval[channel] = cmd, cmdval
                offset += 2
            else:
                cmd, cmdval = None, None

            if maskvar & 16:
                note = prev_note[channel]
            if maskvar & 32:
                ins = prev_ins[channel]
            if maskvar & 64:
                vol = prev_vol[channel]
            if maskvar & 128:
                cmd = prev_cmd[channel]
                cmdval = prev_cmdval[channel]

            cells[channel][row] = Cell(note, ins, vol, cmd, cmdval)

    return cells

def read_patterns(data):
    offsets = pattern_offsets(data)
    patterns = []
    for offset in offsets:
        pattern = read_pattern(data, offset)
        patterns.append(pattern)
    return tuple(patterns)

def read_module(filename):
    try:
        with open(filename, 'rb') as f:
            data = f.read()
    except BaseException as ex:
        die(ex)

    if data[:4].decode('ascii') != 'IMPM':
        die("Invalid IT module: '{}'".format(filename))

    speed, tempo = unpack_from('BB', data, 0x32)
    orders = read_orders(data)
    instruments = read_instruments(data)
    samples = read_samples(data)
    patterns = read_patterns(data)
    return Module(speed, tempo, orders, instruments, samples, patterns)

def row_has_break(pattern, row):
    for channel in pattern:
        if channel[row] is not None and channel[row].cmd in (2, 3):
            return True
    return False

def format_envelope(env):
    if env.flags & ENVELOPE_SUSTAIN_LOOP_ON:
        return 'env $0 $00'
    return 'env $0 $%02x' % max(1, (env.points[1][1] - 1) // 12)

def flush_queue(f, queue, ticks):
    if queue:
        while queue:
            f.write(queue.pop(0))
        f.write('%x\n' % min(0xff, ticks))
        while ticks > 0xff:
            ticks -= 0xff
            f.write('\twait1 $%x\n' % min(0xff, ticks))
    elif ticks:
        flush_queue(f, ['\twait1 $'], ticks)

def convert_channel(module, channel, filename):
    try:
        outfile = open(filename, 'w')
    except BaseException as ex:
        die(ex)

    note, vol, duty, env, ticks = None, None, None, None, 0
    queue = []

    for order in (x for x in module.orders if x != 255):
        pattern = module.patterns[order]
        for row in range(len(pattern[channel])):
            cell = pattern[channel][row]
            if cell is not None and cell.note is not None:
                flush_queue(outfile, queue, ticks)
                ticks = 0
                if cell.ins is not None:
                    smp_index = module.instruments[cell.ins - 1].table[cell.note][1] - 1
                    smp_name = module.samples[smp_index].name
                    if 'noise' not in smp_name:
                        new_duty = int(smp_name[-2:], 16)
                        if new_duty != duty:
                            duty = new_duty
                            queue.append('\tduty $%x\n' % duty)
                        new_env = format_envelope(module.instruments[cell.ins - 1].volume_env)
                        if channel < 2 and new_env != env:
                            env = new_env
                            queue.append('\t%s\n' % env)
                if channel != 2 and cell.vol is not None:
                    new_vol = round(cell.vol * 15/64)
                    if new_vol != vol:
                        vol = new_vol
                        queue.append('\tvol $%x\n' % vol)
                note = cell.note
                if cell.note & 0x80:
                    queue.append('\twait1 $')
                elif channel == 3:
                    smp_index = module.instruments[cell.ins - 1].table[cell.note][1] - 1
                    smp_name = module.samples[smp_index].name
                    queue.append('\tnote $%x $' % int(smp_name[-2:], 16))
                else:
                    queue.append('\tnote $%x $' % (cell.note - 12))
            ticks += module.speed
            if row_has_break(pattern, row):
                break

    if queue:
        flush_queue(outfile, queue, ticks)
    else:
        outfile.write('\tcmdff\n')

    outfile.close()

module = read_module(inpath)
for i in range(4):
    convert_channel(module, i, '{}.{}.txt'.format(inpath, i))
