# Import necessary libraries
import tempfile  # Library for temporary file operations
from audiocraft.models import MusicGen  # Import the MusicGen model for music generation
from audiocraft.data.audio import audio_write  # Import audio_write function for saving audio
import gradio as gr  # Import Gradio for creating a web interface
import torch  # Import PyTorch for deep learning operations
import uuid  # Library for generating unique identifiers (UUIDs)
import os  # Operating system-related functions
from scipy.io.wavfile import write  # Function for writing audio to a WAV file

# Load the model and set parameters
model = MusicGen.get_pretrained("facebook/musicgen-small")  # Load a pre-trained MusicGen model
model.set_generation_params(duration=5)  # Set the duration for audio generation in seconds

# Function to generate music from text descriptions
def generate_music(description):
    # Generate audio based on the provided description
    wav = model.generate([description])  # Generate audio using the model
    audio_array = wav.cpu().numpy().squeeze()  # Convert the audio to a NumPy array
    sample_rate = model.sample_rate  # Get the audio sample rate (e.g., 44100 Hz)

    # Generate a unique file identifier (UUID) for the temporary audio file
    file_id = uuid.uuid1()

    # Define the path for the temporary audio file using the identifier
    file_path = os.path.join(
        tempfile.gettempdir(),  # Get the temporary directory path
        f'{file_id}.wav'  # Create a unique file name with the UUID and .wav extension
    )
    
    # Add debugging statements to check file paths and directory permissions
    print(f"Temporary directory: {tempfile.gettempdir()}")
    print(f"File path: {file_path}")

    # Write the generated audio to the temporary file as a WAV file
    write(file_path, rate=sample_rate, data=audio_array)

    # Return the path to the temporary audio file
    return file_path

# Define Gradio interface with temporary file output
iface = gr.Interface(
    fn=generate_music,  # Use the generate_music function for processing input
    inputs="text",  # Accept text input from the user
    outputs=gr.components.Audio(type="filepath", label="Audio"),  # Display the generated audio as output
    title="Text to Audio Generation",  # Set the title of the web interface
    description="Generate audio based on text descriptions.",  # Provide a description
    live=False,  # Set to False if you don't want real-time updates (for beginner-friendly interaction)
)

# Start the Gradio interface
iface.launch(debug=True)