#!/usr/bin/env python3
# Part of JoeyLib, Copyright (c) 2026 Scott Duensing, under the MIT License.
# See LICENSE.md for license information.
# SPDX-License-Identifier: MIT

"""vgmCheck.py - checks of tools/vgm2song.py, for scripts/check-vgm.sh.

    vgmCheck.py synthetic REPO WORK
        VGM files written here, with register writes whose results are
        worked out by hand (EXPECT_*), go through vgm2song.py and songbake;
        the baked JYM1 is decoded and every row's voices must be the ones
        expected. Covers the SN76489 (tones, white noise, periodic noise on a
        free tone voice, the 447 kHz no-prescaler clock), the AY-3-8910
        (tones, mixed noise, an envelope decay followed row by row, a buzz
        sounding at the envelope's frequency), every wait command, a loop
        point, skipped commands, .vgz, --tick, the fit to 64 KB, and the
        warnings and refusals.

    vgmCheck.py pitch WAV JYM
        Furnace's own audio of a song (WAV) against the JYM1 baked from its
        VGM: the share of WINDOW-sample windows, loud enough and not
        dominated by noise, whose strongest pitch is a tone the JYM1 has
        sounding at the window's middle (within PITCH_CENTS, or an octave
        or the third harmonic of it). Prints "PITCH n/m".
"""

import gzip
import math
import os
import struct
import subprocess
import sys

VGM_RATE = 44100
FRAME = 735
FRAME_PAL = 882
WINDOW = 4096
HOP = 2048
PITCH_CENTS = 50
PITCH_BAND = (50, 5000)
PITCH_LOUD_SHARE = 0.1
HARMONICS = (1, 2, 3, 0.5)
OP_END = 0x00
OP_NOTE = 0x10
OP_OFF = 0x20
OP_TEMPO = 0x40
OP_WAIT = 0x80
VOICES = 4


# ----- JYM1 ------------------------------------------------------------------

def jymRows(path):
    """(tickMs, loop row or None, [voices after each row]) of a JYM1 file:
    a row's voices are four of (Hz or pitch, attenuation) or None."""
    data = open(path, "rb").read()
    if data[:4] != b"JYM1":
        raise SystemExit("vgmCheck: %s is no JYM1" % path)
    tick = data[4] | (data[5] << 8)
    loopOfs = data[6] | (data[7] << 8)
    events = data[8:]
    state = [None] * VOICES
    rows = []
    loopRow = None
    at = 0
    while at < len(events):
        if at == loopOfs:
            loopRow = len(rows)
        op = events[at]
        if op == OP_END:
            break
        if OP_NOTE <= op < OP_NOTE + VOICES:
            state[op - OP_NOTE] = (events[at + 1] | (events[at + 2] << 8), events[at + 3])
            at += 4
        elif OP_OFF <= op < OP_OFF + VOICES:
            state[op - OP_OFF] = None
            at += 1
        elif op == OP_TEMPO:
            at += 3
        elif op & OP_WAIT:
            rows.extend([tuple(state)] * (op & 0x7F))
            at += 1
        else:
            raise SystemExit("vgmCheck: bad JYM1 op 0x%02X" % op)
    return tick, loopRow, rows


# ----- VGM writer --------------------------------------------------------------

def vgm(commands, snClock=0, ayClock=0, loopAt=None, snFlags=0, shiftWidth=16, rate=60):
    """A VGM 1.71 file: commands (bytes) after a 0x100-byte header, with the
    loop at byte loopAt of the commands; rate is its update rate, as
    Furnace writes it."""
    header = bytearray(0x100)
    header[0:4] = b"Vgm "
    body = bytes(commands) + b"\x66"
    struct.pack_into("<I", header, 0x04, len(header) + len(body) - 4)
    struct.pack_into("<I", header, 0x08, 0x171)
    struct.pack_into("<I", header, 0x0C, snClock)
    struct.pack_into("<I", header, 0x24, rate)
    struct.pack_into("<H", header, 0x28, 0x0009)
    header[0x2A] = shiftWidth
    header[0x2B] = snFlags
    struct.pack_into("<I", header, 0x34, len(header) - 0x34)
    struct.pack_into("<I", header, 0x74, ayClock)
    if loopAt is not None:
        struct.pack_into("<I", header, 0x1C, len(header) + loopAt - 0x1C)
    return bytes(header) + body


