Update appback.py3

appback.py3

@@ -549,8 +549,14 @@ def init_session_state():
         st.session_state.clear_input = False
     if 'text_input' not in st.session_state:
         st.session_state.text_input = ""
-    if 'story_input_area' not in st.session_state:
-        st.session_state.story_input_area = ""
+    if 'ending_mode' not in st.session_state:
+        st.session_state.ending_mode = False
+    if 'sentences_to_end' not in st.session_state:
+        st.session_state.sentences_to_end = 0
+    if 'ending_type' not in st.session_state:
+        st.session_state.ending_type = None
+    if 'story_completed' not in st.session_state:
+        st.session_state.story_completed = False
 
 def init_theme_state():
     """Initialize theme-specific state variables"""
@@ -746,9 +752,9 @@ def generate_story_continuation(user_input: str, level: str) -> str:
     }
 
     try:
-        # Get recent story context
+        # Get recent story context (last 5 entries for more coherence)
         story_context = '\n'.join([
-            entry['content'] for entry in st.session_state.story[-3:]
+            entry['content'] for entry in st.session_state.story[-5:]
         ]) if st.session_state.story else "Story just started"
         
         # Create prompt
@@ -811,6 +817,7 @@ def generate_story_continuation(user_input: str, level: str) -> str:
         logging.error(f"Error generating story continuation: {str(e)}")
         return "I'm having trouble continuing the story. Please try again."
 
+
 def generate_dynamic_story_starter(theme_id: str, level: str) -> Dict[str, str]:
     """
     Dynamically generate a story starter based on theme and level.
@@ -848,47 +855,6 @@ def generate_dynamic_story_starter(theme_id: str, level: str) -> Dict[str, str]:
             'th': 'กาลครั้งหนึ่ง...'
         }
 
-def handle_application_error(error: Exception):
-    """Handle application-wide errors"""
-    logging.error(f"Application error: {str(error)}")
-    
-    error_message = f"""
-        <div style="
-            background-color: #ffebee;
-            padding: 20px;
-            border-radius: 10px;
-            border-left: 4px solid #c62828;
-            margin: 20px 0;
-        ">
-            <h3 style="color: #c62828; margin: 0 0 10px 0;">
-                ⚠️ เกิดข้อผิดพลาดในระบบ
-            </h3>
-            <p style="color: #333; margin: 0;">
-                กรุณาลองใหม่อีกครั้ง หรือติดต่อผู้ดูแลระบบ
-            </p>
-            <div style="
-                margin-top: 10px;
-                padding: 10px;
-                background-color: #f5f5f5;
-                border-radius: 5px;
-                font-family: monospace;
-                font-size: 0.9em;
-            ">
-                Error: {str(error)}
-            </div>
-        </div>
-    """
-    
-    st.markdown(error_message, unsafe_allow_html=True)
-    
-    # Show technical details in expander
-    with st.expander("📋 รายละเอียดข้อผิดพลาด (สำหรับผู้ดูแลระบบ)"):
-        st.code(f"""
-        Error Type: {type(error).__name__}
-        Error Message: {str(error)}
-        Timestamp: {datetime.now().isoformat()}
-        """)
-
 def update_points(is_correct_first_try: bool):
     """อัพเดตคะแนนตามผลการเขียน"""
     try:
@@ -1001,7 +967,7 @@ def update_achievements():
         return False
 
 def provide_feedback(text: str, level: str) -> Dict[str, str]:
-    """Provide feedback on the user's writing with appropriate level"""
+    """Provide more detailed feedback on the user's writing with appropriate level"""
     
     level_prompts = {
         'Beginner': """
@@ -1011,6 +977,7 @@ def provide_feedback(text: str, level: str) -> Dict[str, str]:
         - Basic vocabulary
         - Capitalization and periods
         Provide very simple, encouraging feedback in Thai.
+        Explain any mistakes and how to correct them.
         """,
         'Intermediate': """
         Focus on:
@@ -1020,6 +987,7 @@ def provide_feedback(text: str, level: str) -> Dict[str, str]:
         - Basic punctuation
         - Simple conjunctions
         Provide moderately detailed feedback in Thai.
+        Explain any mistakes and suggest improvements.
         """,
         'Advanced': """
         Focus on:
@@ -1029,6 +997,7 @@ def provide_feedback(text: str, level: str) -> Dict[str, str]:
         - All punctuation
         - Style and flow
         Provide comprehensive feedback in Thai.
+        Explain any mistakes and provide examples for improvement.
         """
     }
 
@@ -1059,8 +1028,8 @@ def provide_feedback(text: str, level: str) -> Dict[str, str]:
                             "content": f"Review this sentence: {text}"
                         }
                     ],
-                    max_tokens=200,
-                    temperature=0.3
+                    max_tokens=300,  # Increased from 200
+                    temperature=0.5  # Adjusted for balanced variety
                 )
                 
                 # Parse response
@@ -1091,6 +1060,7 @@ def provide_feedback(text: str, level: str) -> Dict[str, str]:
             "difficulty_score": 5
         }
 
+
 def get_vocabulary_suggestions(context: str = "", level: str = DEFAULT_LEVEL) -> List[str]:
     """Get contextual vocabulary suggestions with Thai translations"""
     try:
@@ -1488,7 +1458,7 @@ def show_story():
         for idx, entry in enumerate(st.session_state.story):
             if entry['role'] == 'AI':
                 if entry.get('is_starter'):
-                    # Story Starter
+                    # Story Starter (คงเดิม)
                     st.markdown(f"""
                         <div style="
                             background-color: #f0f7ff; 
@@ -1509,8 +1479,39 @@ def show_story():
                             </div>
                         </div>
                     """, unsafe_allow_html=True)
+                elif entry.get('is_final'):  # เพิ่มเงื่อนไขสำหรับประโยคจบ
+                    # Final Ending Display
+                    st.markdown(f"""
+                        <div style="
+                            background-color: #e8f5e9;
+                            padding: 25px;
+                            border-radius: 10px;
+                            margin: 15px 0;
+                            border: 2px solid #4caf50;
+                            text-align: center;
+                            box-shadow: 0 4px 8px rgba(0,0,0,0.1);
+                        ">
+                            <div style="font-size: 1.3em; color: #2e7d32; margin-bottom: 15px;">
+                                🎭 จบเรื่อง | The End
+                            </div>
+                            <div style="color: #333; font-size: 1.1em; font-style: italic;">
+                                {entry['content']}
+                            </div>
+                            <div style="margin-top: 15px; padding-top: 15px; border-top: 1px solid #a5d6a7;">
+                                <span style="color: #666;">
+                                    ✨ รูปแบบการจบ: {st.session_state.ending_type}
+                                </span>
+                            </div>
+                        </div>
+                    """, unsafe_allow_html=True)
                 else:
-                    # AI Response
+                    # Regular AI Response with ending mode indicator
+                    ending_info = ""
+                    # แสดงจำนวนประโยคที่เหลือเฉพาะเมื่ออยู่ในโหมดจบเรื่อง
+                    if (st.session_state.get('ending_mode') and 
+                        'remaining_sentences' in entry):
+                        ending_info = f"🎭 เหลืออีก {entry['remaining_sentences']} ประโยค"
+                    
                     st.markdown(f"""
                         <div style="
                             background-color: #f8f9fa;
@@ -1520,7 +1521,10 @@ def show_story():
                             border-left: 4px solid #4caf50;
                         ">
                             <div style="color: #2e7d32;">
-                                🤖 AI: {entry['content']}
+                                <div style="display: flex; justify-content: space-between; align-items: center;">
+                                    <span>🤖 AI: {entry['content']}</span>
+                                    <span style="color: #f57c00; font-size: 0.9em;">{ending_info}</span>
+                                </div>
                             </div>
                         </div>
                     """, unsafe_allow_html=True)
@@ -1531,6 +1535,12 @@ def show_story():
                 bg_color = "#e8f5e9" if entry.get('is_correct') else "#fff"
                 border_color = "#4caf50" if entry.get('is_correct') else "#1e88e5"
                 
