import streamlit as st
import openai
from datetime import datetime
import json
from collections import defaultdict

# Custom CSS for layout
st.markdown("""
<style>
    .box-container {
        background-color: #ffffff;
        border-radius: 10px;
        padding: 20px;
        margin: 10px 0;
        border: 2px solid #eee;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
    }
    
    .story-box {
        max-height: 300px;
        overflow-y: auto;
        padding: 15px;
        background-color: #f8f9fa;
        border-radius: 8px;
    }
    
    .input-box {
        background-color: #ffffff;
        padding: 15px;
        border-radius: 8px;
    }
    
    .help-box {
        background-color: #f0f7ff;
        padding: 15px;
        border-radius: 8px;
    }
    
    .rewards-box {
        background-color: #fff9f0;
        padding: 15px;
        border-radius: 8px;
    }
    
    .stButton button {
        width: 100%;
        border-radius: 20px;
        margin: 5px 0;
    }
    
    .story-text {
        padding: 10px;
        border-radius: 5px;
        margin: 5px 0;
    }
    
    .ai-text {
        background-color: #e3f2fd;
    }
    
    .player-text {
        background-color: #f0f4c3;
    }
    
    .prompt-text {
        color: #666;
        font-style: italic;
        margin: 5px 0;
    }
    
    .badge-container {
        display: inline-block;
        margin: 5px;
        padding: 10px;
        background-color: #ffd700;
        border-radius: 50%;
        text-align: center;
    }
</style>
""", unsafe_allow_html=True)

# Initialize OpenAI
openai.api_key = st.secrets["OPENAI_API_KEY"]

# Initialize session state variables
if 'story' not in st.session_state:
    st.session_state.story = []
if 'current_level' not in st.session_state:
    st.session_state.current_level = 'beginner'
if 'story_started' not in st.session_state:
    st.session_state.story_started = False
if 'first_turn' not in st.session_state:
    st.session_state.first_turn = None
if 'current_turn' not in st.session_state:
    st.session_state.current_turn = None
if 'badges' not in st.session_state:
    st.session_state.badges = defaultdict(int)
if 'vocabulary_used' not in st.session_state:
    st.session_state.vocabulary_used = set()

