#!/usr/bin/env python3
"""
ff_extract.py - Extract Havoc ".FF" (FlashFile) archives.

The .FF format (used by MUSIC.FF, SOUND.FF, OBJECTS.FF, INTRFACE.FF):

    uint32  file_count
    repeat file_count times:
        uint32  absolute_offset_into_file
        char[]  name, NUL-terminated
    <concatenated file data; each entry's size = next_offset - this_offset>

MUSIC.FF additionally contains:
    WAVE<grp><n>.DAT   RIFF WAVE PCM segments (mono 22050Hz 8-bit, typically)
    PLST<grp>.DAT      playlists that sequence the WAVE<grp>* segments into a song:
        uint16  segment_count          (how many WAVE segments in the group, 1-based)
        uint16  step_count
        uint16[step_count]  sequence    (1-based indices into the group's segments)

Usage:
    python ff_extract.py MUSIC.FF -o out_music              # dump every entry raw
    python ff_extract.py MUSIC.FF -o out_music --songs      # also assemble songs from playlists
    python ff_extract.py MUSIC.FF -o out_music --songs --convert ogg   # needs ffmpeg on PATH
"""
import argparse, os, struct, shutil, subprocess
from collections import defaultdict


def parse_ff(data):
    """Return list of (name, offset, size)."""
    (count,) = struct.unpack_from("<I", data, 0)
    pos = 4
    raw = []
    for _ in range(count):
        (off,) = struct.unpack_from("<I", data, pos); pos += 4
        end = data.index(b"\x00", pos)
        name = data[pos:end].decode("latin1")
        pos = end + 1
        raw.append((name, off))
    out = []
    for i, (name, off) in enumerate(raw):
        nxt = raw[i + 1][1] if i + 1 < len(raw) else len(data)
        out.append((name, off, nxt - off))
    return out


def wav_parts(buf):
    """Return (fmt_chunk_bytes, data_bytes) from a RIFF WAVE buffer."""
    assert buf[:4] == b"RIFF" and buf[8:12] == b"WAVE", "not a RIFF WAVE"
    pos = 12
    fmt = data = None
    while pos + 8 <= len(buf):
        cid = buf[pos:pos + 4]
        (csz,) = struct.unpack_from("<I", buf, pos + 4)
        body = buf[pos + 8: pos + 8 + csz]
        if cid == b"fmt ":
            fmt = body
        elif cid == b"data":
            data = body
        pos += 8 + csz + (csz & 1)  # chunks are word-aligned
    return fmt, data


def build_wav(fmt, data):
    riff_len = 4 + (8 + len(fmt)) + (8 + len(data))
    out = bytearray()
    out += b"RIFF" + struct.pack("<I", riff_len) + b"WAVE"
    out += b"fmt " + struct.pack("<I", len(fmt)) + fmt
    out += b"data" + struct.pack("<I", len(data)) + data
    return bytes(out)


def wave_by_id(blobs):
    """hex-id (e.g. 0x07D1) -> wave name. A playlist's segment N maps to base_id + N - 1."""
    out = {}
    for name in blobs:
        if name.upper().startswith("WAVE"):
            out[int(name[4:8], 16)] = name
    return out