def sn(*values):
    return b"".join(bytes((0x50, v)) for v in values)


def snTone(channel, period, atten):
    return sn(0x80 | (channel << 5) | (period & 0x0F), (period >> 4) & 0x3F, 0x90 | (channel << 5) | atten)


def ay(register, value):
    return bytes((0xA0, register, value))


def wait(samples):
    return bytes((0x61, samples & 0xFF, samples >> 8))


# ----- running the tools ------------------------------------------------------

def bake(repo, work, name, data, extra=()):
    """Write data as name.vgm, convert and bake it; (rows, stderr)."""
    vgmPath = os.path.join(work, name + ".vgm")
    songPath = os.path.join(work, name + ".song")
    jymPath = os.path.join(work, name + ".jym")
    open(vgmPath, "wb").write(data)
    conv = subprocess.run([sys.executable, os.path.join(repo, "tools/vgm2song.py")] + list(extra) + [vgmPath, songPath], capture_output=True, text=True)
    if conv.returncode != 0:
        return None, conv.stderr
    baked = subprocess.run([sys.executable, os.path.join(repo, "tools/songbake/songbake.py"), songPath, jymPath], capture_output=True, text=True)
    if baked.returncode != 0:
        return None, conv.stderr + baked.stderr
    return jymRows(jymPath), conv.stderr