def get_ai_story_continuation(story_context, level):
    """Get story continuation from OpenAI API"""
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": f"""You are a creative storyteller helping a {level} level English student write a story. 

                Rules:
                1. READ the story context carefully
                2. ADD meaningful story elements (new events, character development, dialogue)
                3. AVOID generic transitions like 'suddenly' or 'then'
                4. ENSURE story progression and character development
                5. USE vocabulary appropriate for {level} level

                Current story elements:
                - Characters introduced: {extract_characters(story_context)}
                - Current setting: {extract_setting(story_context)}
                - Recent events: {extract_recent_events(story_context)}

                Provide JSON response with:
                1. continuation: Create an interesting next story event (1-2 sentences)
                2. creative_prompt: Ask about specific details/choices for the next part
                3. vocabulary_suggestions: 3 contextual words for the next part
                4. translation: Thai translation of continuation
                5. story_elements: New elements introduced (characters/places/objects)"""},
                {"role": "user", "content": f"Story context: {story_context}\nContinue the story creatively and maintain narrative flow."}
            ],
            temperature=0.7
        )
        
        result = json.loads(response.choices[0].message.content)
        
        # Validate story continuation
        if is_generic_response(result['continuation']):
            return get_ai_story_continuation(story_context, level)
            
        return result
        
    except Exception as e:
        st.error(f"Error with AI response: {e}")
        return create_fallback_response(story_context)

def extract_characters(story_context):
    """Extract character information from story"""
    try:
        char_response = openai.ChatCompletion.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "List all characters mentioned in the story."},
                {"role": "user", "content": story_context}
            ]
        )
        return char_response.choices[0].message.content
    except:
        return "Unknown characters"

def extract_setting(story_context):
    """Extract setting information from story"""
    try:
        setting_response = openai.ChatCompletion.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "Describe the current setting of the story."},
                {"role": "user", "content": story_context}
            ]
        )
        return setting_response.choices[0].message.content
    except:
        return "Unknown setting"

def extract_recent_events(story_context):
    """Extract recent events from story"""
    try:
        events_response = openai.ChatCompletion.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "Summarize the most recent events in the story."},
                {"role": "user", "content": story_context}
            ]
        )
        return events_response.choices[0].message.content
    except:
        return "Unknown events"

def is_generic_response(continuation):
    """Check if the response is too generic"""
    generic_phrases = [
        "suddenly", "then", "something happened",
        "unexpectedly", "the story continues"
    ]
    return any(phrase in continuation.lower() for phrase in generic_phrases)

def create_fallback_response(story_context):
    """Create contextual fallback response"""
    # Extract last event or setting from story_context
    last_event = story_context.split('.')[-2] if len(story_context.split('.')) > 1 else story_context
    
    return {
        "continuation": f"The characters decided to explore further into the story.",
        "creative_prompt": "What interesting detail should we add to the scene?",
        "vocabulary_suggestions": ["explore", "adventure", "discover"],
        "translation": "ตัวละครตัดสินใจสำรวจเรื่องราวต่อไป",
        "story_elements": {"new_elements": ["exploration"]}
    }

def validate_story_quality(story_context, new_continuation):
    """Validate that the story maintains quality and coherence"""
    try:
        validation = openai.ChatCompletion.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": """Validate story coherence and quality.
                Return JSON with:
                1. is_coherent: boolean
                2. reason: string explanation if not coherent"""},
                {"role": "user", "content": f"""Previous: {story_context}
                New continuation: {new_continuation}
                Is this a logical and non-repetitive continuation?"""}
            ]
        )
        
        result = json.loads(validation.choices[0].message.content)
        return result['is_coherent'], result.get('reason', '')
        
    except Exception:
        return True, ""  # Default to accepting if validation fails

def start_story(level):
    """Get the first sentence from AI"""
    # Define level-specific starting prompts
    level_prompts = {
        'beginner': {
            'guidance': "Start with a simple, engaging sentence about a person, animal, or place.",
            'examples': "Example: 'A happy dog lived in a blue house.'"
        },
        'intermediate': {
            'guidance': "Start with an interesting situation using compound sentences.",
            'examples': "Example: 'On a sunny morning, Sarah found a mysterious letter in her garden.'"
        },
        'advanced': {
            'guidance': "Start with a sophisticated hook using complex sentence structure.",
            'examples': "Example: 'Deep in the ancient forest, where shadows danced beneath ancient trees, a peculiar sound echoed through the mist.'"
        }
    }

    try:
        response = openai.ChatCompletion.create(
            model="gpt-4o-mini",  # หรือ model ที่ถูกต้องตามที่คุณใช้
            messages=[
                {"role": "system", "content": f"""You are starting a story for a {level} level English student.
                {level_prompts[level]['guidance']}
                {level_prompts[level]['examples']}
                
                Provide a JSON response with:
                1. opening: An engaging opening sentence (matching the specified level)
                2. creative_prompt: A short question (3-8 words) to guide the student's first sentence
                3. vocabulary_suggestions: 3 words they might use in their response (appropriate for {level} level)
                4. translation: Thai translation of the opening sentence
                
                Make the opening engaging but appropriate for the student's level."""},
                {"role": "user", "content": "Please start a new story."}
            ],
            temperature=0.7
        )
        return json.loads(response.choices[0].message.content)
    except Exception as e:
        st.error(f"Error with AI response: {e}")
        return {
            "opening": "Once upon a time...",
            "creative_prompt": "Who is our main character?",
            "vocabulary_suggestions": ["brave", "curious", "friendly"],
            "translation": "กาลครั้งหนึ่ง..."
        }

