import os
import glob
import tqdm
import argparse
from suno_utils.audio import Audio


# get file list
filelist = glob.glob("/app/suno/minz/datasets/fma/fma_large/*/*.mp3")
print("%d files exsist" % len(filelist))

# parse
parser = argparse.ArgumentParser()
parser.add_argument("--index", required=True, help="partial index")
parser.add_argument("--num_partial", required=True, help="number of partial data")
args = parser.parse_args()
index = int(args.index)
num_partial = int(args.num_partial)

# split data
hop = len(filelist) // num_partial
if index == num_partial:
    filelist = filelist[hop * index - 1:]
else:
    filelist = filelist[hop * index : (index + 1) * hop]

# resample audio
save_path = "/app/suno/data/audio_mono_24khz/fma/audio"

for fn in tqdm.tqdm(filelist):
    save_fn = os.path.join(save_path, fn.split("/")[-1][:-3] + "wav")
    if not os.path.exists(save_fn):
        try:
            audio = Audio.from_file(fn, sample_rate=24000)
            audio.write_wav(save_fn)
        except:
            continue