+                ending_info = ""
+                # แสดงจำนวนประโยคที่เหลือเฉพาะเมื่ออยู่ในโหมดจบเรื่อง
+                if (st.session_state.get('ending_mode') and 
+                    'remaining_sentences' in entry):
+                    ending_info = f"🎭 เหลืออีก {entry['remaining_sentences']} ประโยค"
+                
                 st.markdown(f"""
                     <div style="
                         background-color: {bg_color};
@@ -1540,206 +1550,595 @@ def show_story():
                         border-left: 4px solid {border_color};
                     ">
                         <div style="color: #333;">
-                            👤 You: {status_icon} {entry['content']}
+                            <div style="display: flex; justify-content: space-between; align-items: center;">
+                                <span>👤 You: {status_icon} {entry['content']}</span>
+                                <span style="color: #f57c00; font-size: 0.9em;">{ending_info}</span>
+                            </div>
                         </div>
                         {f'<div style="font-size: 0.9em; color: #666; margin-top: 5px;">{entry.get("feedback", "")}</div>' if entry.get('feedback') else ''}
                     </div>
                 """, unsafe_allow_html=True)
-
-def show_settings_hint():
-    """Show floating hint for settings sidebar"""
-    st.markdown("""
-        <div style="
-            position: fixed;
-            left: 10px;
-            top: 50%;
-            transform: translateY(-50%);
-            background-color: #1e88e5;
-            color: white;
-            padding: 10px;
-            border-radius: 0 8px 8px 0;
-            box-shadow: 2px 2px 10px rgba(0,0,0,0.1);
-            cursor: pointer;
-            z-index: 1000;
-            animation: pulse 2s infinite;
-        ">
-            <div style="writing-mode: vertical-rl; text-orientation: mixed;">
-                ⚙️ ตั้งค่าระดับความยาก
-            </div>
-        </div>
         
-        <style>
-        @keyframes pulse {
-            0% { transform: translateY(-50%) scale(1); }
-            50% { transform: translateY(-50%) scale(1.05); }
-            100% { transform: translateY(-50%) scale(1); }
-        }
-        </style>
-    """, unsafe_allow_html=True)
-
-def show_feedback_section():
-    """Display writing feedback section"""
-    if not st.session_state.feedback:
-        return
-
-    st.markdown("""
-        <div style="margin-bottom: 20px;">
-            <div class="thai-eng">
-                <div class="thai">📝 คำแนะนำจากครู</div>
-                <div class="eng">Writing Feedback</div>
-            </div>
-        </div>
-    """, unsafe_allow_html=True)
-    
-    feedback_data = st.session_state.feedback
-    if isinstance(feedback_data, dict) and feedback_data.get('has_errors'):
-        # Show error feedback
-        st.markdown(f"""
-            <div style="
-                background-color: #fff3e0;
-                padding: 20px;
-                border-radius: 10px;
-                border-left: 4px solid #ff9800;
-                margin: 10px 0;
-            ">
-                <p style="color: #e65100; margin-bottom: 15px;">
-                    {feedback_data['feedback']}
-                </p>
+        # Show completion message if story is done
+        if st.session_state.get('story_completed'):
+            st.markdown("""
                 <div style="
-                    background-color: #fff;
-                    padding: 15px;
-                    border-radius: 8px;
-                    margin-top: 10px;
+                    background-color: #f3e5f5;
+                    padding: 20px;
+                    border-radius: 10px;
+                    margin: 20px 0;
+                    text-align: center;
+                    border: 2px dashed #9c27b0;
                 ">
-                    <p style="color: #666; margin-bottom: 5px;">
-                        ประโยคที่ถูกต้อง:
-                    </p>
-                    <p style="
-                        color: #2e7d32;
-                        font-weight: 500;
-                        font-size: 1.1em;
-                        margin: 0;
-                    ">
-                        {feedback_data['corrected']}
+                    <h2 style="color: #9c27b0; margin-bottom: 10px;">
+                        🎉 ยินดีด้วย! คุณเขียนเรื่องราวจบสมบูรณ์แล้ว
+                    </h2>
+                    <p style="color: #666;">
+                        คุณสามารถบันทึกเรื่องราวหรือเริ่มเรื่องใหม่ได้
                     </p>
                 </div>
-            </div>
-        """, unsafe_allow_html=True)
-        
-        # Correction button
-        if st.button(
-            "✍️ แก้ไขประโยคให้ถูกต้อง",
-            key="correct_button",
-            help="คลิกเพื่อแก้ไขประโยคให้ถูกต้องตามคำแนะนำ"
-        ):
-            last_user_entry_idx = next(
-                (i for i, entry in reversed(list(enumerate(st.session_state.story))) 
-                 if entry['role'] == 'You'),
-                None
-            )
-            if last_user_entry_idx is not None:
-                apply_correction(last_user_entry_idx, feedback_data['corrected'])
-                st.rerun()
-    else:
-        # Show success feedback
+            """, unsafe_allow_html=True)
+
+def show_story_progress():
+    """Display story progress metrics"""
+    if st.session_state.story:
+        total_sentences = len(st.session_state.story)
         st.markdown(f"""
             <div style="
-                background-color: #e8f5e9;
-                padding: 20px;
+                background-color: #f0f7ff;
+                padding: 15px;
                 border-radius: 10px;
-                border-left: 4px solid #4caf50;
                 margin: 10px 0;
             ">
-                <p style="color: #2e7d32; margin: 0;">
-                    {feedback_data.get('feedback', '✨ เขียนได้ถูกต้องแล้วค่ะ!')}
-                </p>
+                <div style="margin-bottom: 10px;">
+                    📊 ความยาวเรื่อง: {total_sentences} ประโยค
+                </div>
+                <div class="progress-bar">
+                    <div class="progress-bar-fill" style="width: {min(total_sentences * 10, 100)}%;"></div>
+                </div>
+                <div style="font-size: 0.9em; color: #666; margin-top: 5px;">
+                    เรื่องควรยาว 10-20 ประโยค เพื่อความสมบูรณ์
+                </div>
             </div>
         """, unsafe_allow_html=True)
 
-def show_writing_tools():
-    """Display writing tools section"""
-    with st.expander("✨ เครื่องมือช่วยเขียน | Writing Tools"):
-        col1, col2 = st.columns(2)
+def show_story_ending_options():
+    """Display story ending options and guidance"""
+    if len(st.session_state.story) >= 10:  # Show options when story is long enough
+        st.markdown("### 🎭 ต้องการจบเรื่องหรือไม่?")
         
-        with col1:
-            if st.button("🎯 ขอคำใบ้", use_container_width=True):
-                with st.spinner("กำลังสร้างคำใบ้..."):
-                    prompt = get_creative_prompt()
-                    st.markdown(f"""
-                        <div style="
-                            background-color: #f3e5f5;
-                            padding: 15px;
-                            border-radius: 8px;
-                            margin-top: 10px;
-                        ">
-                            <div style="color: #6a1b9a;">💭 {prompt['thai']}</div>
-                            <div style="color: #666; font-style: italic;">
-                                💭 {prompt['eng']}
-                            </div>
-                        </div>
-                    """, unsafe_allow_html=True)
+        # Display ending type options
+        ending_type = st.radio(
+            "เลือกวิธีจบเรื่อง:",
+            options=[
+                "Happy Ending - จบแบบมีความสุข",
+                "Mysterious Ending - จบแบบทิ้งท้ายให้คิดต่อ",
+                "Lesson Learned - จบแบบได้ข้อคิด",
+                "Surprise Ending - จบแบบพลิกความคาดหมาย"
+            ],
+            index=0,
+            help="เลือกรูปแบบการจบเรื่องที่คุณต้องการ"
+        )
         
-        with col2:
-            if st.button("📚 คำศัพท์แนะนำ", use_container_width=True):
-                with st.spinner("กำลังค้นหาคำศัพท์..."):
-                    vocab_suggestions = get_vocabulary_suggestions()
-                    st.markdown("#### 📚 คำศัพท์น่ารู้")
-                    for word in vocab_suggestions:
-                        st.markdown(f"""
-                            <div style="
-                                background-color: #f5f5f5;
-                                padding: 10px;
-                                border-radius: 5px;
-                                margin: 5px 0;
-                            ">
-                                • {word}
-                            </div>
-                        """, unsafe_allow_html=True)
+        if st.button("🎬 เริ่มจบเรื่อง", use_container_width=True):
+            st.session_state.ending_mode = True
+            st.session_state.ending_type = ending_type
+            st.session_state.sentences_to_end = 5
+            st.rerun()  # Using st.rerun() instead of st.experimental_rerun()
 
