import os
import re
import struct
import subprocess
import sys

FLUIDSYNTH_LOCATION = "./fluidsynth/fluidsynth.exe"
# gzdoom.sf2 sounds weird compared to gzdoom itself (synths too loud, drums too quiet);
# if you don't have a particular soundfont in mind then Helion's works fine:
# https://github.com/Helion-Engine/Helion/blob/master/Client/SoundFonts/Default.sf2
SOUNDFONT_LOCATION = "./soundfont.sf2"
ADJUST_MP3_VOLUME = True # see convert()
OUTPUT_BASE_FOLDER = "./output"

MUSIC_NAME_MAP = {
    "D_DM2TTL": "title",
    "D_READ_M": "text screen",
    "D_DM2INT": "intermission",
    "D_RUNNIN": "MAP01",
    "D_STALKS": "MAP02",
    "D_COUNTD": "MAP03",
    "D_BETWEE": "MAP04",
    "D_DOOM": "MAP05",
    "D_THE_DA": "MAP06",
    "D_SHAWN": "MAP07",
    "D_DDTBLU": "MAP08",
    "D_IN_CIT": "MAP09",
    "D_DEAD": "MAP10",
    "D_STLKS2": "MAP11",
    "D_THEDA2": "MAP12",
    "D_DOOM2": "MAP13",
    "D_DDTBL2": "MAP14",
    "D_RUNNI2": "MAP15",
    "D_DEAD2": "MAP16",
    "D_STLKS3": "MAP17",
    "D_ROMERO": "MAP18",
    "D_SHAWN2": "MAP19",
    "D_MESSAG": "MAP20",
    "D_COUNT2": "MAP21",
    "D_DDTBL3": "MAP22",
    "D_AMPIE": "MAP23",
    "D_THEDA3": "MAP24",
    "D_ADRIAN": "MAP25",
    "D_MESSG2": "MAP26",
    "D_ROMER2": "MAP27",
    "D_TENSE": "MAP28",
    "D_SHAWN3": "MAP29",
    "D_OPENIN": "MAP30",
    "D_EVIL": "MAP31 & cast sequence",
    "D_ULTIMA": "MAP32",
}


def get_wad_musics(wad_file: str) -> list[tuple[str, bytes]]:
    musics: list[tuple[str, bytes]] = []
    with open(wad_file, "rb") as f:
        HEADER_SIZE = 12
        DIRECTORY_ENTRY_SIZE = 16
        magic_num, directory_entry_count, directory_offset = struct.unpack("<4sII", f.read(HEADER_SIZE))
        assert magic_num in (b"IWAD", b"PWAD")

        lump_info = []
        f.seek(directory_offset)
        for _ in range(directory_entry_count):
            lump_index, lump_size, lump_name_raw = struct.unpack("<II8s", f.read(DIRECTORY_ENTRY_SIZE))
            lump_name = lump_name_raw.decode("ascii").rstrip("\0")
            lump_info.append((lump_index, lump_size, lump_name))

        for lump_index, lump_size, lump_name in lump_info:
            print(lump_name)
            f.seek(lump_index)
            lump_data = f.read(lump_size)
            if lump_data[:4] == b'MThd': # TODO: get MUS as well
                musics.append((lump_name, lump_data))

    return musics


def convert(midi_path: str, mp3_path: str) -> None:

    # render the midi
    midi_proc = subprocess.run([
        FLUIDSYNTH_LOCATION,
        "--quiet",
        "-T", "raw",
        "-F", "-",
        SOUNDFONT_LOCATION,
        midi_path
    ], stdout=subprocess.PIPE, check=True)
    raw_audio = midi_proc.stdout

    # determine how much to boost volume if desired;
    # we could use replaygain tags for personal listening but for internet stuff
    # an adjustment should be made because it's pretty quiet by default
    volume_diff: float = 0
    if ADJUST_MP3_VOLUME:
        volume_proc = subprocess.run([
            "ffmpeg",
            "-hide_banner",
            "-f", "s32le",
            "-i", "-",
            "-filter:a", "volumedetect",
            "-f", "null",
            "-"
        ], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, input=raw_audio, check=True)
        volume_info = volume_proc.stderr
        mean_volume = float(re.findall(rb"mean_volume: ([\d.-]+) dB", volume_info)[0])
        # "ReplayGain nominally plays at -14 dB relative to full-scale leaving
        # 14 dB of headroom for reproduction of dynamic material."
        volume_diff = -14 - mean_volume

    # convert to mp3 with the volume adjustment
    subprocess.run([
        "ffmpeg",
        "-hide_banner",
        "-loglevel", "warning",
        "-f", "s32le",
        "-i", "-",
        "-filter:a", f"volume={volume_diff}dB",
        "-b:a", "192k",
        "-y",
        mp3_path
    ], input=raw_audio, check=True)


def ensure_folder_exists(path: str) -> None:
    if not os.path.exists(path):
        os.makedirs(path)


def process_wads(wads: list[str]) -> None:
    for fpath in wads:
        fn = os.path.split(fpath)[-1]
        fn_stub, ext = os.path.splitext(fn)
        assert ext.lower() == ".wad"

        wad_musics = get_wad_musics(fpath)

        output_folder = os.path.join(OUTPUT_BASE_FOLDER, fn_stub)
        midi_folder = os.path.join(output_folder, "midi")
        mp3_folder = output_folder # os.path.join(output_folder, "mp3")
        ensure_folder_exists(midi_folder)
        ensure_folder_exists(mp3_folder)

        for name, data in wad_musics:
            file_stub = MUSIC_NAME_MAP.get(name, name)
            midi_path = os.path.join(midi_folder, f"{file_stub}.mid")
            mp3_path = os.path.join(mp3_folder, f"{file_stub}.mp3")
            print(f"{midi_path} ({name})")
            with open(midi_path, "wb") as f2:
                f2.write(data)
            convert(midi_path, mp3_path)


if __name__ == "__main__":
    process_wads([x for x in sys.argv[1:] if x.lower().endswith(".wad")])