def main():
    st.title("🎨 Interactive Story Adventure")
    
    # Initial setup (level selection and first turn choice)
    if not st.session_state.story_started:
        with st.container():
            st.markdown('<div class="box-container">', unsafe_allow_html=True)
            st.write("### Welcome to Story Adventure!")
            col1, col2 = st.columns(2)
            with col1:
                level = st.selectbox(
                    "Choose your English level:",
                    ['beginner', 'intermediate', 'advanced']
                )
            with col2:
                first_turn = st.radio(
                    "Who should start the story?",
                    ['AI', 'Player']
                )
            if st.button("Start Story!", key="start_button"):
                st.session_state.current_level = level
                st.session_state.first_turn = first_turn
                st.session_state.current_turn = first_turn
                st.session_state.story_started = True
                if first_turn == 'AI':
                    response = start_story(level)
                    st.session_state.story.append({
                        'author': 'AI',
                        'text': response['opening'],
                        'prompt': response['creative_prompt'],
                        'vocabulary': response['vocabulary_suggestions'],
                        'timestamp': datetime.now().isoformat()
                    })
                    st.session_state.current_turn = 'Player'
                st.rerun()
            st.markdown('</div>', unsafe_allow_html=True)
    
    # Main story interface with 4-box layout
    if st.session_state.story_started:
        # Box 1: Story Display
        with st.container():
            st.markdown('<div class="box-container story-box">', unsafe_allow_html=True)
            st.write("### 📖 Our Story So Far:")
            
            for entry in st.session_state.story:
                # แสดงเนื้อเรื่อง
                css_class = "ai-text" if entry['author'] == 'AI' else "player-text"
                st.markdown(f'<div class="story-text {css_class}">{entry["text"]}</div>', 
                           unsafe_allow_html=True)
                
                # ถ้าเป็น AI response
                if entry['author'] == 'AI':
                    # แสดงคำแปล
                    if 'translation' in entry:
                        st.markdown(f'''
                            <div class="translation-box">
                                🇹🇭 {entry["translation"]}
                            </div>
                        ''', unsafe_allow_html=True)
                    
                    # แสดงคำถามกระตุ้นความคิด
                    if 'prompt' in entry:
                        st.markdown(f'''
                            <div class="prompt-box">
                                💭 {entry["prompt"]}
                            </div>
                        ''', unsafe_allow_html=True)
                    
                    # แสดงคำศัพท์แนะนำ
                    if 'vocabulary' in entry:
                        st.markdown(f'''
                            <div class="vocabulary-box">
                                📚 Try using: {', '.join(entry["vocabulary"])}
                            </div>
                        ''', unsafe_allow_html=True)
                    
                    # แสดงองค์ประกอบใหม่ในเรื่อง
                    if 'story_elements' in entry:
                        st.markdown(f'''
                            <div class="elements-box">
                                🎭 New elements: {', '.join(entry["story_elements"]["new_elements"])}
                            </div>
                        ''', unsafe_allow_html=True)

        # Box 2: Writing Area
        if st.session_state.current_turn == 'Player':
            with st.container():
                st.markdown('<div class="box-container input-box">', unsafe_allow_html=True)
                st.write("### ✏️ Your Turn!")
                user_input = st.text_area("Write your next sentence:", key='story_input')
                if st.button("✨ Submit", key="submit_button"):
                    if user_input:
                        process_player_input(user_input)
                        st.rerun()
                st.markdown('</div>', unsafe_allow_html=True)
        else:
            st.info("🤖 AI is thinking...")  # แสดงสถานะเมื่อเป็น turn ของ AI
        
        # Box 3: Help Buttons
        with st.container():
            st.markdown('<div class="box-container help-box">', unsafe_allow_html=True)
            st.write("### 💡 Need Help?")
            col1, col2, col3 = st.columns(3)
            with col1:
                if st.button("📚 Vocabulary Help"):
                    show_vocabulary_help()
            with col2:
                if st.button("❓ Writing Tips"):
                    show_writing_tips()
            with col3:
                if st.button("🎯 Story Ideas"):
                    show_story_ideas()
            st.markdown('</div>', unsafe_allow_html=True)
        
        # Box 4: Rewards and Progress
        with st.container():
            st.markdown('<div class="box-container rewards-box">', unsafe_allow_html=True)
            st.write("### 🏆 Your Achievements")
            show_achievements()
            col1, col2 = st.columns(2)
            with col1:
                if st.button("📥 Save Story"):
                    save_story()
            with col2:
                if st.button("🔄 Start New Story"):
                    reset_story()
                    st.rerun()
            st.markdown('</div>', unsafe_allow_html=True)

