import streamlit as st
import os
import librosa
import soundfile as sf
import numpy as np
from inference import generate_drum_kit
from audio_utils import play_audio

# Streamlit UI
st.title("Generate Drum Kits with Text")

# User Inputs
prompt = st.text_input("Describe your drum kit (e.g., 'warm vintage')", "8-bit video game")
kit_size = st.slider("Number of sounds per instrument:", 1, 10, 4)

# Run the inference
if st.button("Generate Drum Kit"):
    drum_kit = generate_drum_kit(prompt, kit_size)
    st.session_state["drum_kit"] = drum_kit  # Store results

# Display results
if "drum_kit" in st.session_state:
    drum_kit = st.session_state["drum_kit"]
    st.subheader("Generated Drum Kit")

    for instrument, sounds in drum_kit.items():
        st.write(f"**{instrument}**")
        cols = st.columns(len(sounds))

        for i, sound_file in enumerate(sounds):
            with cols[i]:
                if st.button(f"▶️ {os.path.basename(sound_file)}", key=sound_file):
                    play_audio(sound_file)