import streamlit as st
import random
from PIL import Image
import os

# Load images from the 'images' folder
def load_images():
    image_folder = "images"  
    images = {
        "Mark": "mark.jpeg", "Renjun": "renjun.jpeg", "Jeno": "jeno.jpeg",
        "Haechan": "haechan.jpeg", "Jaemin": "jaemin.jpeg", "Chenle": "chenle.jpeg", "Jisung": "jisung.jpeg"
    }
    loaded_images = {}
    for label, filename in images.items():
        img_path = os.path.join(image_folder, filename)
        if os.path.exists(img_path):
            loaded_images[label] = Image.open(img_path)
    return loaded_images

# Load images and store labels
images = load_images()
image_labels = list(images.keys())

# Initialize game state
if "score" not in st.session_state:
    st.session_state.score = 0  # Tracks correct answers
    st.session_state.image_pool = image_labels.copy()  # Pool of images to show
    random.shuffle(st.session_state.image_pool)  # Shuffle images
    st.session_state.selected_image = None  # Current image to guess
    st.session_state.wrong_message = ""  # Stores wrong answer message
    st.session_state.correct_message = ""  # Stores correct answer message

st.title("7DREAM Image Matching Game")

# Check if all images have been correctly matched
if st.session_state.score == len(image_labels):
    st.success("🎉 Congratulations! You are a Dreamzen. 🎉")

    col1, col2 = st.columns(2)
    with col1:
        if st.button("🔄 Play Again"):
            st.session_state.score = 0
            st.session_state.image_pool = image_labels.copy()
            random.shuffle(st.session_state.image_pool)
            st.session_state.selected_image = None
            st.session_state.wrong_message = ""
            st.session_state.correct_message = ""
            st.rerun()

    with col2:
        if st.button("❌ Exit Game"):
            st.warning("You can now close this tab to exit the game.")
            st.stop()

# Continue game if not finished
else:
    st.write("Match the image to the correct member!")

    # Pick a new image from the pool
    if not st.session_state.selected_image:
        if not st.session_state.image_pool:  # Reshuffle if all images were shown
            st.session_state.image_pool = image_labels.copy()
            random.shuffle(st.session_state.image_pool)
        
        st.session_state.selected_image = st.session_state.image_pool.pop()  # Select next image

    # Display the image
    st.image(images[st.session_state.selected_image], caption="Who is this?", use_container_width=True)

    # Provide answer choices
    selected_label = st.radio("Choose the correct member:", image_labels, key="selected_label")

    if st.button("Submit"):
        if selected_label == st.session_state.selected_image:
            st.session_state.correct_message = "✅ Correct! 🎉"
            st.session_state.wrong_message = ""  # Clear wrong message
            st.session_state.score += 1  # Increase score
            st.session_state.selected_image = None  # Move to next image
            st.rerun()
        else:
            st.session_state.wrong_message = "❌ Wrong! Try again."
            st.session_state.correct_message = ""  # Clear correct message

    # Show messages persistently
    if st.session_state.correct_message:
        st.success(st.session_state.correct_message)
    
    if st.session_state.wrong_message:
        st.error(st.session_state.wrong_message)

st.write(f"**Score:** {st.session_state.score} / {len(image_labels)}")