#!/usr/bin/env python3

# Uses ffprobe to read sound file length.
# So uh, have ffprobe accessible somehow.

import argparse
import asyncio
import json
import os
import subprocess
from typing import Dict, Awaitable


DIRECTORIES = [
    "vox",
    "vox_fem",
    "vox_sfx",
    "vox_mas"
]

PROCESS_BATCH_COUNT = 32

async def main() -> Awaitable[None]:
    parser = argparse.ArgumentParser()
    parser.add_argument("outfile", help="The .dm file to output the length info to.")
    parser.add_argument("sounds_directory", default="../sound/", help="The sound/ directory")
    args = parser.parse_args()

    sounds_directory = args.sounds_directory
    outfile = args.outfile

    durations = {}
    tasks = []

    count = 0

    for directory_name in DIRECTORIES:
        directory = os.path.join(sounds_directory, directory_name)

        for filename in os.listdir(directory):
            print(filename)

            tasks.append(do_work(sounds_directory, directory_name, filename, durations))

            count += 1
            # We have to process these in batches.
            # We can't just start 5000 ffprobe processes.
            if count == PROCESS_BATCH_COUNT:
                await asyncio.gather(*tasks)
                count = 0
                tasks = []

    if tasks:
        await asyncio.gather(*tasks)

    with open(outfile, "w") as f:
        f.write("// Automatically generated by vox_sound_lengths.py\n")
        f.write("var/list/vox_sound_lengths = list()\n")
        # We can't put this stuff inside a regular list() call because it's too damn big apparently.
        f.write("/__vox_sound_lengths_init/New()\n")
        for (name, length) in durations.items():
            f.write(f"\tvox_sound_lengths['{name}'] = {length}\n")
        f.write("/var/__vox_sound_lengths_init/__vox_sound_lengths_init_instance = new\n")

async def do_work(sound_dir: str, dir_name: str, file_name: str, durations: Dict[str, float]) -> Awaitable[None]:
    file_path = os.path.join(sound_dir, dir_name, file_name)
    duration = await get_audio_file_length(file_path)
    durations[f"sound/{dir_name}/{file_name}"] = duration

async def get_audio_file_length(file_name: str) -> Awaitable[float]:
    ffprobe_args = ["-i", file_name, "-show_entries", "format=duration", "-v", "quiet", "-of", "json"]
    proc = await asyncio.create_subprocess_exec("ffprobe", *ffprobe_args, stdout=subprocess.PIPE, encoding="utf-8")
    await proc.wait()

    j = json.loads(await proc.stdout.read())
    return float(j["format"]["duration"])

asyncio.run(main())