import os
import glob
import torch
import torchaudio

from tqdm import tqdm

if __name__ == "__main__":

    # get all song directories
    song_dirs = glob.glob(os.path.join(f"/app/suno/christian/data/ursing/**"))
    print(len(song_dirs))

    vocal_filepaths = []
    for song_dir in song_dirs:
        # get the vocal from each song
        vocal_filepath = os.path.join(song_dir, "Vocal.wav")
        vocal_filepaths.append(vocal_filepath)

    os.makedirs("/app/suno/christian/data/ursing_vocal_48khz/", exist_ok=True)
    os.makedirs("/app/suno/christian/data/ursing_vocal_48khz/audio", exist_ok=True)
    os.makedirs(
        "/app/suno/christian/data/ursing_vocal_48khz/audio/train", exist_ok=True
    )
    os.makedirs("/app/suno/christian/data/ursing_vocal_48khz/audio/val", exist_ok=True)

    # first 60 for train and the rest for val
    train_filepaths = vocal_filepaths[:60]
    val_filepaths = vocal_filepaths[60:]

    idx = 0

    for vocal_filepath in tqdm(train_filepaths):
        vocal_filename = os.path.basename(vocal_filepath)
        # use ffmpeg system call to convert to 48khz
        os.system(
            f"ffmpeg -i {vocal_filepath} -ar 48000 /app/suno/christian/data/ursing_vocal_48khz/audio/train/{idx:03d}-{vocal_filename}"
        )
        idx += 1

    for vocal_filepath in tqdm(val_filepaths):
        vocal_filename = os.path.basename(vocal_filepath)
        # use ffmpeg system call to convert to 48khz
        os.system(
            f"ffmpeg -i {vocal_filepath} -ar 48000 /app/suno/christian/data/ursing_vocal_48khz/audio/val/{idx:03d}-{vocal_filename}"
        )
        idx += 1