import gradio as gr
from musiclang_predict import MusicLangPredictor  # Assuming this is the correct import
from midi2audio import FluidSynth
import os

def musiclang(nb_tokens, temperature, chord_progression):
    top_p = 1.0
    seed = 0

    # If chord progression is empty, we might set it to None or handle it accordingly
    if chord_progression.strip() == "":
        chord_progression = None  # Assuming the predictor can handle None as "generate freely"

    # Initialize the MusicLangPredictor
    ml = MusicLangPredictor('musiclang/musiclang-v2')
    # Generate the score
    if chord_progression is None or chord_progression.strip() == "":
        score = ml.predict(
            nb_tokens=int(nb_tokens),
            temperature=float(temperature),
            topp=top_p,
            chord_progression=chord_progression,  # Pass the chord progression, can be None
            rng_seed=seed
        )
    else:
        score = ml.predict_chords(
            chord_progression,
            time_signature=(4, 4),
            temperature=temperature,
            topp=top_p,
            rng_seed=seed)

    # Save the score as a MIDI file
    midi_path = 'test.mid'
    score.to_midi(midi_path)

    # Convert MIDI to WAV then WAV to MP3
    wav_path = 'result.wav'
    mp3_path = 'result.mp3'
    FluidSynth("/usr/share/sounds/sf2/FluidR3_GM.sf2").midi_to_audio(midi_path, wav_path)

    os.system(f'ffmpeg -i {wav_path} -acodec libmp3lame -y -loglevel quiet -stats {mp3_path}')

    # Return the path to the MP3 for Gradio to display

    # Return both the MP3 path for Gradio to display and the MIDI file path for download
    return mp3_path, midi_path


# Gradio interface with inputs for temperature, nb_tokens, chord progression, and tempo
iface = gr.Interface(
    fn=musiclang,
    inputs=[
        gr.Number(label="Number of Tokens", value=1024, minimum=256, maximum=2048, step=256),
        gr.Slider(label="Temperature", value=0.9, minimum=0.1, maximum=1.0, step=0.1),
        gr.Textbox(label="Chord Progression", placeholder="Am CM Dm/F E7 Am", lines=2, value=""),
        gr.Slider(label="Tempo", value=120, minimum=60, maximum=240, step=1)
    ],
    outputs=[
        gr.Audio(label="Generated Music"),
        gr.File(label="Download MIDI")
    ],
    title="Music Generation with MusicLang",
    description="""Customize the music generation by specifying the number of tokens, temperature, chord progression, and tempo.
    \nChord qualities: M, m, 7, m7b5, sus2, sus4, m7, M7, dim, dim7. You can also specify the bass if it belongs to the chord (e.g., Bm/D).
    \nIf no chord progression is given, it generates a free sample with the specified number of tokens."""
)

iface.launch()