-def show_achievements():
-    """Display achievements and stats"""
-    with st.container():
-        # Points and Streak
-        st.markdown(f"""
-            <div style="
-                background-color: #e3f2fd;
-                padding: 20px;
-                border-radius: 10px;
-                margin-bottom: 20px;
-                text-align: center;
-            ">
-                <h3 style="color: #1e88e5; margin: 0;">
-                    🌟 คะแนนรวม: {st.session_state.points['total']}
-                </h3>
-                <p style="margin: 5px 0;">
-                    🔥 Streak ปัจจุบัน: {st.session_state.points['streak']} ประโยค
-                </p>
-                <p style="margin: 5px 0;">
-                    ⭐ Streak สูงสุด: {st.session_state.points['max_streak']} ประโยค
-                </p>
-            </div>
-        """, unsafe_allow_html=True)
+def handle_ending_mode(text: str):
+    """Handle story submission during ending mode"""
+    try:
+        # Get current remaining sentences
+        remaining = st.session_state.sentences_to_end
 
-        # Writing Stats
-        st.markdown("""
-            <div style="margin-bottom: 15px;">
-                <h3>📊 สถิติการเขียน</h3>
-            </div>
-        """, unsafe_allow_html=True)
+        # Get feedback
+        feedback_data = provide_feedback(text, st.session_state.level)
+        st.session_state.feedback = feedback_data
+        is_correct = not feedback_data.get('has_errors', False)
         