def extract(ff_path, outdir, make_songs, convert):
    data = open(ff_path, "rb").read()
    entries = parse_ff(data)
    os.makedirs(outdir, exist_ok=True)

    blobs = {}
    raw_dir = os.path.join(outdir, "raw")
    os.makedirs(raw_dir, exist_ok=True)
    for name, off, size in entries:
        safe = name.strip().replace("/", "_").replace("\\", "_")
        if not safe:
            continue  # skip empty/terminator TOC entries
        blob = data[off:off + size]
        blobs[safe] = blob
        with open(os.path.join(raw_dir, safe), "wb") as fh:
            fh.write(blob)
    print(f"[{os.path.basename(ff_path)}] {len(entries)} entries -> {raw_dir}")

    songs = []
    if make_songs:
        song_dir = os.path.join(outdir, "songs")
        os.makedirs(song_dir, exist_ok=True)
        waves = wave_by_id(blobs)
        for name in sorted(blobs):
            if not name.upper().startswith("PLST"):
                continue
            grp = name[4:8]
            base = int(grp, 16)
            pl = blobs[name]
            seg_count, step_count = struct.unpack_from(">HH", pl, 0)  # PLST payload is big-endian
            seq = struct.unpack_from(">%dH" % step_count, pl, 4)
            fmt0 = None
            pcm = bytearray()
            ok = True
            for idx in seq:
                seg_name = waves.get(base + idx - 1)
                if seg_name is None:
                    print(f"  ! {name}: missing segment {idx} (wave id {base + idx - 1:04X})")
                    ok = False
                    break
                fmt, dat = wav_parts(blobs[seg_name])
                if fmt0 is None:
                    fmt0 = fmt
                pcm += dat
            if not ok or fmt0 is None:
                continue
            song_wav = os.path.join(song_dir, f"SONG_{grp}.wav")
            with open(song_wav, "wb") as fh:
                fh.write(build_wav(fmt0, bytes(pcm)))
            songs.append(song_wav)
            print(f"  song {grp}: {step_count} steps -> {os.path.basename(song_wav)}")

    if convert:
        ff = shutil.which("ffmpeg")
        if not ff:
            print("\n! ffmpeg not found on PATH; skipping conversion.")
            print("  Install with:  scoop install ffmpeg   (or)   choco install ffmpeg")
            return
        targets = songs if (make_songs and songs) else [
            os.path.join(raw_dir, n) for n in blobs if blobs[n][:4] == b"RIFF"
        ]
        conv_dir = os.path.join(outdir, convert)
        os.makedirs(conv_dir, exist_ok=True)
        codec = {"ogg": ["-c:a", "libvorbis", "-q:a", "5"],
                 "mp3": ["-c:a", "libmp3lame", "-q:a", "2"]}[convert]
        for src in targets:
            base = os.path.splitext(os.path.basename(src))[0]
            dst = os.path.join(conv_dir, base + "." + convert)
            subprocess.run([ff, "-y", "-loglevel", "error", "-i", src, *codec, dst], check=True)
        print(f"\nConverted {len(targets)} file(s) -> {conv_dir}")


def _varlen(n):
    """MIDI variable-length quantity."""
    out = bytearray([n & 0x7F])
    n >>= 7
    while n:
        out.insert(0, (n & 0x7F) | 0x80)
        n >>= 7
    return bytes(out)


def _smf(notes, ppq=480, bpm=120):
    """Build a format-0 SMF from [(note, ticks)] played back-to-back, monophonic."""
    trk = bytearray()
    # tempo meta
    us = int(60_000_000 / bpm)
    trk += _varlen(0) + bytes([0xFF, 0x51, 0x03]) + us.to_bytes(3, "big")
    for note, ticks in notes:
        trk += _varlen(0) + bytes([0x90, note, 100])        # note on
        trk += _varlen(max(1, ticks)) + bytes([0x80, note, 0])  # note off after duration
    trk += _varlen(0) + bytes([0xFF, 0x2F, 0x00])           # end of track
    head = b"MThd" + (6).to_bytes(4, "big") + (0).to_bytes(2, "big") + \
        (1).to_bytes(2, "big") + ppq.to_bytes(2, "big")
    return head + b"MTrk" + len(trk).to_bytes(4, "big") + bytes(trk)


