#!/usr/bin/env python3
"""Extract Space Taxi's digitised speech from a C64 memory dump.

FORMAT (reverse-engineered from bankedSfxLoad at $9802, see
stuff/spacetaxi/live-level1.lst).  It is NOT 4-bit $D418 digi and it is
NOT the SID speech synth.  It is 1-BIT RUN-LENGTH audio:

  * An index table at $9B00 holds [code, ptrLo, ptrHi] triples,
    terminated by a zero code.  `code` is the PETSCII character the game
    asks for -- 'H' 'T' '!' spell "Hey Taxi!", 'P' + digit + '?' spell
    "Pad two, please", 'U' is "up".
  * Each utterance is a byte stream.  For every byte the player busy-waits
    that many passes of a fixed delay loop and then TOGGLES bit 0 of
    $D417, which flips voice 1 in and out of the filter -- a 1-bit
    speaker.  So each byte is the length of one half-cycle.
  * $FE is a raster sync (the player waits on $D011), $FF ends the
    utterance.

  The delay loop at $98A1..$98AA costs 27 cycles per pass and the
  toggle/advance/fetch tail ($98AC..$987F) another 33, hence CYCLES().

LIMITATION: a plain VICE memory dump captures the BANKED-IN view, so
$A000-$BFFF comes back as BASIC ROM ("CBMBASIC" at $A004) and most of the
speech is missing -- the player banks BASIC out ($01 = $06) precisely
because the samples live in the RAM underneath.  Dump it with the ROMs
banked out, e.g. in the VICE monitor:

    bank ram
    save "staxi-ram.bin" 0 0000 ffff

Usage:  extractSpeech.py <dump.bin> [outDir]
        The dump may have a 2-byte load address or not; both are handled.
"""

import os
import struct
import sys
import wave

NTSC_HZ    = 1022727.0     # the port already uses this (ST_SID_HZ_NUM)
OUT_RATE   = 22050
TABLE_ADDR = 0x9B00
END_MARK   = 0xFF
SYNC_MARK  = 0xFE


def cyclesFor(count):
    """Half-cycle length in CPU cycles for one stream byte."""
    return 27 * count + 33


def loadDump(path):
    """Return a 64K-addressable bytes object from a dump with or without
    a load-address header."""
    raw = open(path, "rb").read()
    if len(raw) in (65538, 65536 + 2):
        return raw[2:]
    if len(raw) == 65536:
        return raw
    # Headerless dump starting at $0801 (what mem0801.prg turned out to be).
    pad = bytearray(65536)
    pad[0x0801:0x0801 + len(raw)] = raw[:65536 - 0x0801]
    return bytes(pad)


def readTable(mem):
    """[(code, pointer)] from the $9B00 index."""
    out = []
    y = 0
    while y < 0x300:
        code = mem[TABLE_ADDR + y]
        if code == 0 or code == 0xFF:
            break
        out.append((code, mem[TABLE_ADDR + y + 1] | (mem[TABLE_ADDR + y + 2] << 8)))
        y += 3
    return out


def decode(mem, ptr, limit=0x4000):
    """Return (runs, byteCount, sawEnd). runs is [(level, cycles)]."""
    runs  = []
    level = 0
    addr  = ptr
    while addr - ptr < limit:
        b = mem[addr]
        addr += 1
        if b == END_MARK:
            return runs, addr - ptr, True
        if b == SYNC_MARK:
            runs.append((level, int(NTSC_HZ / 60.0)))
            continue
        runs.append((level, cyclesFor(b)))
        level ^= 1
    return runs, addr - ptr, False


def writeWav(runs, path):
    samples = bytearray()
    carry   = 0.0
    perOut  = NTSC_HZ / OUT_RATE
    for level, cyc in runs:
        carry += cyc
        n = int(carry / perOut)
        carry -= n * perOut
        samples += struct.pack("<h", 12000 if level else -12000) * n
    w = wave.open(path, "wb")
    w.setnchannels(1)
    w.setsampwidth(2)
    w.setframerate(OUT_RATE)
    w.writeframes(bytes(samples))
    w.close()
    return len(samples) // 2 / float(OUT_RATE)