-        col1, col2 = st.columns(2)
-        with col1:
-            st.metric(
-                "ประโยคที่เขียนทั้งหมด", 
-                st.session_state.stats['total_sentences'],
-                help="จำนวนประโยคทั้งหมดที่คุณได้เขียน"
-            )
-            st.metric(
-                "ถูกต้องตั้งแต่แรก", 
-                st.session_state.stats['correct_first_try'],
-                help="จำนวนประโยคที่ถูกต้องโดยไม่ต้องแก้ไข"
+        # Add user's sentence
+        st.session_state.story.append({
+            "role": "You",
+            "content": text,
+            "is_corrected": False,
+            "is_correct": is_correct,
+            "timestamp": datetime.now().isoformat(),
+            "remaining_sentences": remaining
+        })
+        
+        # Update vocabulary
+        words = set(text.lower().split())
+        st.session_state.stats['vocabulary_used'].update(words)
+        
+        # Update points and achievements
+        update_points(is_correct)
+        update_achievements()
+        
+        # Decrease remaining count for user's sentence
+        st.session_state.sentences_to_end -= 1
+        remaining = st.session_state.sentences_to_end
+        
+        if remaining > 0:
+            # Use corrected text if there are errors
+            text_for_continuation = feedback_data['corrected'] if feedback_data.get('has_errors') else text
+            
+            # Generate AI's ending continuation
+            ai_response = generate_ending_continuation(
+                text_for_continuation,
+                ending_type=st.session_state.ending_type,
+                remaining_sentences=remaining,
+                level=st.session_state.level
+            )
+            
+            # Add AI response
+            st.session_state.story.append({
+                "role": "AI",
+                "content": ai_response,
+                "timestamp": datetime.now().isoformat(),
+                "remaining_sentences": remaining
+            })
+            
+            # Decrease remaining sentences
+            st.session_state.sentences_to_end -= 1
+            
+        else:
+            # Generate final ending
+            final_response = generate_final_ending(
+                st.session_state.story,
+                st.session_state.ending_type,
+                level=st.session_state.level
+            )
+            
+            # Add final AI response
+            st.session_state.story.append({
+                "role": "AI",
+                "content": final_response,
+                "timestamp": datetime.now().isoformat(),
+                "is_final": True
+            })
+            
+            # Mark story as complete
+            st.session_state.story_completed = True
+            
+        # Clear input and rerun
+        st.session_state.text_input = ""
+        st.session_state.clear_input = True
+        st.rerun()
+        
+    except Exception as e:
+        logging.error(f"Error in ending mode: {str(e)}")
+        st.error("เกิดข้อผิดพลาดในโหมดจบเรื่อง กรุณาลองใหม่อีกครั้ง")
+
+
+
+def generate_final_ending(story: List[dict], ending_type: str, level: str) -> str:
+    """Generate the final ending sentences based on the story and chosen ending type"""
+    try:
+        # Create a comprehensive prompt for the final ending
+        story_summary = " ".join([entry['content'] for entry in story[-5:]])  # Last 5 sentences
+        
+        ending_prompts = {
+            "Happy Ending": """
+                Create a final happy ending that:
+                - Brings joy and satisfaction
+                - Resolves the main story elements
+                - Ends on a positive note
+                - Use 2-3 sentences
+                """,
+            "Mysterious Ending": """
+                Create a final mysterious ending that:
+                - Leaves an intriguing question
+                - Creates a sense of wonder
+                - Maintains some mystery
+                - Use 2-3 sentences
+                """,
+            "Lesson Learned": """
+                Create a final ending that:
+                - Shows what was learned
+                - Provides a moral lesson
+                - Connects to the story's events
+                - Use 2-3 sentences
+                """,
+            "Surprise Ending": """
+                Create a final twist ending that:
+                - Provides an unexpected but logical conclusion
+                - Connects to previous story elements
+                - Creates a satisfying surprise
+                - Use 2-3 sentences
+                """
+        }
+
+        level_context = {
+            'Beginner': {
+                'instructions': """
+                Additional Guidelines:
+                - Use simple sentences
+                - Use Present Simple Tense
+                - Basic vocabulary
+                - 7-10 words per sentence maximum
+                """,
+                'max_tokens': 100,  # Increased to allow for longer endings
+                'temperature': 0.6
+            },
+            'Intermediate': {
+                'instructions': """
+                Additional Guidelines:
+                - Use 2-3 sentences
+                - Can use Present or Past Tense
+                - Keep each sentence under 15 words
+                - Grade-appropriate vocabulary
+                """,
+                'max_tokens': 120,  # Increased
+                'temperature': 0.7
+            },
+            'Advanced': {
+                'instructions': """
+                Additional Guidelines:
+                - Use 2-3 sentences
+                - Various tenses allowed
+                - Natural sentence length but keep overall response concise
+                - More sophisticated vocabulary and structures
+                """,
+                'max_tokens': 150,  # Increased
+                'temperature': 0.8
+            }
+        }
+
+        level_settings = level_context[level]
+
+        response = client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[
+                {
+                    "role": "system",
+                    "content": f"""
+                    {ending_prompts[ending_type]}
+                    
+                    {level_settings['instructions']}
+                    
+                    This must be the absolute final sentences of the story.
+                    Make it conclusive and satisfying.
+                    """
+                },
+                {
+                    "role": "user",
+                    "content": f"Story context:\n{story_summary}\n\nCreate the final ending:"
+                }
+            ],
+            max_tokens=level_settings['max_tokens'],
+            temperature=level_settings['temperature']
+        )
+        
+        return response.choices[0].message.content.strip()
+        
+    except Exception as e:
+        logging.error(f"Error generating final ending: {str(e)}")
+        return "And so, the story came to an end."
+
+
+
+def generate_ending_continuation(text: str, ending_type: str, remaining_sentences: int, level: str) -> str:
+    """Generate AI continuation focusing on story conclusion"""
+    try:
+        ending_prompts = {
+            "Happy Ending": """
+                Role: Story concluder aiming for a happy ending
+                Goal: Create a satisfying, positive conclusion
+                Rules:
+                - Build towards joy, success, or resolution
+                - Use uplifting and positive language
+                - Connect to previous story elements
+            """,
+            "Mysterious Ending": """
+                Role: Mystery writer creating intrigue
+                Goal: Leave readers thinking and wondering
+                Rules:
+                - Add subtle hints and clues
+                - Create atmospheric descriptions
+                - Leave some questions unanswered
+            """,
+            "Lesson Learned": """
+                Role: Moral story concluder
+                Goal: Incorporate meaningful life lessons
+                Rules:
+                - Connect actions to consequences
+                - Show character growth
+                - Express the moral naturally
+            """,
+            "Surprise Ending": """
+                Role: Plot twist creator
+                Goal: Deliver unexpected but satisfying conclusion
+                Rules:
+                - Plant subtle hints earlier
+                - Subvert expectations logically
+                - Maintain story coherence
+            """
+        }
+
+        level_context = {
+            'Beginner': {
+                'instructions': """
+                Additional Guidelines:
+                - Use only 1-2 VERY simple sentences
+                - Use Present Simple Tense only
+                - Basic vocabulary
+                - 5-7 words per sentence maximum
+                - Focus on clear, basic responses
+                """,
+                'max_tokens': 30,
+                'temperature': 0.6
+            },
+            'Intermediate': {
+                'instructions': """
+                Additional Guidelines:
+                - Use exactly 2 sentences maximum
+                - Can use Present or Past Tense
+                - Keep each sentence under 12 words
+                - Grade-appropriate vocabulary
+                - Add simple descriptions but stay concise
+                """,
+                'max_tokens': 40,
+                'temperature': 0.7
+            },
+            'Advanced': {
+                'instructions': """
+                Additional Guidelines:
+                - Use 2-3 sentences maximum
+                - Various tenses allowed
+                - Natural sentence length but keep overall response concise
+                - More sophisticated vocabulary and structures
+                - Create engaging responses that encourage creative continuation
+                """,
+                'max_tokens': 50,
+                'temperature': 0.8
+            }
+        }
+
+        # Get recent story context
+        story_context = '\n'.join([
+            entry['content'] for entry in st.session_state.story[-5:]
+        ]) if st.session_state.story else "Story just started"
+
+        level_settings = level_context[level]
+
+        response = client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[
+                {
+                    "role": "system",
+                    "content": f"""
+                    {ending_prompts[ending_type]}
+                    
+                    {level_settings['instructions']}
+                    
+                    Additional Rules:
+                    - You have {remaining_sentences} sentences left to conclude the story.
+                    - Each response should be maximum 2 sentences.
+                    - Build towards the finale naturally.
+                    - Ensure coherence with previous story events.
+                    """
+                },
+                {
+                    "role": "user",
+                    "content": f"Story context:\n{story_context}\nStudent's input:\n{text}\nContinue and work towards ending this story:"
+                }
+            ],
+            max_tokens=level_settings['max_tokens'],
+            temperature=level_settings['temperature']
+        )
+
+        return response.choices[0].message.content.strip()
+        
+    except Exception as e:
+        logging.error(f"Error generating ending continuation: {str(e)}")
+        return "The story moved towards its conclusion..."
+
+def complete_story():
+    """Handle story completion and celebration"""
+    st.balloons()  # แสดงเอฟเฟคฉลอง
+    
+    # สร้าง Story Summary
+    story_summary = generate_story_summary(st.session_state.story)
+    
+    # แสดงหน้าจบเรื่อง
+    st.markdown(f"""
+        <div style="
+            background-color: #e8f5e9;
+            padding: 20px;
+            border-radius: 10px;
+            text-align: center;
+            margin: 20px 0;
+        ">
+            <h2>🎉 ยินดีด้วย! คุณเขียนเรื่องราวจบสมบูรณ์แล้ว</h2>
+            
+            <div style="margin: 20px 0;">
+                <h3>📝 สรุปเรื่องราว</h3>
+                <p>{story_summary}</p>
+            </div>
+            
+            <div style="margin: 20px 0;">
+                <h3>🏆 ความสำเร็จ</h3>
+                <p>จำนวนประโยค: {len(st.session_state.story)}</p>
+                <p>คำศัพท์ที่ใช้: {len(st.session_state.stats['vocabulary_used'])}</p>
+                <p>ความแม่นยำ: {st.session_state.stats['accuracy_rate']:.1f}%</p>
+            </div>
+        </div>
+    """, unsafe_allow_html=True)
+    
+    # แสดงตัวเลือกหลังจบเรื่อง
+    col1, col2 = st.columns(2)
+    with col1:
+        if st.button("💾 บันทึกเรื่องราว", use_container_width=True):
+            save_completed_story()
+    with col2:
+        if st.button("🔄 เริ่มเรื่องใหม่", use_container_width=True):
+            reset_story()
+            st.rerun()
+
+def show_feedback_section():
+    """Display writing feedback section"""
+    if not st.session_state.feedback:
+        return
+
+    st.markdown("""
+        <div style="margin-bottom: 20px;">
+            <div class="thai-eng">
+                <div class="thai">📝 คำแนะนำจากครู</div>
+                <div class="eng">Writing Feedback</div>
+            </div>
+        </div>
+    """, unsafe_allow_html=True)
+    
+    feedback_data = st.session_state.feedback
+    if isinstance(feedback_data, dict) and feedback_data.get('has_errors'):
+        # Show error feedback
+        st.markdown(f"""
+            <div style="
+                background-color: #fff3e0;
+                padding: 20px;
+                border-radius: 10px;
+                border-left: 4px solid #ff9800;
+                margin: 10px 0;
+            ">
+                <p style="color: #e65100; margin-bottom: 15px;">
+                    {feedback_data['feedback']}
+                </p>
+                <div style="
+                    background-color: #fff;
+                    padding: 15px;
+                    border-radius: 8px;
+                    margin-top: 10px;
+                ">
+                    <p style="color: #666; margin-bottom: 5px;">
+                        ประโยคที่ถูกต้อง:
+                    </p>
+                    <p style="
+                        color: #2e7d32;
+                        font-weight: 500;
+                        font-size: 1.1em;
+                        margin: 0;
+                    ">
+                        {feedback_data['corrected']}
+                    </p>
+                </div>
+            </div>
+        """, unsafe_allow_html=True)
+        
+        # Correction button
+        if st.button(
+            "✍️ แก้ไขประโยคให้ถูกต้อง",
+            key="correct_button",
+            help="คลิกเพื่อแก้ไขประโยคให้ถูกต้องตามคำแนะนำ"
+        ):
+            last_user_entry_idx = next(
+                (i for i, entry in reversed(list(enumerate(st.session_state.story))) 
+                 if entry['role'] == 'You'),
+                None
+            )
+            if last_user_entry_idx is not None:
+                apply_correction(last_user_entry_idx, feedback_data['corrected'])
+                st.rerun()
+    else:
+        # Show success feedback
+        st.markdown(f"""
+            <div style="
+                background-color: #e8f5e9;
+                padding: 20px;
+                border-radius: 10px;
+                border-left: 4px solid #4caf50;
+                margin: 10px 0;
+            ">
+                <p style="color: #2e7d32; margin: 0;">
+                    {feedback_data.get('feedback', '✨ เขียนได้ถูกต้องแล้วค่ะ!')}
+                </p>
+            </div>
+        """, unsafe_allow_html=True)
+
+def show_writing_tools():
+    """Display writing tools section"""
+    with st.expander("✨ เครื่องมือช่วยเขียน | Writing Tools"):
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            if st.button("🎯 ขอคำใบ้", use_container_width=True):
+                with st.spinner("กำลังสร้างคำใบ้..."):
+                    prompt = get_creative_prompt()
+                    st.markdown(f"""
+                        <div style="
+                            background-color: #f3e5f5;
+                            padding: 15px;
+                            border-radius: 8px;
+                            margin-top: 10px;
+                        ">
+                            <div style="color: #6a1b9a;">💭 {prompt['thai']}</div>
+                            <div style="color: #666; font-style: italic;">
+                                💭 {prompt['eng']}
+                            </div>
+                        </div>
+                    """, unsafe_allow_html=True)
+        
+        with col2:
+            if st.button("📚 คำศัพท์แนะนำ", use_container_width=True):
+                with st.spinner("กำลังค้นหาคำศัพท์..."):
+                    vocab_suggestions = get_vocabulary_suggestions()
+                    st.markdown("#### 📚 คำศัพท์น่ารู้")
+                    for word in vocab_suggestions:
+                        st.markdown(f"""
+                            <div style="
+                                background-color: #f5f5f5;
+                                padding: 10px;
+                                border-radius: 5px;
+                                margin: 5px 0;
+                            ">
+                                • {word}
+                            </div>
+                        """, unsafe_allow_html=True)
+
+def show_achievements():
+    """Display achievements and stats"""
+    with st.container():
+        # Points and Streak
+        st.markdown(f"""
+            <div style="
+                background-color: #e3f2fd;
+                padding: 20px;
+                border-radius: 10px;
+                margin-bottom: 20px;
+                text-align: center;
+            ">
+                <h3 style="color: #1e88e5; margin: 0;">
+                    🌟 คะแนนรวม: {st.session_state.points['total']}
+                </h3>
+                <p style="margin: 5px 0;">
+                    🔥 Streak ปัจจุบ���น: {st.session_state.points['streak']} ประโยค
+                </p>
+                <p style="margin: 5px 0;">
+                    ⭐ Streak สูงสุด: {st.session_state.points['max_streak']} ประโยค
+                </p>
+            </div>
+        """, unsafe_allow_html=True)
+
+        # Writing Stats
+        st.markdown("""
+            <div style="margin-bottom: 15px;">
+                <h3>📊 สถิติการเขียน</h3>
+            </div>
+        """, unsafe_allow_html=True)
+        
+        col1, col2 = st.columns(2)
+        with col1:
+            st.metric(
+                "ประโยคที่เขียนทั้งหมด", 
+                st.session_state.stats['total_sentences'],
+                help="จำนวนประโยคทั้งหมดที่คุณได้เขียน"
             )