def process_player_input(user_input):
    try:
        # Get story context for AI
        story_context = " ".join([entry['text'] for entry in st.session_state.story])
        
        # Add player's sentence first
        st.session_state.story.append({
            'author': 'Player',
            'text': user_input,
            'timestamp': datetime.now().isoformat()
        })
        
        # Get AI's continuation immediately after player's input
        ai_response = get_ai_story_continuation(
            story_context + "\n" + user_input,
            st.session_state.current_level
        )

        # Validate the continuation
        is_coherent, reason = validate_story_quality(story_context, ai_response['continuation'])
        
        if not is_coherent:
            # Try getting a new response
            ai_response = get_ai_story_continuation(story_context + "\n" + user_input, st.session_state.current_level)
        
        # Add AI's response to story
        st.session_state.story.append({
            'author': 'AI',
            'text': ai_response['continuation'],
            'prompt': ai_response['creative_prompt'],
            'vocabulary': ai_response['vocabulary_suggestions'],
            'translation': ai_response.get('translation', ''),
            'timestamp': datetime.now().isoformat()
        })
        
        # Get feedback for player's writing
        feedback_response = openai.ChatCompletion.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": f"""You are an English teacher helping a {st.session_state.current_level} level student.
                Analyze their sentence and provide feedback in JSON format with:
                1. grammar_check: List of grammar issues found (if any)
                2. spelling_check: List of spelling issues found (if any)
                3. improvement_suggestions: Specific suggestions for improvement
                4. positive_feedback: What they did well
                Be encouraging but thorough in your feedback."""},
                {"role": "user", "content": f"Student's sentence: {user_input}"}
            ]
        )
        feedback = json.loads(feedback_response.choices[0].message.content)
        
        # Update vocabulary used
        words = set(user_input.lower().split())
        st.session_state.vocabulary_used.update(words)
        
        # Check for achievements
        check_achievements(user_input, feedback)
        
        # Show feedback to player
        display_feedback(feedback)
        
    except Exception as e:
        st.error(f"Error processing input: {e}")
        st.session_state.story.append({
            'author': 'Player',
            'text': user_input,
            'timestamp': datetime.now().isoformat()
        })

def check_achievements(user_input, feedback):
    """Check and award achievements based on player's writing"""
    # Story length achievements
    if len(st.session_state.story) >= 10:
        award_badge("Storyteller")
    if len(st.session_state.story) >= 20:
        award_badge("Master Storyteller")
    
    # Vocabulary achievements
    if len(st.session_state.vocabulary_used) >= 20:
        award_badge("Vocabulary Explorer")
    if len(st.session_state.vocabulary_used) >= 50:
        award_badge("Word Master")
    
    # Writing quality achievements
    if not feedback['grammar_check'] and not feedback['spelling_check']:
        award_badge("Perfect Writing")
    
    # Creativity achievements
    if len(user_input.split()) >= 15:
        award_badge("Detailed Writer")

def award_badge(badge_name):
    """Award a new badge if not already earned"""
    if badge_name not in st.session_state.badges:
        st.session_state.badges[badge_name] = 1
        st.balloons()
        st.success(f"🏆 New Achievement Unlocked: {badge_name}!")
    else:
        st.session_state.badges[badge_name] += 1

def display_feedback(feedback):
    """Display writing feedback to the player"""
    with st.container():
        # Show positive feedback first
        st.success(f"👏 {feedback['positive_feedback']}")
        
        # Show any grammar or spelling issues
        if feedback['grammar_check']:
            st.warning("📝 Grammar points to consider:")
            for issue in feedback['grammar_check']:
                st.write(f"- {issue}")
        
        if feedback['spelling_check']:
            st.warning("✍️ Spelling points to consider:")
            for issue in feedback['spelling_check']:
                st.write(f"- {issue}")
        
        # Show improvement suggestions
        if feedback['improvement_suggestions']:
            st.info("💡 Suggestions for next time:")
            for suggestion in feedback['improvement_suggestions']:
                st.write(f"- {suggestion}")

def show_vocabulary_help():
    """Show vocabulary suggestions and meanings"""
    if st.session_state.story:
        last_entry = st.session_state.story[-1]
        if 'vocabulary' in last_entry:
            # Get word meanings and examples
            try:
                word_info = openai.ChatCompletion.create(
                    model="gpt-4o-mini",
                    messages=[
                        {"role": "system", "content": f"""For each word, provide:
                        1. Simple definition
                        2. Thai translation
                        3. Example sentence
                        Format in JSON with word as key."""},
                        {"role": "user", "content": f"Words: {', '.join(last_entry['vocabulary'])}"}
                    ]
                )
                word_details = json.loads(word_info.choices[0].message.content)
                
                st.write("### 📚 Suggested Words:")
                for word in last_entry['vocabulary']:
                    with st.expander(f"🔤 {word}"):
                        if word in word_details:
                            details = word_details[word]
                            st.write(f"**Meaning:** {details['definition']}")
                            st.write(f"**แปลไทย:** {details['thai']}")
                            st.write(f"**Example:** _{details['example']}_")
            except Exception as e:
                # Fallback to simple display
                st.write("Try using these words:")
                for word in last_entry['vocabulary']:
                    st.markdown(f"- *{word}*")