def emitC(mem, table, romLo, outDir):
    """Write stSpeechData.c/.h: the raw run-length streams plus an index.

    The streams ship as-is (about 7.4 KB for all 15) and the game expands
    them to PCM on the fly -- rendering them at bake time would be ~38 KB
    at 5 kHz, which will not fit an IIgs bank alongside everything else.
    """
    blobs = []
    index = []
    for code, ptr in table:
        if ptr >= romLo:
            continue
        runs, nbytes, sawEnd = decode(mem, ptr)
        nxt = next((q for _, q in sorted(table, key=lambda e: e[1]) if q > ptr), None)
        if not sawEnd or (nxt is not None and nbytes > nxt - ptr):
            continue
        index.append((code, sum(len(b) for b in blobs), nbytes))
        blobs.append(bytes(mem[ptr:ptr + nbytes]))
    data = b"".join(blobs)

    h = os.path.join(outDir, "stSpeechData.h")
    c = os.path.join(outDir, "stSpeechData.c")
    with open(h, "w") as fp:
        fp.write("// Generated by assets/extractSpeech.py. Do not hand-edit.\n"
                 "//\n"
                 "// Space Taxi's digitised speech as the C64 stores it: 1-bit\n"
                 "// run-length streams. Each byte is one half-cycle length; the\n"
                 "// player toggles its output after each. $FE is a raster sync\n"
                 "// and $FF ends an utterance. stAudio.c expands these to PCM\n"
                 "// through the jlAudioPlaySfxStream fill callback.\n\n"
                 "#ifndef ST_SPEECH_DATA_H\n#define ST_SPEECH_DATA_H\n\n"
                 "#include <stdint.h>\n\n"
                 "typedef struct {\n"
                 "    uint8_t  code;      // PETSCII: 'H' 'T' '!' 'P' '1'..'9' '?' 'U'\n"
                 "    uint16_t offset;    // into kStSpeechRle\n"
                 "    uint16_t length;\n"
                 "} StSpeechEntryT;\n\n"
                 "#define ST_SPEECH_COUNT     %du\n"
                 "#define ST_SPEECH_RLE_BYTES %du\n\n"
                 "extern const StSpeechEntryT kStSpeechIndex[ST_SPEECH_COUNT];\n"
                 "extern const uint8_t        kStSpeechRle[ST_SPEECH_RLE_BYTES];\n\n"
                 "#endif\n" % (len(index), len(data)))
    with open(c, "w") as fp:
        fp.write('// Generated by assets/extractSpeech.py. Do not hand-edit.\n\n'
                 '#include "stSpeechData.h"\n\n'
                 'const StSpeechEntryT kStSpeechIndex[ST_SPEECH_COUNT] = {\n')
        for code, off, ln in index:
            fp.write("    { 0x%02Xu, %5du, %5du },   // '%s'\n"
                     % (code, off, ln, chr(code) if 33 <= code < 127 else "?"))
        fp.write("};\n\nconst uint8_t kStSpeechRle[ST_SPEECH_RLE_BYTES] = {\n")
        for i in range(0, len(data), 16):
            fp.write("    " + " ".join("0x%02X," % b for b in data[i:i + 16]) + "\n")
        fp.write("};\n")
    print("\nemitted %s (%d entries) and %s (%d bytes)" % (h, len(index), c, len(data)))


def main():
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    emit   = "--c" in sys.argv
    args   = [a for a in sys.argv[1:] if a != "--c"]
    mem    = loadDump(args[0])
    outDir = args[1] if len(args) > 1 else "speech"
    os.makedirs(outDir, exist_ok=True)

    if mem[0xA004:0xA00C] == b"CBMBASIC":
        sys.stderr.write("WARNING: $A000 holds BASIC ROM, not the RAM under it.\n"
                         "         Utterances at $A000 and above will not decode.\n"
                         "         Re-dump with `bank ram` (see the module docstring).\n\n")

    table   = readTable(mem)
    romLo   = 0xA000 if mem[0xA004:0xA00C] == b"CBMBASIC" else 0x10000
    ordered = sorted(p for _, p in table)
    good    = 0

    print("%-4s %-6s %7s %7s  %s" % ("code", "ptr", "bytes", "dur(s)", "file"))
    for code, ptr in table:
        label = chr(code) if 33 <= code < 127 else "?"
        # Utterances are stored back to back, so the next pointer is the
        # expected end. Both checks matter: BASIC ROM is full of $FF, so
        # "found a terminator" alone happily reports success on garbage --
        # it claimed all 15 decoded from a ROM-shadowed dump.
        nxt      = next((q for q in ordered if q > ptr), None)
        expected = (nxt - ptr) if nxt else None
        if ptr >= romLo:
            print("%-4s $%04X %7s %7s  -- SKIPPED (shadowed by BASIC ROM in this dump)"
                  % (label, ptr, "-", "-"))
            continue
        runs, nbytes, sawEnd = decode(mem, ptr)
        if not sawEnd:
            print("%-4s $%04X %7d %7s  -- SKIPPED (no $FF terminator)" % (label, ptr, nbytes, "-"))
            continue
        # The terminator must land BEFORE the next utterance starts.
        # (Entries are back to back but may carry a pad byte, so this is
        # <= rather than ==.) Overrunning means the $FF we found was not
        # ours -- which is exactly what happens against a ROM shadow.
        if expected is not None and nbytes > expected:
            print("%-4s $%04X %7d %7s  -- SKIPPED (ran %d bytes, next entry is %d away)"
                  % (label, ptr, nbytes, "-", nbytes, expected))
            continue
        path = os.path.join(outDir, "%02X_%s.wav" % (code, label if label != "?" else "x"))
        dur  = writeWav(runs, path)
        good += 1
        print("%-4s $%04X %7d %7.2f  %s" % (label, ptr, nbytes, dur, os.path.basename(path)))
    print("\n%d of %d utterances decoded into %s" % (good, len(table), outDir))
    if emit:
        emitC(mem, table, romLo, outDir)


if __name__ == "__main__":
    main()