-        with col2:
+            st.metric(
+                "ถูกต้องตั้งแต่แรก", 
+                st.session_state.stats['correct_first_try'],
+                help="จำนวนประโยคที่ถูกต้องโดยไม่ต้องแก้ไข"
+            )
+        with col2:
             st.metric(
                 "ความแม่นยำ", 
                 f"{st.session_state.stats['accuracy_rate']:.1f}%",
@@ -1884,35 +2283,46 @@ def show_story_input():
             <div class="eng">Your Turn</div>
         </div>
     """, unsafe_allow_html=True)
+    
+    # Initialize clear_input flag if not exists
+    if 'clear_input' not in st.session_state:
+        st.session_state.clear_input = False
+    
+    # Default value for text input
+    default_value = "" if st.session_state.clear_input else st.session_state.get('text_input', "")
+    
+    # If clear_input flag is True, reset it
+    if st.session_state.clear_input:
+        st.session_state.clear_input = False
 
-    # ใช้ callback function เพื่อเคลียร์ input
-    def clear_text():
-        st.session_state.story_input_area = ""
+    # Show remaining sentences if in ending mode
+    if st.session_state.get('ending_mode'):
+        remaining = st.session_state.sentences_to_end
+        st.info(f"🎭 โหมดจบเรื่อง - เหลืออีก {remaining} ประโยค")
     
-    # Input area with default empty value and callback
+    # Input area
     text_input = st.text_area(
         "เขียนต่อจากเรื่องราว | Continue the story:",
-        value="",
+        value=st.session_state.get('text_input', ""),
         height=100,
         key="story_input_area",
         help="พิมพ์ประโยคภาษาอังกฤษเพื่อต่อเรื่อง",
         label_visibility="collapsed"
     )
-
+    
+    # Save current input to session state
+    st.session_state.text_input = text_input
+    
     # Submit button with character count
     col1, col2 = st.columns([3, 1])
     with col1:
-        if st.button("📝 ส่งคำตอบ | Submit", use_container_width=True, on_click=clear_text):
+        if st.button("📝 ส่งคำตอบ | Submit", use_container_width=True):
             if not text_input.strip():
                 st.warning("กรุณาเขียนข้อความก่อนส่ง")
                 return
-                
             try:
                 with st.spinner("กำลังวิเคราะห์ประโยค..."):
-                    if st.session_state.get('ending_mode'):
-                        handle_ending_mode(text_input.strip())
-                    else:
-                        handle_story_submission(text_input.strip())
+                    handle_story_submission(text_input.strip())
             except Exception as e:
                 logging.error(f"Error submitting story: {str(e)}")
                 st.error("เกิดข้อผิดพลาดในการส่งคำตอบ กรุณาลองใหม่อีกครั้ง")
@@ -1931,6 +2341,12 @@ def handle_story_submission(text: str):
         st.error("กรุณาเลือกธีมเรื่องราวก่อนเริ่มเขียน")
         return
 
+    # Check if in ending mode
+    if st.session_state.get('ending_mode'):
+        handle_ending_mode(text)
+        return
+
+    # Regular processing
     try:
         # Get feedback
         feedback_data = provide_feedback(text, st.session_state.level)
@@ -1956,18 +2372,26 @@ def handle_story_submission(text: str):
         
         # Generate AI continuation
         try:
-            ai_response = generate_story_continuation(text, st.session_state.level)
+            logging.info("Attempting to generate AI continuation...")
+            # Use the corrected text if there are errors
+            text_for_continuation = feedback_data['corrected'] if feedback_data.get('has_errors') else text
+            
+            ai_response = generate_story_continuation(text_for_continuation, st.session_state.level)
+            
+            # Log the AI response
+            logging.info(f"AI Response generated: {ai_response}")
+            
             if ai_response and ai_response.strip():
                 st.session_state.story.append({
                     "role": "AI",
                     "content": ai_response,
                     "timestamp": datetime.now().isoformat()
                 })
+                logging.info("AI response added to story successfully")
             else:
-                fallback_response = generate_fallback_response(
-                    st.session_state.current_theme, 
-                    st.session_state.level
-                )
+                logging.error("AI generated empty response")
+                # Create a fallback response
+                fallback_response = generate_fallback_response(st.session_state.current_theme, st.session_state.level)
                 st.session_state.story.append({
                     "role": "AI",
                     "content": fallback_response,
@@ -1975,239 +2399,52 @@ def handle_story_submission(text: str):
                 })
         except Exception as e:
             logging.error(f"Error generating AI continuation: {str(e)}")
-            fallback_response = generate_fallback_response(
-                st.session_state.current_theme, 
-                st.session_state.level
-            )
+            # Create a fallback response
+            fallback_response = generate_fallback_response(st.session_state.current_theme, st.session_state.level)
             st.session_state.story.append({
                 "role": "AI",
                 "content": fallback_response,
                 "timestamp": datetime.now().isoformat()
             })
         
-        # Clear input after successful submission
-        st.session_state.story_input_area = ""
-        
         # Update session stats
         update_session_stats()
         
-        # Rerun to update UI
-        st.rerun()
-        
-    except Exception as e:
-        logging.error(f"Error in story submission: {str(e)}")
-        raise
-
-def generate_fallback_response(theme_id: str, level: str) -> str:
-    """Generate a simple fallback response when AI continuation fails"""
-    try:
-        theme = story_themes.get(theme_id, {})
-        if theme:
-            # ใช้คำศัพท์จาก theme และ level ที่เลือก
-            vocab = theme.get('vocabulary', {}).get(level, [])
-            if vocab:
-                word = random.choice(vocab)
-                
-                # สร้างประโยคตาม level
-                if level == 'Beginner':
-                    return f"The story continues with {word}..."
-                elif level == 'Intermediate':
-                    return f"Something interesting happened with the {word}."
-                else:  # Advanced
-                    return f"The mystery deepens as we discover more about the {word}."
-        
-        # Default fallback if no theme-specific response can be generated
-        return "The story continues..."
-        
-    except Exception as e:
-        logging.error(f"Error generating fallback response: {str(e)}")
-        return "What happens next?"
-
-def show_story_progress():
-    """Display story progress metrics"""
-    if st.session_state.story:
-        total_sentences = len(st.session_state.story)
-        st.markdown(f"""
-            <div style="
-                background-color: #f0f7ff;
-                padding: 15px;
-                border-radius: 10px;
-                margin: 10px 0;
-            ">
-                <div style="margin-bottom: 10px;">
-                    📊 ความยาวเรื่อง: {total_sentences} ประโยค
-                </div>
-                <div class="progress-bar">
-                    <div class="progress-bar-fill" style="width: {min(total_sentences * 10, 100)}%;"></div>
-                </div>
-                <div style="font-size: 0.9em; color: #666; margin-top: 5px;">
-                    เรื่องควรยาว 10-20 ประโยค เพื่อความสมบูรณ์
-                </div>
-            </div>
-        """, unsafe_allow_html=True)
-
-def show_story_ending_options():
-    """Display story ending options and guidance"""
-    if len(st.session_state.story) >= 10:  # แสดงตัวเลือกเมื่อเรื่องยาวพอ
-        st.markdown("### 🎭 ต้องการจบเรื่องหรือไม่?")
-        
-        # แสดงตัวเลือกวิธีจบเรื่อง
-        ending_type = st.radio(
-            "เลือกวิธีจบเรื่อง:",
-            options=[
-                "Happy Ending - จบแบบมีความสุข",
-                "Mysterious Ending - จบแบบทิ้งท้ายให้คิดต่อ",
-                "Lesson Learned - จบแบบได้ข้อคิด",
-                "Surprise Ending - จบแบบพลิกความคาดหมาย"
-            ],
-            index=0,
-            help="เลือกรูปแบบการจบเรื่องที่คุณต้องการ"
-        )
-        
-        if st.button("🎬 เริ่มจบเรื่อง", use_container_width=True):
-            st.session_state.ending_mode = True
-            st.session_state.ending_type = ending_type
-            st.session_state.sentences_to_end = 5
-            st.rerun()
+        # Set flag to clear input on next rerun
+        st.session_state.clear_input = True
+        
+        # Rerun to update UI
+        st.rerun()
+        
+    except Exception as e:
+        logging.error(f"Error in story submission: {str(e)}")
+        raise
 
-def handle_ending_mode(text: str):
-    """Handle story submission during ending mode"""
-    remaining = st.session_state.sentences_to_end
-    
-    # แสดงการแจ้งเตือนจำนวนประโยคที่เหลือ
-    st.info(f"🎯 เหลืออีก {remaining} ประโยคในการจบเรื่อง")
-    
-    # สร้าง prompt พิเศษสำหรับ AI ในโหมดจบเรื่อง
-    ending_prompts = {
-        "Happy Ending": "Work towards a positive and uplifting conclusion",
-        "Mysterious Ending": "Create an intriguing open-ended conclusion",
-        "Lesson Learned": "Incorporate a meaningful life lesson",
-        "Surprise Ending": "Build up to an unexpected twist"
-    }
-    
-    # ปรับ AI response ให้พยายามจบเรื่อง
-    modified_continuation = generate_ending_continuation(
-        text,
-        ending_type=st.session_state.ending_type,
-        remaining_sentences=remaining
-    )
-    
-    # อัพเดทจำนวนประโยคที่เหลือ
-    st.session_state.sentences_to_end -= 1
-    
-    # ตรวจสอบว่าถึงจบเรื่องหรือยัง
-    if st.session_state.sentences_to_end <= 0:
-        complete_story()
 
-def generate_ending_continuation(text: str, ending_type: str, remaining_sentences: int) -> str:
-    """Generate AI continuation focusing on story conclusion"""
+def generate_fallback_response(theme_id: str, level: str) -> str:
+    """Generate a simple fallback response when AI continuation fails"""
     try:
-        ending_prompts = {
-            "Happy Ending": """
-                Role: Story concluder aiming for a happy ending
-                Goal: Create a satisfying, positive conclusion
-                Rules:
-                - Build towards joy, success, or resolution
-                - Use uplifting and positive language
-                - Connect to previous story elements
-                """,
-            "Mysterious Ending": """
-                Role: Mystery writer creating intrigue
-                Goal: Leave readers thinking and wondering
-                Rules:
-                - Add subtle hints and clues
-                - Create atmospheric descriptions
-                - Leave some questions unanswered
-                """,
-            "Lesson Learned": """
-                Role: Moral story concluder
-                Goal: Incorporate meaningful life lessons
-                Rules:
-                - Connect actions to consequences
-                - Show character growth
-                - Express the moral naturally
-                """,
-            "Surprise Ending": """
-                Role: Plot twist creator
-                Goal: Deliver unexpected but satisfying conclusion
-                Rules:
-                - Plant subtle hints earlier
-                - Subvert expectations logically
-                - Maintain story coherence
-                """
-        }
-
-        response = client.chat.completions.create(
-            model="gpt-4",
-            messages=[
-                {
-                    "role": "system",
-                    "content": f"""
-                    {ending_prompts[ending_type]}
-                    
-                    Additional Rules:
-                    - You have {remaining_sentences} sentences to conclude
-                    - Each response should be max 2 sentences
-                    - Build towards the finale naturally
-                    - Connect to previous story elements
-                    """
-                },
-                {
-                    "role": "user",
-                    "content": f"Continue and work towards ending this story: {text}"
-                }
-            ],
-            max_tokens=100,
-            temperature=0.7
-        )
+        theme = story_themes.get(theme_id, {})
+        if theme:
+            # ใช้คำศัพท์จาก theme และ level ที่เลือก
+            vocab = theme.get('vocabulary', {}).get(level, [])
+            if vocab:
+                word = random.choice(vocab)
+                
+                # สร้างประโยคตาม level
+                if level == 'Beginner':
+                    return f"The story continues with {word}..."
+                elif level == 'Intermediate':
+                    return f"Something interesting happened with the {word}."
+                else:  # Advanced
+                    return f"The mystery deepens as we discover more about the {word}."
         
-        return response.choices[0].message.content.strip()
+        # Default fallback if no theme-specific response can be generated
+        return "The story continues..."
         
     except Exception as e:
-        logging.error(f"Error generating ending: {str(e)}")
-        return "The story moved towards its conclusion..."
-
-def complete_story():
-    """Handle story completion and celebration"""
-    st.balloons()  # แสดงเอฟเฟคฉลอง
-    
-    # สร้าง Story Summary
-    story_summary = generate_story_summary(st.session_state.story)
-    
-    # แสดงหน้าจบเรื่อง
-    st.markdown(f"""
-        <div style="
-            background-color: #e8f5e9;
-            padding: 20px;
-            border-radius: 10px;
-            text-align: center;
-            margin: 20px 0;
-        ">
-            <h2>🎉 ยินดีด้วย! คุณเขียนเรื่องราวจบสมบูรณ์แล้ว</h2>
-            
-            <div style="margin: 20px 0;">
-                <h3>📝 สรุปเรื่องราว</h3>
-                <p>{story_summary}</p>
-            </div>
-            
-            <div style="margin: 20px 0;">
-                <h3>🏆 ความสำเร็จ</h3>
-                <p>จำนวนประโยค: {len(st.session_state.story)}</p>
-                <p>คำศัพท์ที่ใช้: {len(st.session_state.stats['vocabulary_used'])}</p>
-                <p>ความแม่นยำ: {st.session_state.stats['accuracy_rate']:.1f}%</p>
-            </div>
-        </div>
-    """, unsafe_allow_html=True)
-    
-    # แสดงตัวเลือกหลังจบเรื่อง
-    col1, col2 = st.columns(2)
-    with col1:
-        if st.button("💾 บันทึกเรื่องราว", use_container_width=True):
-            save_completed_story()
-    with col2:
-        if st.button("🔄 เริ่มเรื่องใหม่", use_container_width=True):
-            reset_story()
-            st.rerun()
+        logging.error(f"Error generating fallback response: {str(e)}")
+        return "What happens next?"
 
 def show_main_interface():
     """Display main story interface"""
@@ -2223,9 +2460,14 @@ def show_main_interface():
         """, unsafe_allow_html=True)
         
         show_story()
