import os
import csv
import glob
from typing import List


def find_all_audio_files(audio_dir: str, exts: List[str] = ["wav"]):
    # find all audio files
    filepaths = []
    for ext in exts:
        filepaths += glob.glob(
            os.path.join(audio_dir, "**", f"*.{ext}"), recursive=True
        )
    return filepaths


if __name__ == "__main__":

    for subset in ["val", "train"]:
        print(subset)
        if subset == "train":
            top_dir = "/app/suno/data/audio_2ch_12khz_lg/train"
        else:
            top_dir = "/app/suno/data/audio_2ch_12khz_lg/val"
        # find all directories in subset
        datasets = glob.glob(os.path.join(top_dir, "*"))
        all_filepaths = []
        for dataset in datasets:
            print(dataset)
            dataset_dir = os.path.join(top_dir, dataset)
            filepaths = glob.glob(os.path.join(dataset_dir, "*.wav"))
            all_filepaths.extend(filepaths)

        manifest_path = os.path.join(os.path.dirname(top_dir), f"{subset}_manifest.csv")
        with open(manifest_path, "w") as fp:
            fp.write(f"path\n")
            for filepath in all_filepaths:
                fp.write(f"{filepath}\n")