def build_flstudio(ff_path, outdir, ppq=480, bpm=120):
    """Sample pack (all WAVE segments as .wav) + per-song MIDI reconstructing each playlist,
    with a single global note->sample mapping so one FL sampler covers every song."""
    data = open(ff_path, "rb").read()
    entries = parse_ff(data)
    blobs = {n: data[o:o + s] for n, o, s in entries if n.strip()}

    samp_dir = os.path.join(outdir, "samples")
    midi_dir = os.path.join(outdir, "midi")
    os.makedirs(samp_dir, exist_ok=True)
    os.makedirs(midi_dir, exist_ok=True)

    # global mapping: every WAVE id -> a unique MIDI note (base C1=24), + duration in ticks
    wave_ids = sorted(int(n[4:8], 16) for n in blobs if n.upper().startswith("WAVE"))
    BASE = 24
    tps = ppq * bpm / 60.0  # ticks per second
    note_of, ticks_of, file_of = {}, {}, {}
    rows = []
    for i, wid in enumerate(wave_ids):
        name = next(n for n in blobs if n.upper().startswith("WAVE") and int(n[4:8], 16) == wid)
        note = BASE + i
        fmt, dat = wav_parts(blobs[name])
        af, ch, sr, br, ba, bits = struct.unpack_from("<HHIIHH", fmt, 0)
        secs = len(dat) / (sr * ch * bits / 8)
        out_name = f"WAVE_{wid:04X}.wav"
        with open(os.path.join(samp_dir, out_name), "wb") as fh:
            fh.write(build_wav(fmt, dat))
        note_of[wid] = note
        ticks_of[wid] = round(secs * tps)
        file_of[wid] = out_name
        rows.append((wid, note, out_name, round(secs, 3)))

    with open(os.path.join(outdir, "mapping.csv"), "w") as fh:
        fh.write("wave_id,midi_note,sample_file,seconds\n")
        for wid, note, fn, secs in rows:
            fh.write(f"{wid:04X},{note},{fn},{secs}\n")

    n_songs = 0
    for name in sorted(blobs):
        if not name.upper().startswith("PLST"):
            continue
        base = int(name[4:8], 16)
        pl = blobs[name]
        seg_count, step_count = struct.unpack_from(">HH", pl, 0)
        seq = struct.unpack_from(">%dH" % step_count, pl, 4)
        notes, ok = [], True
        for idx in seq:
            wid = base + idx - 1
            if wid not in note_of:
                ok = False; break
            notes.append((note_of[wid], ticks_of[wid]))
        if not ok:
            continue
        with open(os.path.join(midi_dir, f"SONG_{name[4:8]}.mid"), "wb") as fh:
            fh.write(_smf(notes, ppq, bpm))
        n_songs += 1

    readme = os.path.join(outdir, "README_FLStudio.txt")
    with open(readme, "w") as fh:
        fh.write(
            "Havoc music -> FL Studio pack\n"
            "=============================\n\n"
            f"samples/   {len(rows)} WAVE segments (PCM mono 22050Hz 8-bit), named WAVE_<id>.wav\n"
            f"midi/      {n_songs} songs, each a MIDI that triggers samples in the original order\n"
            "mapping.csv  wave_id -> MIDI note -> sample file -> length(s)\n\n"
            "How to use:\n"
            "1. Add a sampler that maps keys chromatically starting at C1 (MIDI note 24):\n"
            "   - DirectWave or Slicex: drop all samples/ in, set root keys per mapping.csv, or\n"
            "   - Channel rack: load each sample on its own channel and key it per mapping.csv.\n"
            "2. Import a midi/SONG_*.mid into the playlist; it plays segments back-to-back\n"
            "   exactly as the game sequenced them. Tempo is a neutral 120 BPM; note lengths\n"
            "   equal each sample's real duration, so playback matches the assembled song.\n"
            "3. Remix freely - rearrange segments, layer, repitch, etc.\n"
        )
    # zip it
    import zipfile
    zpath = os.path.join(os.path.dirname(outdir) or ".",
                         os.path.basename(outdir.rstrip("/\\")) + ".zip")
    with zipfile.ZipFile(zpath, "w", zipfile.ZIP_DEFLATED) as z:
        for root, _, files in os.walk(outdir):
            for f in files:
                full = os.path.join(root, f)
                z.write(full, os.path.relpath(full, outdir))
    print(f"FL pack: {len(rows)} samples, {n_songs} MIDI songs -> {outdir}")
    print(f"Zipped -> {zpath}")


def main():
    ap = argparse.ArgumentParser(description="Extract Havoc .FF archives")
    ap.add_argument("ff", help="path to a .FF archive")
    ap.add_argument("-o", "--out", default=None, help="output dir (default: out_<name>)")
    ap.add_argument("--songs", action="store_true", help="assemble songs from PLST playlists")
    ap.add_argument("--convert", choices=["ogg", "mp3"], help="convert to ogg/mp3 (needs ffmpeg)")
    ap.add_argument("--fl", action="store_true", help="build FL Studio pack (samples + MIDI + zip)")
    a = ap.parse_args()
    out = a.out or ("out_" + os.path.splitext(os.path.basename(a.ff))[0].lower())
    if a.fl:
        build_flstudio(a.ff, out)
    else:
        extract(a.ff, out, a.songs, a.convert)


if __name__ == "__main__":
    main()