-        
-        if st.session_state.story:
+
+        # Only show input if story is not completed
+        if st.session_state.story and not st.session_state.get('story_completed'):
             show_story_input()
+        
+        # Show completion options if story is done
+        if st.session_state.get('story_completed'):
+            show_completion_options()
     
     with col2:
         # Feedback section
@@ -2241,6 +2483,312 @@ def show_main_interface():
         # Save options
         show_save_options()
 
+def save_completed_story():
+    """Save completed story with all available formats"""
+    try:
+        st.markdown("### 💾 บันทึกเรื่องราว", unsafe_allow_html=True)
+        
+        save_col1, save_col2 = st.columns(2)
+        
+        with save_col1:
+            # Create and offer PDF download
+            try:
+                pdf_data = create_story_pdf()
+                st.download_button(
+                    label="📥 ดาวน์โหลด PDF",
+                    data=pdf_data,
+                    file_name=f"joystory_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf",
+                    mime="application/pdf",
+                    key="download_story_pdf",
+                    use_container_width=True
+                )
+            except Exception as e:
+                logging.error(f"PDF creation failed: {str(e)}")
+                st.error("ไม่สามารถสร้างไฟล์ PDF ได้")
+        
+        with save_col2:
+            # Create and offer JSON download
+            try:
+                story_data = {
+                    'timestamp': datetime.now().isoformat(),
+                    'level': st.session_state.level,
+                    'theme': st.session_state.current_theme,
+                    'story': st.session_state.story,
+                    'stats': {
+                        key: list(value) if isinstance(value, set) else value
+                        for key, value in st.session_state.stats.items()
+                    },
+                    'points': st.session_state.points,
+                    'achievements': st.session_state.achievements,
+                    'ending_type': st.session_state.ending_type
+                }
+                
+                json_str = json.dumps(story_data, ensure_ascii=False, indent=2)
+                
+                st.download_button(
+                    label="📥 ดาวน์โหลดเรื่องราว",
+                    data=json_str,
+                    file_name=f"joystory_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
+                    mime="application/json",
+                    key="download_story_json",
+                    use_container_width=True
+                )
+            except Exception as e:
+                logging.error(f"JSON creation failed: {str(e)}")
+                st.error("ไม่สามารถสร้างไฟล์บันทึกได้")
+
+    except Exception as e:
+        logging.error(f"Error in save_completed_story: {str(e)}")
+        st.error("เกิดข้อผิดพลาดในการบันทึกเรื่องราว")
+        
+def show_completion_options():
+    """Display options and summary when story is completed"""
+    try:
+        st.balloons()
+        
+        # Show completion message
+        st.markdown("""
+            <div style="
+                background-color: #e8f5e9;
+                padding: 20px;
+                border-radius: 10px;
+                text-align: center;
+                margin: 20px 0;
+                border: 2px solid #4caf50;
+            ">
+                <h2 style="color: #2e7d32; margin-bottom: 15px;">
+                    🎉 ยินดีด้วย! คุณเขียนเรื่องราวจบสมบูรณ์แล้ว
+                </h2>
+                <p style="color: #1b5e20;">
+                    เรื่องราวของคุณจบลงด้วย: {st.session_state.ending_type}
+                </p>
+            </div>
+        """, unsafe_allow_html=True)
+        
+        # Show statistics
+        metrics_col1, metrics_col2 = st.columns(2)
+        with metrics_col1:
+            st.metric(
+                "จำนวนประโยคทั้งหมด",
+                len(st.session_state.story),
+                help="จำนวนประโยคในเรื่องทั้งหมด"
+            )
+            st.metric(
+                "คำศัพท์ที่ใช้",
+                len(st.session_state.stats['vocabulary_used']),
+                help="จำนวนคำศัพท์ที่ไม่ซ้ำกัน"
+            )
+        
+        with metrics_col2:
+            st.metric(
+                "คะแนนรวม",
+                st.session_state.points['total'],
+                help="คะแนนรวมที่ได้รับ"
+            )
+            st.metric(
+                "ความแม่นยำ",
+                f"{st.session_state.stats['accuracy_rate']:.1f}%",
+                help="อัตราการเขียนถูกต้อง"
+            )
+        
+        # Save options
+        st.markdown("### 🎯 ต้องการทำอะไรต่อ?")
+        save_completed_story()
+        
+        # New story option
+        if st.button("🔄 เริ่มเรื่องใหม่", 
+                    key="new_story_button", 
+                    use_container_width=True):
+            if st.checkbox("✅ ยืนยันการเริ่มใหม่", key="confirm_reset"):
+                st.warning("การเริ่มใหม่จะลบเรื่องราวปัจจุบันทั้งหมด")
+                if st.button("🔄 ยืนยันการเริ่มใหม่", 
+                            key="final_confirm_reset", 
+                            use_container_width=True):
+                    reset_story()
+                    st.rerun()
+
+    except Exception as e:
+        logging.error(f"Error showing completion options: {str(e)}")
+        st.error("เกิดข้อผิดพลาดในการแสดงตัวเลือกหลังจบเรื่อง")
+
+def show_save_dialog():
+    """Display save options dialog"""
+    try:
+        st.markdown("### 💾 บันทึกเรื่องราว")
+        
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            # PDF Export with unique key
+            if st.button("📑 บันทึกเป็น PDF", 
+                        key="save_pdf_button", 
+                        use_container_width=True):
+                try:
+                    pdf = create_story_pdf()
+                    st.download_button(
+                        "📥 ดาวน์โหลด PDF",
+                        data=pdf,
+                        file_name=f"story_{datetime.now().strftime('%Y%m%d')}.pdf",
+                        mime="application/pdf",
+                        key="download_pdf_button"
+                    )
+                except Exception as e:
+                    logging.error(f"Error creating PDF: {str(e)}")
+                    st.error("เกิดข้อผิดพลาดในการสร้างไฟล์ PDF")
+        
+        with col2:
+            # JSON Save with unique key
+            if st.button("💾 บันทึกความก้าวหน้า", 
+                        key="save_progress_button", 
+                        use_container_width=True):
+                try:
+                    story_data = {
+                        'timestamp': datetime.now().isoformat(),
+                        'level': st.session_state.level,
+                        'story': st.session_state.story,
+                        'achievements': st.session_state.achievements,
+                        'stats': {
+                            key: list(value) if isinstance(value, set) else value
+                            for key, value in st.session_state.stats.items()
+                        },
+                        'points': st.session_state.points
+                    }
+                    
+                    st.download_button(
+                        "📥 ดาวน์โหลดไฟล์บันทึก",
+                        data=json.dumps(story_data, ensure_ascii=False, indent=2),
+                        file_name=f"story_progress_{datetime.now().strftime('%Y%m%d')}.json",
+                        mime="application/json",
+                        key="download_json_button"
+                    )
+                except Exception as e:
+                    logging.error(f"Error saving progress: {str(e)}")
+                    st.error("เกิดข้อผิดพลาดในการบันทึกความก้าวหน้า")
+    except Exception as e:
+        logging.error(f"Error showing save dialog: {str(e)}")
+        st.error("เกิดข้อผิดพลาดในการแสดงตัวเลือกการบันทึก")
+
+def reset_story_with_confirmation():
+    """Reset story with confirmation dialog"""
+    try:
+        if st.session_state.story:
+            if st.checkbox("✅ ยืนยันการเริ่มใหม่", 
+                          key="reset_confirmation_checkbox"):
+                st.warning("การเริ่มใหม่จะลบเรื่องราวปัจจุบันทั้งหมด")
+                if st.button("🔄 ยืนยันการเริ่มใหม่", 
+                            key="confirm_reset_button", 
+                            use_container_width=True):
+                    reset_story()
+                    st.success("เริ่มต้นใหม่เรียบร้อยแล้ว!")
+                    st.rerun()
+    except Exception as e:
+        logging.error(f"Error in reset confirmation: {str(e)}")
+        st.error("เกิดข้อผิดพลาดในการรีเซ็ตเรื่อง")
+
+def create_story_pdf():
+    """Create PDF from story content"""
+    try:
+        buffer = io.BytesIO()
+        doc = SimpleDocTemplate(
+            buffer, 
+            pagesize=A4,
+            rightMargin=72,
+            leftMargin=72,
+            topMargin=72,
+            bottomMargin=72
+        )
+        
+        # Prepare styles
+        styles = getSampleStyleSheet()
+        title_style = ParagraphStyle(
+            'CustomTitle',
+            parent=styles['Title'],
+            fontSize=24,
+            spaceAfter=30,
+            alignment=1  # Center alignment
+        )
+        normal_style = ParagraphStyle(
+            'CustomBody',
+            parent=styles['Normal'],
+            fontSize=12,
+            spaceBefore=6,
+            spaceAfter=6
+        )
+        
+        # Create story elements list
+        elements = []
+        
+        # Add title
+        title_text = "My Story - JoyStory"
+        elements.append(Paragraph(title_text, title_style))
+        elements.append(Spacer(1, 12))
+        
+        # Add metadata
+        metadata_style = ParagraphStyle(
+            'Metadata',
+            parent=styles['Normal'],
+            fontSize=10,
+            textColor=colors.gray,
+            alignment=1
+        )
+        metadata_text = f"Level: {st.session_state.level}<br/>Date: {datetime.now().strftime('%Y-%m-%d')}"
+        elements.append(Paragraph(metadata_text, metadata_style))
+        elements.append(Spacer(1, 20))
+        
+        # Add story content
+        for entry in st.session_state.story:
+            # Skip system messages or empty entries
+            if not entry.get('content'):
+                continue
+                
+            # Format based on role
+            if entry['role'] == 'AI':
+                text = f"🤖 AI: {entry['content']}"
+                color = colors.blue
+            else:
+                text = f"👤 You: {entry['content']}"
+                color = colors.black
+            
+            # Create paragraph style with color
+            style = ParagraphStyle(
+                'ColoredText',
+                parent=normal_style,
+                textColor=color
+            )
+            
+            # Add paragraph
+            elements.append(Paragraph(text, style))
+            elements.append(Spacer(1, 6))
+        
+        # Add statistics
+        elements.append(Spacer(1, 20))
+        stats_style = ParagraphStyle(
+            'Stats',
+            parent=styles['Normal'],
+            fontSize=10,
+            textColor=colors.gray
+        )
+        
+        stats_text = f"""
+        Story Statistics:<br/>
+        Total Sentences: {len(st.session_state.story)}<br/>
+        Unique Words: {len(st.session_state.stats['vocabulary_used'])}<br/>
+        Accuracy Rate: {st.session_state.stats['accuracy_rate']:.1f}%<br/>
+        Total Points: {st.session_state.points['total']}
+        """
+        elements.append(Paragraph(stats_text, stats_style))
+        
+        # Build PDF
+        doc.build(elements)
+        pdf = buffer.getvalue()
+        buffer.close()
+        
+        return pdf
+        
+    except Exception as e:
+        logging.error(f"Error creating PDF: {str(e)}")
+        raise
+
 # === 6. MAIN APPLICATION LOGIC ===
 def main():
     try:
