import streamlit as st
from gradio_client import Client
import time

def generate_image(prompt):
    try:
        client = Client("mukaist/Midjourney")
        result = client.predict(
            prompt=prompt,
            negative_prompt="(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck",
            use_negative_prompt=True,
            style="2560 x 1440",
            seed=0,
            width=1024,
            height=1024,
            guidance_scale=6,
            randomize_seed=True,
            api_name="/run"
        )
        return result
    except Exception as e:
        st.error(f"Error generating image: {str(e)}")
        return None

def handle_prompt_click(prompt_text, key):
    st.session_state[f'selected_prompt_{key}'] = prompt_text
    
    # Show loading message
    with st.spinner('Generating artwork...'):
        # Generate the image
        result = generate_image(prompt_text)
        
        if result:
            # Store the result in session state
            st.session_state[f'generated_image_{key}'] = result
            st.success("Artwork generated successfully!")

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 generated images and prompts
    st.markdown("---")
    st.markdown("### Generated Artwork:")
    
    # Create columns for displaying images and prompts
    display_cols = st.columns(2)
    
    for key in st.session_state:
        if key.startswith('selected_prompt_'):
            idx = key.split('_')[-1]
            image_key = f'generated_image_{idx}'
            
            # Display prompt
            with display_cols[0]:
                st.write("Prompt:", st.session_state[key])
            
            # Display image if it exists
            if image_key in st.session_state:
                with display_cols[1]:
                    # Assuming the result is a path or URL to the image
                    st.image(st.session_state[image_key], caption="Generated Artwork", use_column_width=True)

if __name__ == "__main__":
    main()