import os
import glob
import warnings
import random

procid = int(os.environ.get("SLURM_PROCID", 0))
localid = int(os.environ.get("SLURM_LOCALID", 0))
world_size = int(os.environ.get("SLURM_JOB_NUM_NODES", 1)) * int(
    os.environ.get("SLURM_NTASKS_PER_NODE", 1)
)

assert world_size > 0, "WORLD_SIZE is 0"

print(f"PROCID: {procid}, LOCALID: {localid}, WORLD_SIZE: {world_size}")

dir = "/app2/suno/data/victor/gmd/MIDIs"
output_dir = "/app2/suno/data/sara/gmd/synth_v6"

soundfonts = [
    "/home/victor/neon/sunoData/src/sunodata/midi/soundfonts/FatBoy-v0.786.sf2",
    "/home/victor/neon/sunoData/src/sunodata/midi/soundfonts/Touhou.sf2",
    "/home/sara/soundfonts/good/RLNDGM.sf2",
    "/home/sara/soundfonts/good/arachno.sf2",
    "/home/sara/soundfonts/good/musescore_general.sf2",
]


def process_midi_file(midi_file, sf_path):
    from suno_utils.audio.midi import Midi

    try:
        outpath = midi_file.replace(dir, output_dir).replace(".mid", ".opus")
        os.makedirs(os.path.dirname(outpath), exist_ok=True)
        if not os.path.exists(outpath):
            midi = Midi.from_path(midi_file)
            midi.pmidi.remove_invalid_notes()
            if midi.duration_s() < 30:
                return f"Skipped {midi_file}: too short"
            if midi.duration_s() > 300:
                return f"Skipped {midi_file}: too long"
            note_density = midi.num_notes / midi.duration_s()
            if note_density > 20:
                return f"Skipped {midi_file}: note density too high"

            try:
                audio = midi.to_audio(sf_path=sf_path)
            except Exception as e:
                try:
                    audio = midi.to_audio(sf_path=random.randint(0, len(soundfonts) - 1))
                except Exception as e:
                    audio = midi.to_audio(sf_path=random.randint(0, len(soundfonts) - 1))

            audio.write_opus(outpath)
        return f"Processed {midi_file}"
    except Exception as e:
        # print(e)
        return f"Error loading {midi_file}: {e}"


def main():
    print("main!")
    from tqdm import tqdm

    midi_files = list(glob.iglob(os.path.join(dir, "**", "*.mid"), recursive=True))
    midi_files = midi_files[procid::world_size]
    print(len(midi_files))

    skipped = 0
    for i, midi_file in tqdm(enumerate(midi_files)):
        sf_path = soundfonts[i % len(soundfonts)]  # Cycle through available soundfonts
        result = process_midi_file(midi_file, sf_path)
        if "Error" in result or "Skipped" in result:
            skipped += 1
    print("SKIPPED", skipped)


main()