@@ -2252,9 +2800,11 @@ def main():
         if 'ending_mode' not in st.session_state:
             st.session_state.ending_mode = False
         if 'sentences_to_end' not in st.session_state:
-            st.session_state.sentences_to_end = 5
+            st.session_state.sentences_to_end = 0  # Set to 0 by default
         if 'ending_type' not in st.session_state:
             st.session_state.ending_type = None
+        if 'story_completed' not in st.session_state:
+            st.session_state.story_completed = False
         
         # Add watermark
         st.markdown("""
@@ -2284,36 +2834,104 @@ def main():
             else:
                 # Show story progress if story exists
                 if st.session_state.story:
-                    show_story_progress()
+                    # Display story progress
+                    total_sentences = len(st.session_state.story)
+                    st.markdown(f"""
+                        <div style="
+                            background-color: #f0f7ff;
+                            padding: 15px;
+                            border-radius: 10px;
+                            margin: 10px 0;
+                        ">
+                            <div style="margin-bottom: 10px;">
+                                📊 ความยาวเรื่อง: {total_sentences} ประโยค
+                            </div>
+                            <div class="progress-bar">
+                                <div class="progress-bar-fill" style="width: {min(total_sentences * 5, 100)}%;"></div>
+                            </div>
+                            <div style="font-size: 0.9em; color: #666; margin-top: 5px;">
+                                เรื่องควรยาว 10-20 ประโยค เพื่อความสมบูรณ์
+                            </div>
+                        </div>
+                    """, unsafe_allow_html=True)
                     
-                    # Show ending options if story is long enough and not in ending mode
-                    if (len(st.session_state.story) >= 10 and 
-                        not st.session_state.ending_mode):
-                        show_story_ending_options()
+                    # Show ending options if story is long enough and ending mode is not active
+                    if (total_sentences >= 10 and not st.session_state.get('ending_mode') and not st.session_state.get('story_completed')):
+                        st.markdown("### 🎭 ต้องการจบเรื่องหรือไม่?")
+                        ending_type = st.radio(
+                            "เลือกวิธีจบเรื่อง:",
+                            options=[
+                                "Happy Ending",
+                                "Mysterious Ending",
+                                "Lesson Learned",
+                                "Surprise Ending"
+                            ],
+                            index=0,
+                            key="ending_type_selector"
+                        )
+                        
+                        if st.button("🎬 เริ่มจบเรื่อง", use_container_width=True):
+                            st.session_state.ending_mode = True
+                            st.session_state.ending_type = ending_type
+                            st.session_state.sentences_to_end = 5  # Set the countdown
+                            st.success(f"โหมดจบเรื่องเริ่มต้นแล้ว! รูปแบบการจบ: {ending_type}")
+                            st.rerun()
                     
-                    # Show ending mode warning if active
-                    if st.session_state.ending_mode:
-                        st.warning(f"""
-                            🎭 กำลังอยู่ในโหมดจบเรื่อง - เหลือ {st.session_state.sentences_to_end} ประโยค
-                            \nรูปแบบการจบ: {st.session_state.ending_type}
-                        """)
-                
-                # Show main interface
+                    # Show ending mode notification
+                    if st.session_state.get('ending_mode'):
+                        remaining = st.session_state.sentences_to_end
+                        if remaining > 0:
+                            st.warning(f"""
+                                🎭 กำลังอยู่ในโหมดจบเรื่อง
+                                - เหลืออีก {remaining} ประโยค
+                                - รูปแบบการจบ: {st.session_state.ending_type}
+                                
+                                พยายามเขียนให้เรื่องจบอย่างสมบูรณ์!
+                            """)
+                        elif remaining <= 0 and not st.session_state.get('story_completed'):
+                            # Story should be marked as completed
+                            st.success("🎉 เรื่องราวของคุณจบลงแล้ว!")
+                            st.balloons()
+                            st.session_state.ending_mode = False
+                            st.session_state.story_completed = True
+
+                # Display main interface
                 show_main_interface()
                 
-                # Check if story is complete
-                if (st.session_state.ending_mode and 
-                    st.session_state.sentences_to_end <= 0):
-                    complete_story()
-
+                # Handle story completion
+                if st.session_state.get('story_completed'):
+                    st.markdown("""
+                        <div style="
+                            background-color: #e8f5e9;
+                            padding: 20px;
+                            border-radius: 10px;
+                            text-align: center;
+                            margin: 20px 0;
+                        ">
+                            <h2>🎉 ยินดีด้วย! คุณเขียนเรื่องราวจบสมบูรณ์แล้ว</h2>
+                            <p>คุณสามารถบันทึกเรื่องราวหรือเริ่มเรื่องใหม่ได้</p>
+                        </div>
+                    """, unsafe_allow_html=True)
+                    # Show save and reset options
+                    col1, col2 = st.columns(2)
+                    with col1:
+                        if st.button("💾 บันทึกเรื่องราว", use_container_width=True):
+                            save_completed_story()
+                    with col2:
+                        if st.button("🔄 เริ่มเรื่องใหม่", use_container_width=True):
+                            reset_story()
+                            st.rerun()
+                
         # Handle reset if needed
-        if st.session_state.should_reset:
+        if st.session_state.get('should_reset'):
             reset_story()
-            # Reset ending mode states
+            # Reset ending states
             st.session_state.ending_mode = False
-            st.session_state.sentences_to_end = 5
+            st.session_state.sentences_to_end = 0
             st.session_state.ending_type = None
-            
+            st.session_state.story_completed = False
+            st.rerun()  # Using st.rerun() instead of st.experimental_rerun()
+        
         # Auto-save progress periodically
         if st.session_state.story:
             auto_save_progress()
@@ -2321,6 +2939,7 @@ def main():
     except Exception as e:
         handle_application_error(e)
 
+
 def check_session_status():
     """Check and maintain session status"""
     try:
@@ -2328,23 +2947,16 @@ def check_session_status():
         if 'session_start' not in st.session_state:
             st.session_state.session_start = datetime.now()
         
-        # Calculate session duration
-        current_time = datetime.now()
-        session_duration = (current_time - st.session_state.session_start).total_seconds()
-        
         # Check for session timeout (2 hours)
+        session_duration = (datetime.now() - st.session_state.session_start).total_seconds()
         if session_duration > 7200:  # 2 hours
-            st.warning("⚠️ เซสชันหมดอายุ กรุณาบันทึกความก้าวหน้าและรีเฟรชหน้าเว็บ")
+            st.warning("เซสชันหมดอายุ กรุณาบันทึกความก้าวหน้าและรีเฟรชหน้าเว็บ")
         
         # Check for inactivity (30 minutes)
-        if 'last_interaction' not in st.session_state:
-            st.session_state.last_interaction = datetime.now().isoformat()
-            
         last_interaction = datetime.fromisoformat(st.session_state.last_interaction)
-        inactivity_duration = (current_time - last_interaction).total_seconds()
-        
+        inactivity_duration = (datetime.now() - last_interaction).total_seconds()
         if inactivity_duration > 1800:  # 30 minutes
-            st.info("💤 ไม่มีกิจกรรมเป็นเวลานาน กรุณาบันทึกความก้าวหน้าเพื่อความปลอดภัย")
+            st.info("ไม่มีกิจกรรมเป็นเวลานาน กรุณาบันทึกความก้าวหน้าเพื่อความปลอดภัย")
         
         # Update stats if story exists
         if st.session_state.story: