import streamlit as st

def handle_prompt_click(prompt_text, key):
    st.session_state[f'selected_prompt_{key}'] = prompt_text
    st.write(f"Generated prompt for: {prompt_text}")

def main():
    st.title("🎨 Art Prompt Generator")
    st.markdown("### Select a prompt style to generate artwork:")

    # Dictionary mapping prompts to emojis
    prompt_emojis = {
        "AIart/AIArtistCommunity": "🤖",
        "Black & White": "⚫⚪",
        "Black & Yellow": "⚫💛",
        "Blindfold": "🙈",
        "Break": "💔",
        "Broken": "🔨",
        "Christmas Celebrations art": "🎄",
        "Colorful Art": "🎨",
        "Crimson art": "🔴",
        "Eyes Art": "👁️",
        "Going out with Style": "💃",
        "Hooded Girl": "🧥",
        "Lips": "👄",
        "MAEKHLONG": "🏮",
        "Mermaid": "🧜‍♀️",
        "Morning Sunshine": "🌅",
        "Music Art": "🎵",
        "Owl": "🦉",
        "Pink": "💗",
        "Purple": "💜",
        "Rain": "🌧️",
        "Red Moon": "🌑",
        "Rose": "🌹",
        "Snow": "❄️",
        "Spacesuit Girl": "👩‍🚀",
        "Steampunk": "⚙️",
        "Succubus": "😈",
        "Sunlight": "☀️",
        "Weird art": "🎭",
        "White Hair": "👱‍♀️",
        "Wings art": "👼",
        "Woman with Sword": "⚔️"
    }

    # Create columns for better button layout
    col1, col2, col3 = st.columns(3)
    
    # Distribute buttons across columns
    for idx, (prompt, emoji) in enumerate(prompt_emojis.items()):
        full_prompt = f"QT {prompt}"
        col = [col1, col2, col3][idx % 3]
        
        with col:
            if st.button(f"{emoji} {prompt}", key=f"btn_{idx}"):
                handle_prompt_click(full_prompt, idx)
                
    # Display selected prompt if any
    st.markdown("---")
    st.markdown("### Generated Prompts:")
    for key in st.session_state:
        if key.startswith('selected_prompt_'):
            st.write(st.session_state[key])

if __name__ == "__main__":
    main()