def rowOf(sample, tickMs):
    return int(sample // (tickMs * VGM_RATE / 1000))


class Checks:
    def __init__(self):
        self.failures = []
        self.count = 0

    def expect(self, what, got, want):
        self.count += 1
        if got != want:
            self.failures.append("%s: got %r, want %r" % (what, got, want))


# ----- the synthetic cases -----------------------------------------------------

def snCase(repo, work, checks):
    # Frame 0: tone 0 period 254 (3579545 / (32 * 254) = 440.4 Hz) at
    # attenuation 2, and white noise at clock / 1024 (3495.6 Hz: pitch
    # 125000 / 3495.6 = 35.8, held at 31) at 5. Frame 3: tone 1 period 127
    # (880.8 Hz) at 0, tone 0 off. Frame 6 (the loop): periodic noise clocked
    # by tone 2 (period 100: 1118.6 Hz, attenuation 15, so tone voice 2 is
    # free): a 1118.6 / 16 = 69.9 Hz pulse on tone voice 2 at the noise's 4.
    # Frame 9: a Game Gear stereo write and a data block (skipped). The
    # frames come by every wait form; the song ends at frame 12.
    commands = bytearray()
    commands += snTone(0, 254, 2) + sn(0xE5, 0xF5) + bytes((0x62,))
    commands += bytes((0x62, 0x62))
    commands += snTone(1, 127, 0) + sn(0x9F) + wait(FRAME)
    commands += wait(FRAME - 16) + bytes((0x7F,))
    commands += wait(FRAME)
    loopAt = len(commands)
    commands += snTone(2, 100, 15) + sn(0xE3, 0xF4) + bytes((0x70 + 14,)) + wait(FRAME - 15)
    commands += wait(FRAME) + wait(FRAME)
    commands += bytes((0x4F, 0xFF)) + bytes((0x67, 0x66, 0x00)) + struct.pack("<I", 4) + b"\x01\x02\x03\x04"
    commands += bytes((0x62, 0x62, 0x62))
    data = vgm(commands, snClock=3579545, loopAt=loopAt)
    got, err = bake(repo, work, "sn", data)
    if got is None:
        checks.failures.append("sn: did not convert: " + err.strip())
        return
    tick, loopRow, rows = got
    checks.expect("sn: tick (60 Hz updates)", tick, 16)
    checks.expect("sn: loop row", loopRow, rowOf(6 * FRAME, 16))
    checks.expect("sn: rows", len(rows), math.ceil(12 * FRAME / (16 * VGM_RATE / 1000)))
    for frame, want in ((0, ((440, 2), None, None, (31, 5))),
                        (3, (None, (881, 0), None, (31, 5))),
                        (6, (None, (881, 0), (70, 4), None)),
                        (11, (None, (881, 0), (70, 4), None))):
        checks.expect("sn: frame %d" % frame, rows[rowOf(frame * FRAME, 16)], want)
    checks.expect("sn: the frame before 3", rows[rowOf(3 * FRAME, 16) - 1][0], (440, 2))

    # The same file compressed (.vgz) bakes the same; at --tick 33 it keeps
    # its tick and its length.
    gotGz, _ = bake(repo, work, "snGz", gzip.compress(data))
    checks.expect("sn: .vgz", gotGz, got)
    slow, _ = bake(repo, work, "snSlow", data, ("--tick", "33"))
    checks.expect("sn: --tick 33", slow[0] if slow else None, 33)
    checks.expect("sn: --tick 33 rows", len(slow[2]) if slow else None, math.ceil(12 * FRAME / (33 * VGM_RATE / 1000)))

    # The TI-99/4A's chip: 447443 Hz, no /8 prescaler: period 254 is still
    # 447443 / (4 * 254) = 440.4 Hz.
    ti, _ = bake(repo, work, "ti", vgm(snTone(0, 254, 0) + bytes((0x62,)), snClock=447443))
    checks.expect("sn: 447 kHz tone", ti[2][0][0] if ti else None, (440, 0))


def ayCase(repo, work, checks):
    # A 2 MHz chip, 50 Hz updates. Frame 0: tone A period 284 (2000000 /
    # (16 * 284) = 440.1 Hz) at volume 13 (attenuation 2); channel B noise
    # only (mixer 0x2E: tone A on, noise B on), noise period 10 (pitch 10)
    # at volume 15 (attenuation 0); channel C on the envelope, shape 0 (one
    # decay, then silence), period 2000 (a 256 ms cycle: one level every
    # 16 ms) from frame 0, tone C on. Frame 10: tone C off, the envelope
    # repeating (shape 8) at period 10 (2000000 / 2560 = 781.25 Hz): a buzz,
    # which plays on tone voice 2 at its own frequency.
    commands = bytearray()
    commands += ay(0, 284 & 0xFF) + ay(1, 284 >> 8) + ay(4, 100) + ay(5, 0) + ay(6, 10)
    commands += ay(7, 0x2A) + ay(8, 13) + ay(9, 15) + ay(10, 0x10)
    commands += ay(11, 2000 & 0xFF) + ay(12, 2000 >> 8) + ay(13, 0x00)
    commands += bytes((0x63,)) * 10
    commands += ay(7, 0x2E) + ay(11, 10) + ay(12, 0) + ay(13, 0x08)
    commands += bytes((0x63,)) * 4
    got, err = bake(repo, work, "ay", vgm(commands, ayClock=2000000, rate=50))
    if got is None:
        checks.failures.append("ay: did not convert: " + err.strip())
        return
    tick, loopRow, rows = got
    checks.expect("ay: tick (50 Hz updates)", tick, 20)
    checks.expect("ay: no loop", loopRow, None)
    checks.expect("ay: buzz warning", "envelope repeating faster" in err, True)
    # Tone C (2000000 / 1600 = 1250 Hz) follows the decay: at the end of row
    # r (20 * (r + 1) ms), the level is 15 - floor(time / 16 ms), silent
    # once the decay is over.
    for row in range(10):
        level = 15 - int(20 * (row + 1) / 16)
        want = ((440, 2), None, (1250, 15 - level) if level > 0 else None, (10, 0))
        checks.expect("ay: row %d" % row, rows[row], want)
    checks.expect("ay: buzz row", rows[11], ((440, 2), None, (781, 0), (10, 0)))


def fitCase(repo, work, checks):
    # A note every 60 Hz frame for 25000 frames: about 100 KB at 16 ms rows,
    # so the converter goes to 32 ms rows, which fit.
    commands = bytearray()
    for frame in range(25000):
        commands += snTone(0, 200 + frame % 50, 0) + bytes((0x62,))
    got, err = bake(repo, work, "fit", vgm(commands, snClock=3579545))
    checks.expect("fit: tick", got[0] if got else err.strip(), 32)
    checks.expect("fit: warning", "using 32 ms rows" in err, True)


def refusalCases(repo, work, checks):
    _, err = bake(repo, work, "unknown", vgm(bytes((0x01, 0x62)), snClock=3579545))
    checks.expect("unknown command refused", "unknown VGM command" in err, True)
    _, err = bake(repo, work, "both", vgm(snTone(0, 254, 0) + ay(8, 15) + bytes((0x62,)), snClock=3579545, ayClock=2000000))
    checks.expect("two chips need --chip", "--chip" in err, True)
    got, err = bake(repo, work, "others", vgm(snTone(0, 254, 0) + bytes((0x51, 0x10, 0x20, 0x30, 0x40)) + bytes((0x62,)), snClock=3579545))
    checks.expect("other chip warned", got is not None and "other chips" in err, True)
    got, err = bake(repo, work, "second", vgm(snTone(0, 254, 0) + bytes((0x30, 0x9F, 0x62)), snClock=3579545))
    checks.expect("second chip warned", got is not None and "second chip" in err, True)


def synthetic(repo, work):
    checks = Checks()
    snCase(repo, work, checks)
    ayCase(repo, work, checks)
    fitCase(repo, work, checks)
    refusalCases(repo, work, checks)
    for failure in checks.failures:
        print("vgmCheck: FAIL %s" % failure)
    print("SYNTHETIC %d/%d" % (checks.count - len(checks.failures), checks.count))
    return 1 if checks.failures else 0


# ----- the pitch comparison ----------------------------------------------------

def pitch(wavPath, jymPath):
    import numpy
    import wave

    w = wave.open(wavPath)
    rate = w.getframerate()
    audio = numpy.frombuffer(w.readframes(w.getnframes()), dtype=numpy.int16).astype(numpy.float64)
    if w.getnchannels() > 1:
        audio = audio.reshape(-1, w.getnchannels()).mean(axis=1)
    tick, _, rows = jymRows(jymPath)
    freqs = numpy.fft.rfftfreq(WINDOW, 1 / rate)
    band = (freqs >= PITCH_BAND[0]) & (freqs <= PITCH_BAND[1])
    shape = numpy.hanning(WINDOW)
    loud = numpy.sqrt(numpy.mean(audio ** 2)) * PITCH_LOUD_SHARE
    tested = 0
    matched = 0
    for start in range(0, len(audio) - WINDOW, HOP):
        piece = audio[start:start + WINDOW]
        if numpy.sqrt(numpy.mean(piece ** 2)) < loud:
            continue
        row = int(((start + WINDOW / 2) / rate) * 1000 // tick)
        if row >= len(rows):
            break
        state = rows[row]
        tones = [v for v in state[:3] if v is not None]
        if not tones or (state[3] is not None and state[3][1] <= min(v[1] for v in tones)):
            continue
        spectrum = numpy.abs(numpy.fft.rfft(piece * shape))
        spectrum[~band] = 0
        peak = int(numpy.argmax(spectrum))
        if peak <= 0 or peak >= len(spectrum) - 1:
            continue
        a, b, c = spectrum[peak - 1], spectrum[peak], spectrum[peak + 1]
        bend = 0.5 * (a - c) / (a - 2 * b + c) if a - 2 * b + c else 0
        heard = freqs[peak] + bend * (freqs[1] - freqs[0])
        tested += 1
        if any(abs(1200 * math.log2(heard / (hz * k))) <= PITCH_CENTS for hz, _ in tones for k in HARMONICS):
            matched += 1
    print("PITCH %d/%d" % (matched, tested))
    return 0


def main(argv):
    if len(argv) == 4 and argv[1] == "synthetic":
        return synthetic(argv[2], argv[3])
    if len(argv) == 4 and argv[1] == "pitch":
        return pitch(argv[2], argv[3])
    sys.stderr.write(__doc__)
    return 2


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