import streamlit as st
import openai
import datetime
import json

# Set your OpenAI API key
openai.api_key = 'YOUR_OPENAI_API_KEY'

# Set up the Streamlit page configuration
st.set_page_config(
    page_title="JoyStory - Interactive Story Adventure",
    page_icon=":book:",
    layout="wide",
    initial_sidebar_state="collapsed",
)

# Initialize session state variables
if 'story' not in st.session_state:
    st.session_state.story = []
if 'achievements' not in st.session_state:
    st.session_state.achievements = []
if 'level' not in st.session_state:
    st.session_state.level = 'Beginner'
if 'unique_words' not in st.session_state:
    st.session_state.unique_words = set()
if 'total_words' not in st.session_state:
    st.session_state.total_words = 0
if 'badges' not in st.session_state:
    st.session_state.badges = []

# Function to generate AI response
def generate_ai_response(user_input, level):
    prompt = f"Continue the story in a way that is engaging and suitable for a {level.lower()} English learner.\n\nStory so far:\n{''.join(st.session_state.story)}\nUser input: {user_input}\nAI continuation:"
    response = openai.Completion.create(
        engine='gpt-4o-mini'
        prompt=prompt,
        max_tokens=100,
        temperature=0.7,
    )
    ai_output = response.choices[0].text.strip()
    return ai_output

# Function to get vocabulary suggestions
def get_vocabulary_suggestions():
    suggestions = openai.Completion.create(
        engine='gpt-4o-mini',
        prompt=f"Provide 3 interesting words suitable for a {st.session_state.level.lower()} English learner to use in the next part of the story.",
        max_tokens=20,
        n=1,
        stop=None,
        temperature=0.5,
    )
    words = suggestions.choices[0].text.strip().split(', ')
    return words

# Function to get creative prompts
def get_creative_prompt():
    prompt = openai.Completion.create(
        engine='gpt-4o-mini',
        prompt="Ask an open-ended question to inspire the next part of the story.",
        max_tokens=30,
        n=1,
        stop=None,
        temperature=0.7,
    )
    question = prompt.choices[0].text.strip()
    return question

# Function to provide feedback
def get_feedback(user_input):
    feedback_prompt = f"Provide one positive, simple feedback on the following sentence for a {st.session_state.level.lower()} English learner:\n\n{user_input}"
    feedback = openai.Completion.create(
        engine='gpt-4o-mini',
        prompt=feedback_prompt,
        max_tokens=50,
        temperature=0.5,
    )
    return feedback.choices[0].text.strip()

# Function to update achievements
def update_achievements(user_input):
    words = set(user_input.lower().split())
    st.session_state.unique_words.update(words)
    st.session_state.total_words += len(words)
    # Example achievement: Use 50 unique words
    if len(st.session_state.unique_words) >= 50 and '50 Unique Words' not in st.session_state.badges:
        st.session_state.badges.append('50 Unique Words')
        st.success("Achievement Unlocked: 50 Unique Words!")

# Layout the UI
st.title("📖 JoyStory - Interactive Story Adventure")

col1, col2 = st.columns([3, 1])

with col1:
    # Story Display Box
    st.header("Story So Far")
    story_display = st.empty()
    if st.session_state.story:
        story_display.markdown('\n\n'.join(st.session_state.story))
    else:
        st.info("Your story will appear here as you write it.")

    # User Input Box
    st.header("Your Turn to Write")
    user_input = st.text_area("Continue the story:", height=100)
    if st.button("Submit"):
        if user_input.strip() == "":
            st.warning("Please enter some text to continue the story.")
        else:
            st.session_state.story.append(f"**You:** {user_input}")
            update_achievements(user_input)
            feedback = get_feedback(user_input)
            st.session_state.story.append(f"**Feedback:** {feedback}")

            ai_response = generate_ai_response(user_input, st.session_state.level)
            st.session_state.story.append(f"**AI:** {ai_response}")
            story_display.markdown('\n\n'.join(st.session_state.story))

with col2:
    # Help and Suggestions Box
    st.header("Need Some Help?")
    if st.button("Get Vocabulary Suggestions"):
        vocab_suggestions = get_vocabulary_suggestions()
        st.write("Try using these words in your next sentence:")
        for word in vocab_suggestions:
            st.write(f"- {word}")

    if st.button("Get a Creative Prompt"):
        creative_prompt = get_creative_prompt()
        st.write(creative_prompt)

    # Rewards and Progress Box
    st.header("Achievements")
    if st.session_state.badges:
        for badge in st.session_state.badges:
            st.success(f"🏆 {badge}")
    else:
        st.write("No achievements yet. Keep writing to earn badges!")

    # Story Saving
    if st.button("Save Story"):
        story_data = {
            'level': st.session_state.level,
            'date': datetime.datetime.now().isoformat(),
            'story': st.session_state.story,
            'achievements': st.session_state.badges,
            'total_words': st.session_state.total_words,
            'unique_words': list(st.session_state.unique_words),
        }
        story_json = json.dumps(story_data, indent=4)
        st.download_button(
            label="Download Story",
            data=story_json,
            file_name='joystory.json',
            mime='application/json',
        )

# Sidebar for selecting language level
with st.sidebar:
    st.header("Settings")
    level = st.radio(
        "Select your English level:",
        ('Beginner', 'Intermediate', 'Advanced'),
        index=['Beginner', 'Intermediate', 'Advanced'].index(st.session_state.level)
    )
    st.session_state.level = level
    if st.button("Start New Story"):
        st.session_state.story = []
        st.session_state.achievements = []
        st.session_state.unique_words = set()
        st.session_state.total_words = 0
        st.session_state.badges = []
        st.experimental_rerun()