import os
import glob
import multiprocessing as mp


def ffmpeg_resample_48khz(filepath: str, dst_dir: str):
    basename = os.path.basename(filepath)
    input_ext = os.path.splitext(basename)[1]
    dst_filepath = os.path.join(dst_dir, basename).replace(input_ext, ".wav")
    cmd = f"""ffmpeg -y -i "{filepath}" -ar 48000 "{dst_filepath}" > /dev/null 2>&1"""
    os.system(cmd)


def ffmpeg_resample_24khz(filepath: str, dst_dir: str):
    basename = os.path.basename(filepath)
    dst_filepath = os.path.join(dst_dir, basename)
    cmd = f"""ffmpeg -y -i "{filepath}" -ar 24000 "{dst_filepath}" > /dev/null 2>&1"""
    os.system(cmd)


def ffmpeg_resample_12khz(filepath: str, dst_dir: str):
    basename = os.path.basename(filepath)
    dst_filepath = os.path.join(dst_dir, basename)
    cmd = f"""ffmpeg -y -i "{filepath}" -ar 12000 "{dst_filepath}" > /dev/null 2>&1"""
    os.system(cmd)


if __name__ == "__main__":

    num_workers = 64
    codec_mode = False
    resample_func = ffmpeg_resample_48khz

    if not codec_mode:
        # src_dir = "/app/suno/christian/reference-audio-wav"
        # dst_dir = "/app/suno/christian/reference-audio-wav-24khz"
        src_dir = "/app/suno/christian/data/singing/amazing_grace"
        dst_dir = "/app/suno/christian/data/singing/amazing_grace_48khz"
        os.makedirs(dst_dir, exist_ok=True)

        # find all wav files in source directory
        src_filepaths = glob.glob(os.path.join(src_dir, "*.m4a"))
        if len(src_filepaths) == 0:
            raise ValueError("No files found in source directory")
        args = [(src_filepath, dst_dir) for src_filepath in src_filepaths]

        with mp.Pool(num_workers) as pool:
            pool.starmap(resample_func, args)

    else:
        src_dir = "/app/suno/data/audio_2ch_48khz_lg/"
        dst_dir = "/app/suno/data/audio_2ch_12khz_lg/"
        os.makedirs(dst_dir, exist_ok=True)
        datasets = [
            "genius_hq",
            "youtube_music",
            "podcasts",
            "jamendo",
            "imslp",
            "pond5_music",
            "pond5_sfx",
            "shutter_music",
            "spot_genres",
            "tency",
        ]
        subsets = ["val", "train"]

        for subset in subsets:
            print(subset)
            os.makedirs(os.path.join(dst_dir, subset), exist_ok=True)
            for dataset in datasets:
                print(dataset)
                os.makedirs(os.path.join(dst_dir, subset, dataset), exist_ok=True)
                # find all source audio files
                subset_src_dir = os.path.join(src_dir, subset, dataset)
                subset_dst_dir = os.path.join(dst_dir, subset, dataset)
                src_filepaths = glob.glob(os.path.join(subset_src_dir, "*.wav"))
                args = [
                    (src_filepath, subset_dst_dir) for src_filepath in src_filepaths
                ]

                with mp.Pool(num_workers) as pool:
                    pool.starmap(resample_func, args)
                print()