import streamlit as st
import random

# Initialize session state
if "current_page" not in st.session_state:
    st.session_state.current_page = "main"
if "mode" not in st.session_state:
    st.session_state.mode = ""
if "number" not in st.session_state:
    st.session_state.number = 0

def main_page():
    st.title("Before and After Game")
    st.write("Welcome to the Before and After Game! Click 'Play' to get started.")

    if st.button("Play", key="play_button"):
        st.session_state.current_page = "select_mode"

def select_mode():
    st.title("Choose a Mode")
    st.write("Select 'Before' to find the number before a given number or 'After' to find the number after it.")

    col1, col2 = st.columns(2)
    with col1:
        if st.button("Before", key="before_button"):
            st.session_state.mode = "before"
            st.session_state.current_page = "game"
    with col2:
        if st.button("After", key="after_button"):
            st.session_state.mode = "after"
            st.session_state.current_page = "game"

def start_game():
    st.session_state.number = random.randint(1, 100)
    ask_question()

def ask_question():
    number = st.session_state.number
    mode = st.session_state.mode
    st.title(f"What is {mode.upper()} {number}?")
    st.write("Choose the correct answer from the options below.")

    # Generate options
    correct_answer = number - 1 if mode == "before" else number + 1
    options = [correct_answer]
    while len(options) < 4:
        option = random.randint(1, 100)
        if option not in options:
            options.append(option)
    random.shuffle(options)

    # Create buttons for options
    col1, col2 = st.columns(2)
    for i, option in enumerate(options):
        if i % 2 == 0:
            with col1:
                if st.button(str(option), key=f"option_{option}"):
                    check_answer(option, correct_answer)
        else:
            with col2:
                if st.button(str(option), key=f"option_{option}"):
                    check_answer(option, correct_answer)

def check_answer(selected, correct):
    if selected == correct:
        st.success("Great Job! 🎉")
        st.balloons()
        if st.button("Next Question"):
            start_game()
    else:
        st.error("Oops! Try again!")

def app():
    if st.session_state.current_page == "main":
        main_page()
    elif st.session_state.current_page == "select_mode":
        select_mode()
    elif st.session_state.current_page == "game":
        start_game()

# Run the application
if __name__ == "__main__":
    app()