from suno_utils.utils.text import write_json
from suno_utils.utils.s3 import list_s3_dir, upload_s3_files
from pathlib import Path

midi_s3_paths = list_s3_dir("s3://suno-data/sara/midi/msmt_midi_test_set/")

final_data = []
for s3_midi_path,_ in midi_s3_paths:
    example_id = Path(s3_midi_path).stem
    s3_audio_path = f"s3://suno-data/sara/midi/msmt_midi_test_set_audio/{example_id}.wav"
    s3_midi_path = f"s3://suno-data/sara/midi/msmt_midi_test_set/{example_id}.mid"
    final_data.append(
        dict(
            id=example_id,
            s3_midi_path=s3_midi_path,
            s3_audio_path=s3_audio_path
        )
    )
write_json(final_data, "midi_test_files.json")