def show_writing_tips():
    """Show writing tips based on current level and story context"""
    if st.session_state.story:
        try:
            # Get personalized writing tips
            tips_response = openai.ChatCompletion.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": f"""You are helping a {st.session_state.current_level} level English student.
                    Provide 3 specific writing tips based on their level and story context.
                    Include examples for each tip.
                    Keep tips short and friendly."""},
                    {"role": "user", "content": f"Story so far: {' '.join([entry['text'] for entry in st.session_state.story])}"}
                ]
            )
            
            st.write("### ✍️ Writing Tips")
            tips = tips_response.choices[0].message.content.split('\n')
            for tip in tips:
                if tip.strip():
                    st.info(tip)
            
            # Add quick reference examples
            st.write("### 🎯 Quick Examples:")
            col1, col2 = st.columns(2)
            with col1:
                st.write("**Good:**")
                st.success("The happy dog played in the garden.")
            with col2:
                st.write("**Better:**")
                st.success("The excited golden retriever chased butterflies in the colorful garden.")
                
        except Exception as e:
            # Fallback tips
            st.write("### General Writing Tips:")
            st.write("1. Use descriptive words")
            st.write("2. Keep your sentences clear")
            st.write("3. Try to be creative")

def show_story_ideas():
    """Generate creative story ideas based on current context"""
    if st.session_state.story:
        try:
            # Get creative prompts
            ideas_response = openai.ChatCompletion.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": f"""Generate 3 creative story direction ideas for a {st.session_state.current_level} student.
                    Ideas should follow from the current story.
                    Make suggestions exciting but achievable for their level."""},
                    {"role": "user", "content": f"Story so far: {' '.join([entry['text'] for entry in st.session_state.story])}"}
                ]
            )
            
            st.write("### 💭 Story Ideas")
            ideas = ideas_response.choices[0].message.content.split('\n')
            for i, idea in enumerate(ideas, 1):
                if idea.strip():
                    with st.expander(f"Idea {i}"):
                        st.write(idea)
                        
            # Add story element suggestions
            st.write("### 🎨 Try adding:")
            elements = ["a new character", "a surprise event", "a funny moment", "a problem to solve"]
            cols = st.columns(len(elements))
            for col, element in zip(cols, elements):
                with col:
                    st.button(element, key=f"element_{element}")
                    
        except Exception as e:
            # Fallback ideas
            st.write("### Story Ideas:")
            st.write("1. Introduce a new character")
            st.write("2. Add an unexpected event")
            st.write("3. Create a problem to solve")

def show_achievements():
    """Display achievements and badges with animations"""
    if st.session_state.badges:
        st.write("### 🏆 Your Achievements")
        cols = st.columns(3)
        for i, (badge, count) in enumerate(st.session_state.badges.items()):
            with cols[i % 3]:
                st.markdown(f'''
                    <div class="badge-container" style="animation: bounce 1s infinite;">
                        <div style="font-size: 2em;">🌟</div>
                        <div style="font-weight: bold;">{badge}</div>
                        <div>x{count}</div>
                    </div>
                ''', unsafe_allow_html=True)
        
        # Add progress towards next badge
        st.write("### 📈 Progress to Next Badge:")
        current_words = len(st.session_state.vocabulary_used)
        next_milestone = (current_words // 10 + 1) * 10
        progress = current_words / next_milestone
        st.progress(progress)
        st.write(f"Use {next_milestone - current_words} more unique words for next badge!")

def save_story():
    """Save story with additional information"""
    if st.session_state.story:
        # Prepare story text with formatting
        story_parts = []
        story_parts.append("=== My English Story Adventure ===\n")
        story_parts.append(f"Level: {st.session_state.current_level}")
        story_parts.append(f"Date: {datetime.now().strftime('%Y-%m-%d')}\n")
        
        # Add story content
        for entry in st.session_state.story:
            author = "👤 You:" if entry['author'] == 'Player' else "🤖 AI:"
            story_parts.append(f"{author} {entry['text']}")
        
        # Add achievements
        story_parts.append("\n=== Achievements ===")
        for badge, count in st.session_state.badges.items():
            story_parts.append(f"🌟 {badge}: x{count}")
        
        # Create the full story text
        story_text = "\n".join(story_parts)
        
        # Offer download options
        col1, col2 = st.columns(2)
        with col1:
            st.download_button(
                label="📝 Download Story (Text)",
                data=story_text,
                file_name="my_story.txt",
                mime="text/plain"
            )
        with col2:
            # Create JSON version with more details
            story_data = {
                "metadata": {
                    "level": st.session_state.current_level,
                    "date": datetime.now().isoformat(),
                    "vocabulary_used": list(st.session_state.vocabulary_used),
                    "achievements": dict(st.session_state.badges)
                },
                "story": st.session_state.story
            }
            st.download_button(
                label="💾 Download Progress Report (JSON)",
                data=json.dumps(story_data, indent=2),
                file_name="story_progress.json",
                mime="application/json"
            )

def reset_story():
    """Reset all story-related session state"""
    if st.button("🔄 Start New Story", key="reset_button"):
        st.session_state.story = []
        st.session_state.story_started = False
        st.session_state.first_turn = None
        st.session_state.current_turn = None
        st.session_state.vocabulary_used = set()
        # Keep badges for long-term progress
        st.balloons()
        st.success("Ready for a new adventure! Your achievements have been saved.")
        st.rerun()


if __name__ == "__main__":
    main()