app.py

# === 1. IMPORTS & SETUP === 
import streamlit as st
import json
import datetime
import logging
from openai import OpenAI
from typing import Dict, List, Set, Tuple
import io
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from datetime import datetime 
import random
from sounds import AudioManager, get_sound_commands

# === 2. CONFIGURATIONS ===
# Theme Configuration
story_themes = {
    'fantasy': {
        'id': 'fantasy',
        'icon': '🏰',
        'name_en': 'Fantasy & Magic',
        'name_th': 'แฟนตาซีและเวทมนตร์',
        'description_th': 'ผจญภัยในโลกแห่งเวทมนตร์และจินตนาการ',
        'description_en': 'Adventure in a world of magic and imagination',
        'level_range': ['Beginner', 'Intermediate', 'Advanced'],
        'vocabulary': {
            'Beginner': ['dragon', 'magic', 'wand', 'spell', 'wizard', 'fairy', 'castle', 'king', 'queen'],
            'Intermediate': ['potion', 'enchanted', 'castle', 'creature', 'power', 'scroll', 'portal', 'magical'],
            'Advanced': ['sorcery', 'mystical', 'enchantment', 'prophecy', 'ancient', 'legendary', 'mythical']
        },
        'story_starters': {
            'Beginner': [
                {'th': 'วันหนึ่งฉันเจอไม้วิเศษในสวน...', 'en': 'One day, I found a magic wand in the garden...'},
                {'th': 'มังกรน้อยกำลังมองหาเพื่อน...', 'en': 'The little dragon was looking for a friend...'},
                {'th': 'เจ้าหญิงน้อยมีความลับวิเศษ...', 'en': 'The little princess had a magical secret...'}
            ],
            'Intermediate': [
                {'th': 'ในปราสาทเก่าแก่มีประตูลึกลับ...', 'en': 'In the ancient castle, there was a mysterious door...'},
                {'th': 'เมื่อน้ำยาวิเศษเริ่มส่องแสง...', 'en': 'When the magic potion started to glow...'},
                {'th': 'หนังสือเวทมนตร์เล่มนั้นเปิดออกเอง...', 'en': 'The spellbook opened by itself...'}
            ],
            'Advanced': [
                {'th': 'คำทำนายโบราณกล่าวถึงผู้วิเศษที่จะมา...', 'en': 'The ancient prophecy spoke of a wizard who would come...'},
                {'th': 'ในโลกที่เวทมนตร์กำลังจะสูญหาย...', 'en': 'In a world where magic was fading away...'},
                {'th': 'ณ จุดบรรจบของดวงดาวทั้งห้า...', 'en': 'At the convergence of the five stars...'}
            ]
        },
        'background_color': '#E8F3FF',
        'accent_color': '#1E88E5'
    },
    'nature': {
        'id': 'nature',
        'icon': '🌳',
        'name_en': 'Nature & Animals',
        'name_th': 'ธรรมชาติและสัตว์โลก',
        'description_th': 'เรื่องราวของสัตว์น้อยใหญ่และธรรมชาติอันงดงาม',
        'description_en': 'Stories of animals and beautiful nature',
        'level_range': ['Beginner', 'Intermediate', 'Advanced'],
        'vocabulary': {
            'Beginner': ['tree', 'bird', 'flower', 'cat', 'dog', 'garden', 'rabbit', 'butterfly', 'sun'],
            'Intermediate': ['forest', 'river', 'mountain', 'wildlife', 'season', 'weather', 'rainbow', 'stream'],
            'Advanced': ['ecosystem', 'habitat', 'wilderness', 'environment', 'conservation', 'migration', 'climate']
        },
        'story_starters': {
            'Beginner': [
                {'th': 'แมวน้อยเจอนกในสวน...', 'en': 'The little cat found a bird in the garden...'},
                {'th': 'ดอกไม้สวยกำลังเบ่งบาน...', 'en': 'The beautiful flower was blooming...'},
                {'th': 'กระต่ายน้อยหลงทางในสวน...', 'en': 'The little rabbit got lost in the garden...'}
            ],
            'Intermediate': [
                {'th': 'ในป่าใหญ่มีเสียงลึกลับ...', 'en': 'In the big forest, there was a mysterious sound...'},
                {'th': 'แม่น้ำสายนี้มีความลับ...', 'en': 'This river had a secret...'},
                {'th': 'สายรุ้งพาดผ่านภูเขา...', 'en': 'A rainbow stretched across the mountain...'}
            ],
            'Advanced': [
                {'th': 'ฝูงนกกำลังอพยพย้ายถิ่น...', 'en': 'The birds were migrating...'},
                {'th': 'ป่าฝนกำลังเปลี่ยนแปลง...', 'en': 'The rainforest was changing...'},
                {'th': 'ความลับของระบบนิเวศ...', 'en': 'The secret of the ecosystem...'}
            ]
        },
        'background_color': '#F1F8E9',
        'accent_color': '#4CAF50'
    },
    'space': {
        'id': 'space',
        'icon': '🚀',
        'name_en': 'Space Adventure',
        'name_th': 'ผจญภัยในอวกาศ',
        'description_th': 'เรื่องราวการสำรวจอวกาศและดวงดาวอันน่าตื่นเต้น',
        'description_en': 'Exciting stories of space exploration and celestial discoveries',
        'level_range': ['Beginner', 'Intermediate', 'Advanced'],
        'vocabulary': {
            'Beginner': ['star', 'moon', 'planet', 'sun', 'rocket', 'alien', 'space', 'light'],
            'Intermediate': ['astronaut', 'spacecraft', 'galaxy', 'meteor', 'satellite', 'orbit', 'comet'],
            'Advanced': ['constellation', 'nebula', 'astronomy', 'telescope', 'exploration', 'discovery']
        },
        'story_starters': {
            'Beginner': [
                {'th': 'จรวดลำน้อยพร้อมบินแล้ว...', 'en': 'The little rocket was ready to fly...'},
                {'th': 'ดาวดวงน้อยเปล่งแสงวิบวับ...', 'en': 'The little star twinkled brightly...'},
                {'th': 'มนุษย์ต่างดาวที่เป็นมิตร...', 'en': 'The friendly alien...'}
            ],
            'Intermediate': [
                {'th': 'นักบินอวกาศพบสิ่งประหลาด...', 'en': 'The astronaut found something strange...'},
                {'th': 'ดาวเคราะห์ดวงใหม่ถูกค้นพบ...', 'en': 'A new planet was discovered...'},
                {'th': 'สถานีอวกาศส่งสัญญาณลึกลับ...', 'en': 'The space station sent a mysterious signal...'}
            ],
            'Advanced': [
                {'th': 'การสำรวจดาวหางนำไปสู่การค้นพบ...', 'en': 'The comet exploration led to a discovery...'},
                {'th': 'กาแล็กซี่ที่ไม่มีใครเคยเห็น...', 'en': 'An unknown galaxy appeared...'},
                {'th': 'ความลับของหลุมดำ...', 'en': 'The secret of the black hole...'}
            ]
        },
        'background_color': '#E1F5FE',
        'accent_color': '#0288D1'
    },
    'adventure': {
        'id': 'adventure',
        'icon': '🗺️',
        'name_en': 'Adventure & Quest',
        'name_th': 'การผจญภัยและการค้นหา',
        'description_th': 'ออกผจญภัยค้นหาสมบัติและความลับต่างๆ',
        'description_en': 'Embark on quests to find treasures and secrets',
        'level_range': ['Beginner', 'Intermediate', 'Advanced'],
        'vocabulary': {
            'Beginner': ['map', 'treasure', 'cave', 'island', 'path', 'boat', 'key', 'chest'],
            'Intermediate': ['compass', 'adventure', 'journey', 'mystery', 'explore', 'discover', 'quest'],
            'Advanced': ['expedition', 'archaeology', 'artifact', 'ancient', 'mysterious', 'discovery']
        },
        'story_starters': {
            'Beginner': [
                {'th': 'แผนที่เก่าแก่ชิ้นหนึ่ง...', 'en': 'An old map showed...'},
                {'th': 'บนเกาะเล็กๆ มีสมบัติ...', 'en': 'On a small island, there was a treasure...'},
                {'th': 'ถ้ำลึกลับถูกค้นพบ...', 'en': 'A mysterious cave was found...'}
            ],
            'Intermediate': [
                {'th': 'เข็มทิศวิเศษชี้ไปที่...', 'en': 'The magical compass pointed to...'},
                {'th': 'การเดินทางเริ่มต้นที่...', 'en': 'The journey began at...'},
                {'th': 'ความลับของวัตถุโบราณ...', 'en': 'The secret of the ancient artifact...'}
            ],
            'Advanced': [
                {'th': 'การสำรวจซากปรักหักพังนำไปสู่...', 'en': 'The ruins exploration led to...'},
                {'th': 'นักโบราณคดีค้นพบ...', 'en': 'The archaeologist discovered...'},
                {'th': 'ความลับของอารยธรรมโบราณ...', 'en': 'The secret of the ancient civilization...'}
            ]
        },
        'background_color': '#FFF3E0',
        'accent_color': '#FF9800'
    },
    'school': {
        'id': 'school',
        'icon': '🏫',
        'name_en': 'School & Friends',
        'name_th': 'โรงเรียนและเพื่อน',
        'description_th': 'เรื่องราวสนุกๆ ในโรงเรียนกับเพื่อนๆ',
        'description_en': 'Fun stories about school life and friendship',
        'level_range': ['Beginner', 'Intermediate', 'Advanced'],
        'vocabulary': {
            'Beginner': ['friend', 'teacher', 'book', 'classroom', 'pencil', 'desk', 'lunch', 'play'],
            'Intermediate': ['homework', 'project', 'library', 'playground', 'student', 'lesson', 'study'],
            'Advanced': ['presentation', 'experiment', 'knowledge', 'research', 'collaboration', 'achievement']
        },
        'story_starters': {
            'Beginner': [
                {'th': 'วันแรกในห้องเรียนใหม่...', 'en': 'First day in the new classroom...'},
                {'th': 'เพื่อนใหม่ในโรงเรียน...', 'en': 'A new friend at school...'},
                {'th': 'ที่โต๊ะอาหารกลางวัน...', 'en': 'At the lunch table...'}
            ],
            'Intermediate': [
                {'th': 'โครงงานพิเศษของห้องเรา...', 'en': 'Our class special project...'},
                {'th': 'ในห้องสมุดมีความลับ...', 'en': 'The library had a secret...'},
                {'th': 'การทดลองวิทยาศาสตร์ครั้งนี้...', 'en': 'This science experiment...'}
            ],
            'Advanced': [
                {'th': 'การนำเสนอครั้งสำคัญ...', 'en': 'The important presentation...'},
                {'th': 'การค้นคว้าพิเศษนำไปสู่...', 'en': 'The special research led to...'},
                {'th': 'โครงการความร่วมมือระหว่างห้อง...', 'en': 'The inter-class collaboration project...'}
            ]
        },
        'background_color': '#F3E5F5',
        'accent_color': '#9C27B0'
    },
    'superhero': {
        'id': 'superhero',
        'icon': '🦸',
        'name_en': 'Superheroes',
        'name_th': 'ซูเปอร์ฮีโร่',
        'description_th': 'เรื่องราวของฮีโร่ตัวน้อยผู้ช่วยเหลือผู้อื่น',
        'description_en': 'Stories of young heroes helping others',
        'level_range': ['Beginner', 'Intermediate', 'Advanced'],
        'vocabulary': {
            'Beginner': ['hero', 'help', 'save', 'power', 'mask', 'cape', 'fly', 'strong'],
            'Intermediate': ['rescue', 'protect', 'brave', 'courage', 'mission', 'team', 'secret', 'mighty'],
            'Advanced': ['superhero', 'extraordinary', 'responsibility', 'leadership', 'determination', 'justice']
        },
        'story_starters': {
            'Beginner': [
                {'th': 'ฮีโร่น้อยคนใหม่ของเมือง...', 'en': 'The city\'s new young hero...'},
                {'th': 'พลังพิเศษของฉันทำให้...', 'en': 'My special power made me...'},
                {'th': 'เมื่อต้องช่วยเหลือแมวตัวน้อย...', 'en': 'When I had to save a little cat...'}
            ],
            'Intermediate': [
                {'th': 'ภารกิจลับของทีมฮีโร่...', 'en': 'The hero team\'s secret mission...'},
                {'th': 'พลังใหม่ที่น่าประหลาดใจ...', 'en': 'A surprising new power...'},
                {'th': 'การช่วยเหลือครั้งสำคัญ...', 'en': 'An important rescue mission...'}
            ],
            'Advanced': [
                {'th': 'ความรับผิดชอบของการเป็นฮีโร่...', 'en': 'The responsibility of being a hero...'},
                {'th': 'เมื่อเมืองต้องการฮีโร่...', 'en': 'When the city needed a hero...'},
                {'th': 'การต่อสู้เพื่อความยุติธรรม...', 'en': 'Fighting for justice...'}
            ]
        },
        'background_color': '#FFE0B2',
        'accent_color': '#F57C00'
    },
    'mystery': {
        'id': 'mystery',
        'icon': '🔍',
        'name_en': 'Mystery & Detective',
        'name_th': 'ไขปริศนาและนักสืบ',
        'description_th': 'สืบสวนปริศนาและไขความลับต่างๆ',
        'description_en': 'Solve mysteries and uncover secrets',
        'level_range': ['Beginner', 'Intermediate', 'Advanced'],
        'vocabulary': {
            'Beginner': ['clue', 'find', 'look', 'search', 'mystery', 'hidden', 'secret', 'detective'],
            'Intermediate': ['investigate', 'evidence', 'puzzle', 'solve', 'discover', 'suspicious', 'case'],
            'Advanced': ['investigation', 'deduction', 'enigma', 'cryptic', 'mysterious', 'revelation']
        },
        'story_starters': {
            'Beginner': [
                {'th': 'มีรอยปริศนาในสวน...', 'en': 'There were mysterious footprints in the garden...'},
                {'th': 'จดหมายลึกลับถูกส่งมา...', 'en': 'A mysterious letter arrived...'},
                {'th': 'ของเล่นหายไปอย่างลึกลับ...', 'en': 'The toy mysteriously disappeared...'}
            ],
            'Intermediate': [
                {'th': 'เบาะแสชิ้นแรกนำไปสู่...', 'en': 'The first clue led to...'},
                {'th': 'ความลับในห้องเก่า...', 'en': 'The secret in the old room...'},
                {'th': 'รหัสลับถูกค้นพบ...', 'en': 'A secret code was found...'}
            ],
            'Advanced': [
                {'th': 'คดีปริศนาที่ยากที่สุด...', 'en': 'The most challenging mystery case...'},
                {'th': 'ความลับที่ซ่อนอยู่มานาน...', 'en': 'A long-hidden secret...'},
                {'th': 'การสืบสวนนำไปสู่การค้นพบ...', 'en': 'The investigation led to a discovery...'}
            ]
        },
        'background_color': '#E0E0E0',
        'accent_color': '#616161'
    },
    'science': {
        'id': 'science',
        'icon': '🔬',
        'name_en': 'Science & Discovery',
        'name_th': 'วิทยาศาสตร์และการค้นพบ',
        'description_th': 'การทดลองและค้นพบทางวิทยาศาสตร์ที่น่าตื่นเต้น',
        'description_en': 'Exciting scientific experiments and discoveries',
        'level_range': ['Beginner', 'Intermediate', 'Advanced'],
        'vocabulary': {
            'Beginner': ['experiment', 'science', 'lab', 'test', 'mix', 'observe', 'change', 'result'],
            'Intermediate': ['hypothesis', 'research', 'discovery', 'invention', 'laboratory', 'scientist'],
            'Advanced': ['innovation', 'technological', 'breakthrough', 'analysis', 'investigation']
        },
        'story_starters': {
            'Beginner': [
                {'th': 'การทดลองง่ายๆ เริ่มต้นด้วย...', 'en': 'The simple experiment started with...'},
                {'th': 'ในห้องทดลองมีสิ่งมหัศจรรย์...', 'en': 'In the lab, there was something amazing...'},
                {'th': 'เมื่อผสมสองสิ่งเข้าด้วยกัน...', 'en': 'When mixing the two things together...'}
            ],
            'Intermediate': [
                {'th': 'การค้นพบที่น่าประหลาดใจ...', 'en': 'A surprising discovery...'},
                {'th': 'สิ่งประดิษฐ์ใหม่ทำให้...', 'en': 'The new invention made...'},
                {'th': 'การทดลองที่ไม่คาดคิด...', 'en': 'An unexpected experiment...'}
            ],
            'Advanced': [
                {'th': 'นวัตกรรมที่จะเปลี่ยนโลก...', 'en': 'Innovation that would change the world...'},
                {'th': 'การค้นพบทางวิทยาศาสตร์ครั้งสำคัญ...', 'en': 'An important scientific discovery...'},
                {'th': 'เทคโนโลยีใหม่ที่น่าทึ่ง...', 'en': 'Amazing new technology...'}
            ]
        },
        'background_color': '#E8EAF6',
        'accent_color': '#3F51B5'
    }
}

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# === 2. CONFIGURATIONS ===
# Page Config
st.set_page_config(
    page_title="JoyStory - Interactive Story Adventure",
    page_icon="📖",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Initialize OpenAI client
try:
    client = OpenAI()
except Exception as e:
    logging.error(f"Failed to initialize OpenAI client: {str(e)}")
    st.error("Failed to initialize AI services. Please try again later.")

# Define Constants
MAX_RETRIES = 3
DEFAULT_LEVEL = 'Beginner'
SUPPORTED_LANGUAGES = ['th', 'en']

# Achievement Requirements
ACHIEVEMENT_THRESHOLDS = {
    'perfect_writer': 5,    # 5 correct sentences in a row
    'vocabulary_master': 50, # 50 unique words used
    'story_master': 10,     # 10 sentences in story
    'accuracy_king': 80     # 80% accuracy rate
}

# Level Configuration
level_options = {
    'Beginner': {
        'thai_name': 'ระดับเริ่มต้น (ป.1-3)',
        'age_range': '7-9 ปี',
        'description': 'เหมาะสำหรับน้องๆ ที่เริ่มเรียนรู้การเขียนประโยคภาษาอังกฤษ',
        'features': [
            'ประโยคสั้นๆ ง่ายๆ',
            'คำศัพท์พื้นฐานที่ใช้ในชีวิตประจำวัน',
            'มีคำแนะนำภาษาไทยละเอียด',
            'เน้นการใช้ Present Simple Tense'
        ],
        'max_sentence_length': 10,
        'allowed_tenses': ['present_simple'],
        'feedback_level': 'detailed'
    },
    'Intermediate': {
        'thai_name': 'ระดับกลาง (ป.4-6)',
        'age_range': '10-12 ปี',
        'description': 'เหมาะสำหรับน้องๆ ที่สามารถเขียนประโยคพื้นฐานได้แล้ว',
        'features': [
            'ประโยคซับซ้อนขึ้น',
            'เริ่มใช้ Past Tense ได้',
            'คำศัพท์หลากหลายขึ้น',
            'สามารถเขียนเรื่องราวต่อเนื่องได้'
        ],
        'max_sentence_length': 15,
        'allowed_tenses': ['present_simple', 'past_simple'],
        'feedback_level': 'moderate'
    },
    'Advanced': {
        'thai_name': 'ระดับก้าวหน้า (ม.1-3)',
        'age_range': '13-15 ปี',
        'description': 'เหมาะสำหรับน้องๆ ที่มีพื้นฐานภาษาอังกฤษดี',
        'features': [
            'เขียนเรื่องราวได้หลากหลายรูปแบบ',
            'ใช้ Tense ต่างๆ ได้',
            'คำศัพท์ระดับสูงขึ้น',
            'สามารถแต่งเรื่องที่ซับซ้อนได้'
        ],
        'max_sentence_length': 20,
        'allowed_tenses': ['present_simple', 'past_simple', 'present_perfect', 'past_perfect'],
        'feedback_level': 'concise'
    }
}

# Achievement Configuration
achievements_list = {
    'perfect_writer': {
        'name': '🌟 นักเขียนไร้ที่ติ',
        'description': 'เขียนถูกต้อง 5 ประโยคติดต่อกัน',
        'condition': lambda: st.session_state.points['streak'] >= ACHIEVEMENT_THRESHOLDS['perfect_writer']
    },
    'vocabulary_master': {
        'name': '📚 ราชาคำศัพท์',
        'description': 'ใช้คำศัพท์ไม่ซ้ำกัน 50 คำ',
        'condition': lambda: len(st.session_state.stats['vocabulary_used']) >= ACHIEVEMENT_THRESHOLDS['vocabulary_master']
    },
    'story_master': {
        'name': '📖 นักแต่งนิทาน',
        'description': 'เขียนเรื่องยาว 10 ประโยค',
        'condition': lambda: len(st.session_state.story) >= ACHIEVEMENT_THRESHOLDS['story_master']
    },
    'accuracy_king': {
        'name': '👑 ราชาความแม่นยำ',
        'description': 'มีอัตราความถูกต้อง 80% ขึ้นไป (อย่างน้อย 10 ประโยค)',
        'condition': lambda: (
            st.session_state.stats['total_sentences'] >= 10 and 
            st.session_state.stats['accuracy_rate'] >= ACHIEVEMENT_THRESHOLDS['accuracy_king']
        )
    }
}

# Initial CSS Setup
st.markdown("""
    <style>
    /* Hide audio elements */
    .hide-tag {
        display: none !important;
    }
    
    /* Base Styles */
    @import url('https://fonts.googleapis.com/css2?family=Sarabun:wght@400;600&display=swap');
    
    body {
        font-family: 'Sarabun', sans-serif;
    }
    
    /* Custom Classes */
    .thai-eng {
        font-size: 1.1em;
        padding: 10px;
        background-color: #f8f9fa;
        border-radius: 8px;
        margin: 10px 0;
    }

    .theme-card {
        transition: all 0.3s ease;
    }
    
    .theme-card:hover {
        transform: translateY(-5px);
        box-shadow: 0 4px 15px rgba(0,0,0,0.1);
    }
    
    .theme-header {
        text-align: center;
        margin-bottom: 2rem;
    }
    
    .theme-header h2 {
        font-size: 2rem;
        color: #1e88e5;
        margin-bottom: 1rem;
    }
    
    .theme-header p {
        color: #666;
        font-size: 1.1rem;
    }
    
    /* ปรับแต่งปุ่มเลือกธีม */
    .stButton > button {
        width: 100%;
        padding: 0.75rem;
        border: none;
        border-radius: 8px;
        font-family: 'Sarabun', sans-serif;
        transition: all 0.2s ease;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 2px 8px rgba(0,0,0,0.1);
    }
    
    .thai {
        color: #1e88e5;
    }
    
    .eng {
        color: #333;
    }
    
    /* Error Messages */
    .error-message {
        color: #d32f2f;
        padding: 10px;
        border-left: 4px solid #d32f2f;
        background-color: #ffebee;
        margin: 10px 0;
    }
    
    /* Success Messages */
    .success-message {
        color: #2e7d32;
        padding: 10px;
        border-left: 4px solid #2e7d32;
        background-color: #e8f5e9;
        margin: 10px 0;
    }
    </style>
""", unsafe_allow_html=True)

# Reset counter when page reloads
if 'theme_button_counter' in st.session_state:
    st.session_state.theme_button_counter = 0

# === 3. STATE MANAGEMENT ===
def init_session_state():
    """Initialize all session state variables with default values"""
    
    # Basic states
    default_states = {
        'theme_selection_id': datetime.now().strftime('%Y%m%d%H%M%S'),
        'current_theme': None,
        'theme_button_counter': 0,
        'theme_story_starter': None,
        'story': [],
        'feedback': None,
        'level': DEFAULT_LEVEL,
        'should_reset': False,
        'last_interaction': datetime.now().isoformat(),
    }
    
    # Points system
    points_states = {
        'points': {
            'total': 0,
            'perfect_sentences': 0,
            'corrections_made': 0,
            'streak': 0,
            'max_streak': 0
        }
    }
    
    # Statistics tracking
    stats_states = {
        'stats': {
            'total_sentences': 0,
            'correct_first_try': 0,
            'accuracy_rate': 0.0,
            'vocabulary_used': set(),
            'corrections_made': 0,
            'average_sentence_length': 0,
            'total_words': 0,
            'session_duration': 0
        }
    }
    
    # Game progress
    progress_states = {
        'achievements': [],
        'unlocked_features': set(),
        'current_milestone': 0,
        'next_milestone': 5
    }
    
    # User preferences
    preferences_states = {
        'language': 'th',
        'feedback_level': 'detailed',
        'theme_color': 'light',
        'sound_enabled': True
    }
    
    # Initialize all states if they don't exist
    for state_dict in [default_states, points_states, stats_states, progress_states, preferences_states]:
        for key, value in state_dict.items():
            if key not in st.session_state:
                st.session_state[key] = value

    if 'clear_input' not in st.session_state:
        st.session_state.clear_input = False
    if 'text_input' not in st.session_state:
        st.session_state.text_input = ""
    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"""
    if 'current_theme' not in st.session_state:
        st.session_state.current_theme = None
    if 'theme_story_starter' not in st.session_state:
        st.session_state.theme_story_starter = None

def reset_story():
    """Reset story and related state variables completely"""
    try:
        # Reset all story-related states
        new_states = {
            'story': [],
            'feedback': None,
            'theme_story_starter': None,
            'text_input': "",
            'current_theme': None,
            'ending_mode': False,
            'sentences_to_end': 0,
            'ending_type': None,
            'story_completed': False,
            'stitched_story': None,
            'clear_input': True,
            'last_submission': None
        }
        
        # Reset points
        new_points = {
            'total': 0,
            'perfect_sentences': 0,
            'corrections_made': 0,
            'streak': 0,
            'max_streak': 0
        }
        
        # Reset stats
        new_stats = {
            'total_sentences': 0,
            'correct_first_try': 0,
            'accuracy_rate': 0.0,
            'vocabulary_used': set(),
            'corrections_made': 0,
            'average_sentence_length': 0,
            'total_words': 0,
            'session_duration': 0
        }
        
        # Reset all session states
        for key, value in new_states.items():
            st.session_state[key] = value
        st.session_state.points = new_points
        st.session_state.stats = new_stats
        
        # Reset achievements and progress
        st.session_state.achievements = []
        st.session_state.current_milestone = 0
        st.session_state.next_milestone = 5
        
        # Reset UI elements
        if 'story_input_area' in st.session_state:
            del st.session_state.story_input_area
        
        # Force new input field creation
        st.session_state.clear_count = st.session_state.get('clear_count', 0) + 1
        
        # Reset theme selection
        st.session_state.theme_button_counter = 0
        st.session_state.theme_selection_id = datetime.now().strftime('%Y%m%d%H%M%S')
        
        # Log reset
        logging.info("Story state reset successfully")
        
        return True
    
    except Exception as e:
        logging.error(f"Error resetting story state: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการรีเซ็ตเรื่อง กรุณาลองใหม่อีกครั้ง")
        return False

def save_progress() -> Dict:
    """Save current progress to JSON format"""
    try:
        progress_data = {
            'timestamp': datetime.now().isoformat(),
            'level': st.session_state.level,
            '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,
            'current_theme': st.session_state.current_theme
        }
        
        logging.info("Progress saved successfully")
        return progress_data
    
    except Exception as e:
        logging.error(f"Error saving progress: {str(e)}")
        raise

def load_progress(data: Dict):
    """Load progress from saved data"""
    try:
        # Validate data structure
        required_keys = ['level', 'story', 'stats', 'points', 'achievements']
        if not all(key in data for key in required_keys):
            raise ValueError("Invalid save data format")
        
        # Load basic data
        st.session_state.level = data['level']
        st.session_state.story = data['story']
        st.session_state.achievements = data['achievements']
        st.session_state.points = data['points']
        
        # Load stats (converting lists back to sets where needed)
        st.session_state.stats = {
            'total_sentences': data['stats']['total_sentences'],
            'correct_first_try': data['stats']['correct_first_try'],
            'accuracy_rate': data['stats']['accuracy_rate'],
            'vocabulary_used': set(data['stats']['vocabulary_used']),
            'corrections_made': data['stats']['corrections_made'],
            'average_sentence_length': data['stats'].get('average_sentence_length', 0),
            'total_words': data['stats'].get('total_words', 0),
            'session_duration': data['stats'].get('session_duration', 0)
        }
        
        # Load theme if present
        if 'current_theme' in data:
            st.session_state.current_theme = data['current_theme']
        
        logging.info("Progress loaded successfully")
        st.success("โหลดความก้าวหน้าเรียบร้อย!")
        
        return True
    
    except Exception as e:
        logging.error(f"Error loading progress: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการโหลดข้อมูล กรุณาตรวจสอบไฟล์และลองใหม่อีกครั้ง")
        return False

def update_session_stats():
    """Update session statistics"""
    try:
        if st.session_state.story:
            # Update word counts
            all_text = ' '.join([entry['content'] for entry in st.session_state.story])
            words = all_text.split()
            st.session_state.stats['total_words'] = len(words)
            
            # Update average sentence length
            if st.session_state.stats['total_sentences'] > 0:
                st.session_state.stats['average_sentence_length'] = (
                    st.session_state.stats['total_words'] / 
                    st.session_state.stats['total_sentences']
                )
            
            # Update session duration
            start_time = datetime.fromisoformat(st.session_state.last_interaction)
            current_time = datetime.now()
            duration = (current_time - start_time).total_seconds()
            st.session_state.stats['session_duration'] = duration
            
            # Update last interaction time
            st.session_state.last_interaction = current_time.isoformat()
            
        return True
    
    except Exception as e:
        logging.error(f"Error updating session stats: {str(e)}")
        return False

# === 4. UTILITY FUNCTIONS ===

def initialize_audio():
    """Initialize audio system"""
    if 'audio_manager' not in st.session_state:
        audio_manager = AudioManager()
        st.session_state.audio_manager = audio_manager
        st.session_state.sound_commands = get_sound_commands()
        
        # Add audio elements to the page using st.markdown with hide-tag class
        st.markdown("""
            <style>
            .hide-tag {
                display: none;
            }
            </style>
        """, unsafe_allow_html=True)
        
        st.markdown(
            f"""
            <div class="hide-tag">
                {audio_manager.get_audio_html()}
            </div>
            """,
            unsafe_allow_html=True
        )

def show_audio_controls():
    """Show audio control in sidebar"""
    st.sidebar.markdown("### 🔊 ตั้งค่าเสียง")
    
    # Background music toggle
    if st.sidebar.checkbox("เปิดเพลงประกอบ", value=True, key='bgm_enabled'):
        st.markdown(f"<script>{st.session_state.sound_commands['bgm_play']}</script>", unsafe_allow_html=True)
    else:
        st.markdown(f"<script>{st.session_state.sound_commands['bgm_pause']}</script>", unsafe_allow_html=True)
    
    # Volume control
    volume = st.sidebar.slider("ระดับเสียง", 0, 100, 50, key='volume') / 100
    st.markdown(f"<script>setVolume({volume});</script>", unsafe_allow_html=True)
    
def generate_story_continuation(user_input: str, level: str) -> str:
    """Generate AI story continuation using ChatGPT"""
    
    level_context = {
        'Beginner': {
            'instructions': """
            Role: Teaching assistant for Thai students (grades 1-3)
            Rules:
            - Use only 1-2 VERY simple sentences
            - Use Present Simple Tense only
            - Basic vocabulary (family, school, daily activities)
            - 5-7 words per sentence maximum
            - Focus on clear, basic responses
            """,
            'max_tokens': 30,
            'temperature': 0.6
        },
        'Intermediate': {
            'instructions': """
            Role: Teaching assistant for Thai students (grades 4-6)
            Rules:
            - 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': """
            Role: Teaching assistant for Thai students (grades 7-9)
            Rules:
            - 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
        }
    }

    try:
        # Get recent story context (last 5 entries for more coherence)
        story_context = '\n'.join([
            entry['content'] for entry in st.session_state.story[-5:]
        ]) if st.session_state.story else "Story just started"
        
        # Create prompt
        level_settings = level_context[level]
        
        for _ in range(MAX_RETRIES):
            try:
                response = client.chat.completions.create(
                    model="gpt-4o-mini",
                    messages=[
                        {
                            "role": "system",
                            "content": f"""
                            {level_settings['instructions']}
                            
                            CRUCIAL GUIDELINES:
                            - Never exceed the maximum sentences for the level
                            - Create openings for student's creativity
                            - Do not resolve plot points or conclude the story
                            - Avoid using 'suddenly' or 'then'
                            - Make each sentence meaningful but incomplete
                            - Leave room for the student to develop the story
                            """
                        },
                        {
                            "role": "user",
                            "content": f"Story context:\n{story_context}\nStudent's input:\n{user_input}\nProvide a brief continuation:"
                        }
                    ],
                    max_tokens=level_settings['max_tokens'],
                    temperature=level_settings['temperature'],
                    presence_penalty=0.6,
                    frequency_penalty=0.6
                )
                
                # Process and clean response
                response_text = response.choices[0].message.content.strip()
                sentences = [s.strip() for s in response_text.split('.') if s.strip()]
                
                # Limit sentences based on level
                max_sentences = {'Beginner': 2, 'Intermediate': 2, 'Advanced': 3}
                if len(sentences) > max_sentences[level]:
                    sentences = sentences[:max_sentences[level]]
                
                # Reconstruct response
                final_response = '. '.join(sentences) + '.'
                
                logging.info(f"Generated continuation for level {level}")
                return final_response
                
            except Exception as e:
                if _ < MAX_RETRIES - 1:
                    logging.warning(f"Retry {_+1} failed: {str(e)}")
                    continue
                raise
                
        raise Exception("Max retries exceeded")
        
    except Exception as e:
        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.
    Returns both Thai and English versions.
    """
    try:
        # Get theme data
        theme = story_themes.get(theme_id)
        if not theme:
            raise ValueError(f"Theme {theme_id} not found")

        # Get random starter for the level
        level_starters = theme['story_starters'].get(level, [])
        if not level_starters:
            # Fallback to Beginner level if no starters found for specified level
            level_starters = theme['story_starters'].get('Beginner', [])
        
        if not level_starters:
            raise ValueError(f"No story starters found for theme {theme_id}")
        
        # Select random starter
        starter = random.choice(level_starters)
        
        # Return both languages
        return {
            'en': starter['en'],
            'th': starter['th']
        }
        
    except Exception as e:
        logging.error(f"Error generating story starter: {str(e)}")
        # Provide fallback starter
        return {
            'en': 'Once upon a time...',
            'th': 'กาลครั้งหนึ่ง...'
        }

def update_points(is_correct_first_try: bool):
    """อัพเดตคะแนนตามผลการเขียน"""
    try:
        # คำนวณคะแนนพื้นฐาน
        base_points = 10
        
        if is_correct_first_try:
            # ถูกต้องในครั้งแรก
            points = base_points * 2
            st.session_state.points['perfect_sentences'] += 1
            st.session_state.points['streak'] += 1
            
            # อัพเดต max streak
            if st.session_state.points['streak'] > st.session_state.points['max_streak']:
                st.session_state.points['max_streak'] = st.session_state.points['streak']
        else:
            # ต้องแก้ไข
            points = base_points // 2
            st.session_state.points['corrections_made'] += 1
            st.session_state.points['streak'] = 0
        
        # เพิ่มคะแนนรวม
        st.session_state.points['total'] += points
        
        # อัพเดตสถิติ
        st.session_state.stats['total_sentences'] += 1
        if is_correct_first_try:
            st.session_state.stats['correct_first_try'] += 1
        
        # คำนวณอัตราความถูกต้อง
        st.session_state.stats['accuracy_rate'] = (
            st.session_state.stats['correct_first_try'] / 
            st.session_state.stats['total_sentences'] * 100
        )
        
        logging.info(f"Points updated: +{points} points")
        return True
        
    except Exception as e:
        logging.error(f"Error updating points: {str(e)}")
        return False

def apply_correction(story_index: int, corrected_text: str):
    """แก้ไขประโยคในเรื่อง"""
    try:
        if not (0 <= story_index < len(st.session_state.story)):
            raise ValueError("Invalid story index")
            
        # เก็บข้อความเดิม
        original_text = st.session_state.story[story_index]['content']
        
        # อัพเดทข้อความ
        st.session_state.story[story_index].update({
            'content': corrected_text,
            'is_corrected': True,
            'is_correct': True,
            'original_text': original_text,
            'correction_timestamp': datetime.now().isoformat()
        })
        
        # อัพเดทสถิติ
        st.session_state.stats['corrections_made'] += 1
        
        # Log การแก้ไข
        logging.info(f"Sentence corrected at index {story_index}")
        
        # แสดงข้อความยืนยัน
        st.success("✅ แก้ไขประโยคเรียบร้อยแล้ว!")
        return True
        
    except Exception as e:
        logging.error(f"Error applying correction: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการแก้ไขประโยค กรุณาลองใหม่อีกครั้ง")
        return False

def update_achievements():
    """ตรวจสอบและอัพเดตความสำเร็จ"""
    try:
        current_achievements = st.session_state.achievements
        
        # ตรวจสอบเงื่อนไขต่างๆ
        if (st.session_state.points['streak'] >= 5 and 
            "🌟 นักเขียนไร้ที่ติ" not in current_achievements):
            current_achievements.append("🌟 นักเขียนไร้ที่ติ")
            st.success("🎉 ได้รับความสำเร็จใหม่: นักเขียนไร้ที่ติ!")
            st.markdown(f"<script>{st.session_state.sound_commands['achievement']}</script>", unsafe_allow_html=True)
        
        if (len(st.session_state.stats['vocabulary_used']) >= 50 and 
            "📚 ราชาคำศัพท์" not in current_achievements):
            current_achievements.append("📚 ราชาคำศัพท์")
            st.success("🎉 ได้รับความสำเร็จใหม่: ราชาคำศัพท์!")
            st.markdown(f"<script>{st.session_state.sound_commands['achievement']}</script>", unsafe_allow_html=True)
        
        if (len(st.session_state.story) >= 10 and 
            "📖 นักแต่งนิทาน" not in current_achievements):
            current_achievements.append("📖 นักแต่งนิทาน")
            st.success("🎉 ได้รับความสำเร็จใหม่: นักแต่งนิทาน!")
            st.markdown(f"<script>{st.session_state.sound_commands['achievement']}</script>", unsafe_allow_html=True)
        
        if (st.session_state.stats['total_sentences'] >= 10 and 
            st.session_state.stats['accuracy_rate'] >= 80 and 
            "👑 ราชาความแม่นยำ" not in current_achievements):
            current_achievements.append("👑 ราชาความแม่นยำ")
            st.success("🎉 ได้รับความสำเร็จใหม่: ราชาความแม่นยำ!")
            st.markdown(f"<script>{st.session_state.sound_commands['achievement']}</script>", unsafe_allow_html=True)
        
        # บันทึกความสำเร็จ
        st.session_state.achievements = current_achievements
        logging.info("Achievements updated successfully")
        return True
        
    except Exception as e:
        logging.error(f"Error updating achievements: {str(e)}")
        return False

def provide_feedback(text: str, level: str) -> Dict[str, str]:
    """Provide more detailed feedback on the user's writing with appropriate level"""
    
    level_prompts = {
        'Beginner': """
        Focus on:
        - Basic sentence structure (Subject + Verb + Object)
        - Present Simple Tense usage
        - Basic vocabulary
        - Capitalization and periods
        Provide very simple, encouraging feedback in Thai.
        Explain any mistakes and how to correct them.
        """,
        'Intermediate': """
        Focus on:
        - Sentence variety
        - Past Tense usage
        - Vocabulary appropriateness
        - Basic punctuation
        - Simple conjunctions
        Provide moderately detailed feedback in Thai.
        Explain any mistakes and suggest improvements.
        """,
        'Advanced': """
        Focus on:
        - Complex sentence structures
        - Various tense usage
        - Advanced vocabulary
        - All punctuation
        - Style and flow
        Provide comprehensive feedback in Thai.
        Explain any mistakes and provide examples for improvement.
        """
    }

    try:
        for _ in range(MAX_RETRIES):
            try:
                response = client.chat.completions.create(
                    model="gpt-4o-mini",
                    messages=[
                        {
                            "role": "system",
                            "content": f"""
                            You are a Thai English teacher helping {level} students.
                            {level_prompts[level]}
                            
                            Return your response in this EXACT format (valid JSON):
                            {{
                                "feedback": "(ข้อเสนอแนะภาษาไทย)",
                                "corrected": "(ประโยคภาษาอังกฤษที่ถูกต้อง)",
                                "has_errors": true/false,
                                "error_types": ["grammar", "vocabulary", "spelling", ...],
                                "difficulty_score": 1-10
                            }}
                            """
                        },
                        {
                            "role": "user",
                            "content": f"Review this sentence: {text}"
                        }
                    ],
                    max_tokens=300,  # Increased from 200
                    temperature=0.5  # Adjusted for balanced variety
                )
                
                # Parse response
                feedback_data = json.loads(response.choices[0].message.content.strip())
                
                # Validate required fields
                required_fields = ['feedback', 'corrected', 'has_errors']
                if not all(field in feedback_data for field in required_fields):
                    raise ValueError("Missing required fields in feedback")
                
                logging.info(f"Generated feedback for level {level}")
                return feedback_data
                
            except json.JSONDecodeError:
                if _ < MAX_RETRIES - 1:
                    continue
                raise
                
        raise Exception("Max retries exceeded")
        
    except Exception as e:
        logging.error(f"Error generating feedback: {str(e)}")
        return {
            "feedback": "⚠️ ระบบไม่สามารถวิเคราะห์ประโยคได้ กรุณาลองใหม่อีกครั้ง",
            "corrected": text,
            "has_errors": False,
            "error_types": [],
            "difficulty_score": 5
        }


def get_vocabulary_suggestions(context: str = "", level: str = DEFAULT_LEVEL) -> List[str]:
    """Get contextual vocabulary suggestions with Thai translations"""
    try:
        recent_story = context or '\n'.join([
            entry['content'] for entry in st.session_state.story[-3:]
        ]) if st.session_state.story else "Story just started"
        
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": f"""
                    You are a Thai-English bilingual teacher.
                    Suggest 5 English words appropriate for {level} level students.
                    Format each suggestion as:
                    word (type) - คำแปล | example sentence
                    """
                },
                {
                    "role": "user",
                    "content": f"Story context:\n{recent_story}\n\nSuggest relevant words:"
                }
            ],
            max_tokens=200,
            temperature=0.8
        )
        
        suggestions = response.choices[0].message.content.split('\n')
        return [s.strip() for s in suggestions if s.strip()]
        
    except Exception as e:
        logging.error(f"Error getting vocabulary suggestions: {str(e)}")
        return [
            "happy (adj) - มีความสุข | I am happy today",
            "run (verb) - วิ่ง | The dog runs fast",
            "tree (noun) - ต้นไม้ | A tall tree"
        ]

def get_creative_prompt() -> Dict[str, str]:
    """Generate a bilingual creative writing prompt"""
    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": """
                    Create short story prompts in both English and Thai.
                    Keep it simple and under 6 words each.
                    Make it open-ended and encouraging.
                    """
                },
                {
                    "role": "user",
                    "content": "Generate a creative writing prompt:"
                }
            ],
            max_tokens=50,
            temperature=0.7
        )
        
        prompt_eng = response.choices[0].message.content.strip()
        
        # Get Thai translation
        response_thai = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": "Translate to short Thai prompt, keep it natural:"
                },
                {
                    "role": "user",
                    "content": prompt_eng
                }
            ],
            max_tokens=50,
            temperature=0.7
        )
        
        prompt_thai = response_thai.choices[0].message.content.strip()
        
        return {
            "eng": prompt_eng,
            "thai": prompt_thai
        }
        
    except Exception as e:
        logging.error(f"Error generating creative prompt: {str(e)}")
        return {
            "eng": "What happens next?",
            "thai": "แล้วอะไรจะเกิดขึ้นต่อ?"
        }

def calculate_text_metrics(text: str) -> Dict[str, float]:
    """Calculate various metrics for the given text"""
    try:
        words = text.split()
        sentences = [s.strip() for s in text.split('.') if s.strip()]
        
        metrics = {
            'word_count': len(words),
            'sentence_count': len(sentences),
            'average_word_length': sum(len(word) for word in words) / len(words) if words else 0,
            'unique_words': len(set(words)),
            'complexity_score': calculate_complexity_score(text)
        }
        
        return metrics
        
    except Exception as e:
        logging.error(f"Error calculating text metrics: {str(e)}")
        return {
            'word_count': 0,
            'sentence_count': 0,
            'average_word_length': 0,
            'unique_words': 0,
            'complexity_score': 0
        }

def calculate_complexity_score(text: str) -> float:
    """Calculate a complexity score for the text (0-10)"""
    try:
        # Basic metrics
        words = text.split()
        word_count = len(words)
        unique_words = len(set(words))
        avg_word_length = sum(len(word) for word in words) / word_count if word_count > 0 else 0
        
        # Calculate score components
        vocabulary_score = (unique_words / word_count) * 5 if word_count > 0 else 0
        length_score = min((avg_word_length / 10) * 5, 5)
        
        # Combine scores
        total_score = vocabulary_score + length_score
        
        return min(total_score, 10.0)
        
    except Exception as e:
        logging.error(f"Error calculating complexity score: {str(e)}")
        return 5.0

# === 5. UI COMPONENTS ===
def show_welcome_section():
    """Display welcome message and introduction"""
    st.markdown("""
        <div class="welcome-header" style="text-align: center; padding: 20px;">
            <div class="thai" style="font-size: 1.5em; margin-bottom: 10px;">
            🌟 ยินดีต้อนรับสู่ JoyStory - แอพฝึกเขียนภาษาอังกฤษแสนสนุก!
            </div>
            <div style="color: #666; margin-bottom: 20px;">
            เรียนรู้ภาษาอังกฤษผ่านการเขียนเรื่องราวด้วยตัวเอง 
            พร้อมผู้ช่วย AI ที่จะช่วยแนะนำและให้กำลังใจ
            </div>
            <div class="eng" style="font-style: italic; color: #1e88e5;">
            Welcome to JoyStory - Fun English Writing Adventure!
            </div>
        </div>
    """, unsafe_allow_html=True)

def show_parent_guide():
    """Display guide for parents"""
    with st.expander("📖 คำแนะนำสำหรับผู้ปกครอง | Parent's Guide"):
        st.markdown("""
            <div class="parent-guide" style="background-color: #fff3e0; padding: 20px; border-radius: 10px;">
                <h4 style="color: #f57c00; margin-bottom: 15px;">คำแนะนำในการใช้งาน</h4>
                <ul style="list-style-type: none; padding-left: 0;">
                    <li style="margin-bottom: 10px;">
                        👥 <strong>การมีส่วนร่วม:</strong> แนะนำให้นั่งเขียนเรื่องราวร่วมกับน้องๆ
                    </li>
                    <li style="margin-bottom: 10px;">
                        💡 <strong>การช่วยเหลือ:</strong> ช่วยอธิบายคำแนะนำและคำศัพท์ที่น้องๆ ไม่เข้าใจ
                    </li>
                    <li style="margin-bottom: 10px;">
                        🌟 <strong>การให้กำลังใจ:</strong> ให้กำลังใจและชื่นชมเมื่อน้องๆ เขียนได้ดี
                    </li>
                    <li style="margin-bottom: 10px;">
                        ⏱️ <strong>เวลาที่เหมาะสม:</strong> ใช้เวลาในการเขียนแต่ละครั้งไม่เกิน 20-30 นาที
                    </li>
                </ul>
                <div style="background-color: #ffe0b2; padding: 15px; border-radius: 8px; margin-top: 15px;">
                    <p style="margin: 0;">
                        💡 <strong>เคล็ดลับ:</strong> ให้น้องๆ พูดเรื่องราวที่อยากเขียนเป็นภาษาไทยก่อน 
                        แล้วค่อยๆ ช่วยกันแปลงเป็นภาษาอังกฤษ
                    </p>
                </div>
            </div>
        """, unsafe_allow_html=True)

def show_level_selection():
    """Display level selection interface"""
    st.markdown("""
        <div class="thai" style="margin-bottom: 15px;">
        🎯 เลือกระดับการเรียนรู้
        </div>
    """, unsafe_allow_html=True)

    level = st.radio(
        "ระดับการเรียน",
        options=list(level_options.keys()),
        format_func=lambda x: level_options[x]['thai_name'],
        key="level_selector",
        label_visibility="collapsed"
    )
    
    # Show level details
    st.markdown(f"""
        <div class="level-info" style="background-color: #e3f2fd; padding: 15px; border-radius: 10px; margin: 15px 0;">
            <h4 style="color: #1976d2; margin-bottom: 10px;">
                {level_options[level]['thai_name']}
            </h4>
            <p style="color: #444; margin-bottom: 10px;">
                🎯 {level_options[level]['description']}
            </p>
            <p style="color: #666;">
                👥 เหมาะสำหรับอายุ: {level_options[level]['age_range']}
            </p>
            <div style="margin-top: 15px;">
                <h5 style="color: #1976d2; margin-bottom: 8px;">✨ คุณลักษณะ:</h5>
                <ul style="list-style-type: none; padding-left: 0;">
                    {' '.join([f'<li style="margin-bottom: 5px;">• {feature}</li>' for feature in level_options[level]['features']])}
                </ul>
            </div>
        </div>
    """, unsafe_allow_html=True)

    return level

def show_theme_selection():
    """Display theme selection interface"""
    # Header section
    header_html = '''
        <div style="text-align: center; margin-bottom: 2rem;">
            <h2 style="color: #1e88e5; font-family: 'Sarabun', sans-serif; margin: 0; display: flex; align-items: center; justify-content: center; gap: 8px;">
                <span style="font-size: 1.5em;">🎨</span>
                <span>เลือกธีมเรื่องราว | Choose Story Theme</span>
            </h2>
            <p style="color: #666; max-width: 600px; margin: 1rem auto 0;">
                เลือกโลกแห่งจินตนาการที่คุณต้องการผจญภัย และเริ่มต้นเขียนเรื่องราวของคุณ
            </p>
        </div>
    '''
    st.markdown(header_html, unsafe_allow_html=True)

    # Filter available themes
    available_themes = [
        theme for theme in story_themes.values()
        if st.session_state.level in theme['level_range']
    ]

    # Create grid layout
    num_themes = len(available_themes)
    rows = (num_themes + 3) // 4  # Ceiling division to get number of rows needed
    
    # Create rows with 4 themes each
    for row in range(rows):
        cols = st.columns(4)
        for col_idx, col in enumerate(cols):
            theme_idx = row * 4 + col_idx
            if theme_idx < num_themes:
                theme = available_themes[theme_idx]
                with col:
                    # Theme card container
                    theme_card = f'''
                        <div style="
                            background-color: {theme['background_color']};
                            border-radius: 15px;
                            padding: 1.5rem;
                            margin-bottom: 1rem;
                            height: 100%;
                            position: relative;
                            transition: all 0.3s ease;
                            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
                        ">
                            <div style="position: absolute; top: 1rem; right: 1rem; 
                                      padding: 0.25rem 0.75rem; border-radius: 999px; 
                                      font-size: 0.8rem; background: #E3F2FD; 
                                      color: #1E88E5;">
                                {st.session_state.level}
                            </div>
                            <div style="font-size: 2.5rem; margin-bottom: 1rem;">
                                {theme['icon']}
                            </div>
                            <div style="font-size: 1.2rem; font-weight: 600; 
                                      margin-bottom: 0.5rem; color: {theme['accent_color']};">
                                {theme['name_th']}
                            </div>
                            <div style="font-size: 0.9rem; color: #666; 
                                      line-height: 1.4; margin-bottom: 1rem;">
                                {theme['description_th']}
                            </div>
                        </div>
                    '''
                    st.markdown(theme_card, unsafe_allow_html=True)
                    
                    # Theme selection button
                    if st.button(
                        f"เลือกธีม {theme['name_th']}", 
                        key=f"theme_{theme['id']}_{row}_{col_idx}",
                        use_container_width=True
                    ):
                        handle_theme_selection(theme)

def handle_theme_selection(theme: dict):
    """Handle theme selection and initialization"""
    try:
        with st.spinner("กำลังเตรียมเรื่องราว..."):
            # Reset story completely first
            reset_story()
            
            # Then set new theme
            st.session_state.current_theme = theme['id']
            starter = generate_dynamic_story_starter(
                theme['id'],
                st.session_state.level
            )
            st.session_state.story = [{
                "role": "AI",
                "content": starter['en'],
                "thai_content": starter['th'],
                "is_starter": True
            }]
            st.success(f"เลือกธีม {theme['name_th']} เรียบร้อยแล้ว!")
            st.rerun()
    except Exception as e:
        logging.error(f"Error selecting theme: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการเลือกธีม กรุณาลองใหม่อีกครั้ง")

def show_theme_card(theme: Dict):  # เปลี่ยนจาก display_theme_card เป็น show_theme_card
    """Display a single theme card with proper styling"""
    card_html = f"""
        <div style="
            background-color: {theme['background_color']};
            border-radius: 15px;
            padding: 1.5rem;
            margin-bottom: 1rem;
            height: 100%;
            position: relative;
            transition: all 0.3s ease;
        ">
            <div style="position: absolute; top: 1rem; right: 1rem; padding: 0.25rem 0.75rem; border-radius: 999px; font-size: 0.8rem; background: #E3F2FD; color: #1E88E5;">
                {st.session_state.level}
            </div>
            
            <div style="font-size: 2.5rem; margin-bottom: 1rem;">
                {theme['icon']}
            </div>
            
            <div style="font-size: 1.2rem; font-weight: 600; margin-bottom: 0.5rem; color: {theme['accent_color']};">
                {theme['name_th']}
            </div>
            
            <div style="font-size: 0.9rem; color: #666; line-height: 1.4; margin-bottom: 1rem;">
                {theme['description_th']}
            </div>
        </div>
    """
    
    # แสดงการ์ดธีม
    st.markdown(card_html, unsafe_allow_html=True)
    
    # แสดงปุ่มเลือกธีม
    if st.button(
        f"เลือกธีม {theme['name_th']}", 
        key=f"theme_{theme['id']}",
        help=f"คลิกเพื่อเริ่มเขียนเรื่องราวในธีม {theme['name_th']}",
        use_container_width=True
    ):
        try:
            with st.spinner("กำลังเตรียมเรื่องราว..."):
                st.session_state.current_theme = theme['id']
                starter = generate_dynamic_story_starter(
                    theme['id'],
                    st.session_state.level
                )
                st.session_state.story = [{
                    "role": "AI",
                    "content": starter['en'],
                    "thai_content": starter['th'],
                    "is_starter": True
                }]
                st.success(f"เลือกธีม {theme['name_th']} เรียบร้อยแล้ว!")
                st.rerun()
        except Exception as e:
            logging.error(f"Error selecting theme: {str(e)}")
            st.error("เกิดข้อผิดพลาดในการเลือกธีม กรุณาลองใหม่อีกครั้ง")

def show_story():
    """Display the story with proper formatting"""
    story_display = st.container()
    
    with story_display:
        # Clear previous story if resetting
        if st.session_state.get('should_reset'):
            reset_story()
            st.rerun()
            return
            
        if not st.session_state.story:
            st.info("เลือกธีมเรื่องราวที่ต้องการเพื่อเริ่มต้นการผจญภัย!")
            return
        
        # Story Display Box
        for idx, entry in enumerate(st.session_state.story):
            if entry['role'] == 'AI':
                if entry.get('is_starter'):
                    # Story Starter (คงเดิม)
                    st.markdown(f"""
                        <div style="
                            background-color: #f0f7ff; 
                            padding: 20px; 
                            border-radius: 10px; 
                            margin: 10px 0;
                            border-left: 4px solid #1e88e5;
                            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
                        ">
                            <div style="color: #1e88e5; margin-bottom: 10px; font-size: 1.1em;">
                                🎬 เริ่มเรื่อง:
                            </div>
                            <div style="color: #666; margin-bottom: 8px; font-family: 'Sarabun', sans-serif;">
                                {entry.get('thai_content', '')}
                            </div>
                            <div style="color: #333; font-size: 1.1em;">
                                {entry['content']}
                            </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:
                    # 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;
                            padding: 15px;
                            border-radius: 8px;
                            margin: 8px 0;
                            border-left: 4px solid #4caf50;
                        ">
                            <div style="color: #2e7d32;">
                                <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)
            
            elif entry['role'] == 'You':
                # User Entry
                status_icon = "✅" if entry.get('is_correct') else "✍️"
                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};
                        padding: 15px;
                        border-radius: 8px;
                        margin: 8px 0;
                        border-left: 4px solid {border_color};
                    ">
                        <div style="color: #333;">
                            <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)
        
        # Show completion message if story is done
        if st.session_state.get('story_completed'):
            st.markdown("""
                <div style="
                    background-color: #f3e5f5;
                    padding: 20px;
                    border-radius: 10px;
                    margin: 20px 0;
                    text-align: center;
                    border: 2px dashed #9c27b0;
                ">
                    <h2 style="color: #9c27b0; margin-bottom: 10px;">
                        🎉 ยินดีด้วย! คุณเขียนเรื่องราวจบสมบูรณ์แล้ว
                    </h2>
                    <p style="color: #666;">
                        คุณสามารถบันทึกเรื่องราวหรือเริ่มเรื่องใหม่ได้
                    </p>
                </div>
            """, 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: #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 generate_story_summary(story_entries: List[dict]) -> str:
    """Generate a concise summary of the story for the stitching process"""
    try:
        # Extract main story content
        story_text = []
        for entry in story_entries:
            if entry.get('content'):
                text = entry['content'].strip()
                if text.endswith('.'):
                    text = text[:-1]  # Remove trailing period for better concatenation
                story_text.append(text)
        
        story_summary = '. '.join(story_text) + '.'
        return story_summary
        
    except Exception as e:
        logging.error(f"Error generating story summary: {str(e)}")
        return "Error creating story summary."

def stitch_story(raw_story: List[dict], theme_id: str, level: str) -> Dict[str, str]:
    """Create a polished, coherent version of the story with Thai translation"""
    try:
        # Get theme details for context
        theme = story_themes.get(theme_id, {})
        theme_context = f"Theme: {theme.get('name_en', 'General')} - {theme.get('description_en', '')}"
        
        # Get story summary
        story_summary = generate_story_summary(raw_story)
        
        # Create level-appropriate instructions
        level_context = {
            'Beginner': {
                'instructions': """
                    Create a simple, clear narrative using:
                    - Basic vocabulary
                    - Short, straightforward sentences
                    - Present tense
                    - Clear transitions
                    Maintain the original story's key points but make it flow smoothly.
                """,
                'max_tokens': 500,
                'temperature': 0.7
            },
            'Intermediate': {
                'instructions': """
                    Create an engaging narrative using:
                    - Grade-appropriate vocabulary
                    - Mix of simple and compound sentences
                    - Present and past tense
                    - Natural transitions
                    Preserve the original story's elements while enhancing the flow.
                """,
                'max_tokens': 700,
                'temperature': 0.7
            },
            'Advanced': {
                'instructions': """
                    Create a sophisticated narrative using:
                    - Rich vocabulary
                    - Varied sentence structures
                    - Multiple tenses
                    - Elegant transitions
                    Maintain the story's essence while adding literary flourish.
                """,
                'max_tokens': 1000,
                'temperature': 0.8
            }
        }

        level_settings = level_context[level]
        
        # Generate polished English version
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": f"""
                    You are a professional children's story editor.
                    {level_settings['instructions']}
                    Context: {theme_context}
                    
                    Task: Create a polished, coherent version of this story while:
                    1. Maintaining the original plot points
                    2. Improving flow and transitions
                    3. Adding appropriate descriptive elements
                    4. Making the narrative more engaging
                    5. Keeping the language level appropriate for {level} students
                    """
                },
                {
                    "role": "user",
                    "content": f"Original story:\n{story_summary}\n\nCreate a polished version:"
                }
            ],
            max_tokens=level_settings['max_tokens'],
            temperature=level_settings['temperature']
        )
        
        polished_english = response.choices[0].message.content.strip()
        
        # Generate Thai translation
        translation_response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": """
                    You are a professional Thai-English translator specializing in children's literature.
                    Create a natural, flowing Thai translation that:
                    1. Captures the story's meaning and emotion
                    2. Uses appropriate Thai language for the target age group
                    3. Maintains cultural relevance
                    4. Reads naturally in Thai
                    """
                },
                {
                    "role": "user",
                    "content": f"Translate this story to Thai:\n{polished_english}"
                }
            ],
            max_tokens=1000,
            temperature=0.7
        )
        
        thai_translation = translation_response.choices[0].message.content.strip()
        
        return {
            'original': story_summary,
            'polished_english': polished_english,
            'thai_translation': thai_translation
        }
        
    except Exception as e:
        logging.error(f"Error in story stitching: {str(e)}")
        raise

def generate_story_illustration(story_summary: str) -> str:
    """Generate illustration for the story using DALL-E 3"""
    try:
        # Refined prompt for DALL-E
        illustration_prompt = f"""
        Create a single, cohesive children's book cover illustration for this story:
        {story_summary}
        
        Style requirements:
        - Cute and child-friendly illustration
        - Soft, warm colors
        - Storybook art style
        - Clear and simple composition
        - Safe for children
        - Focus on a single, unified scene
        - No text or sub-scenes in the image
        """
        
        # Generate image
        response = client.images.generate(
            model="dall-e-3",
            prompt=illustration_prompt,
            size="1024x1024",
            quality="standard",
            n=1,  # Ensures only one image is generated
        )
        
        image_url = response.data[0].url
        logging.info("Story illustration generated successfully")
        return image_url
        
    except Exception as e:
        logging.error(f"Error generating story illustration: {str(e)}")
        return None


def show_stitched_story(story_data: Dict[str, str]):
    """Display the stitched story with translation and illustration"""
    try:
        st.markdown("# 📖 Your Polished Story | เรื่องราวฉบับสมบูรณ์")
        
        # English Version
        st.markdown("### 🇬🇧 English Version")
        st.markdown(
            f"""
            <div style="background-color: #f8f9fa; padding: 20px; border-radius: 10px; 
                        border-left: 4px solid #1565c0; margin-bottom: 20px;">
                {story_data['polished_english']}
            </div>
            """,
            unsafe_allow_html=True
        )
        
        # Generate and display illustration
        with st.spinner("🎨 กำลังวาดภาพประกอบ..."):
            illustration_url = generate_story_illustration(story_data['polished_english'])
            if illustration_url:
                st.markdown("### 🎨 Story Illustration | ภาพประกอบเรื่องราว")
                st.image(
                    illustration_url,
                    caption="AI-generated illustration for your story",
                    use_column_width=True
                )
                
                # Add download button for illustration
                st.markdown(f"""
                    <div style="text-align: center; margin: 10px 0;">
                        <a href="{illustration_url}" download="story_illustration.png" 
                           target="_blank" style="text-decoration: none;">
                            📥 ดาวน์โหลดภาพประกอบ
                        </a>
                    </div>
                """, unsafe_allow_html=True)
        
        # Thai Version
        st.markdown("### 🇹🇭 ฉบับภาษาไทย")
        st.markdown(
            f"""
            <div style="background-color: #f8f9fa; padding: 20px; border-radius: 10px; 
                        border-left: 4px solid #28a745; margin-bottom: 20px;">
                {story_data['thai_translation']}
            </div>
            """,
            unsafe_allow_html=True
        )

        # Save Options with illustration
        st.markdown("### 💾 บันทึกเรื่องราวฉบับสมบูรณ์")
        save_col1, save_col2 = st.columns(2)
        
        with save_col1:
            if st.button("📥 บันทึกเป็น PDF", key="save_stitched_pdf", use_container_width=True):
                # Add illustration to PDF if available
                pdf_data = create_bilingual_story_pdf(story_data, illustration_url)
                st.download_button(
                    "ดาวน์โหลด PDF",
                    data=pdf_data,
                    file_name=f"my_story_{datetime.now().strftime('%Y%m%d_%H%M')}.pdf",
                    mime="application/pdf",
                    key="download_stitched_pdf"
                )
        
        with save_col2:
            if st.button("💾 บันทึกข้อความ", key="save_stitched_text", use_container_width=True):
                text_data = {
                    **story_data,
                    'illustration_url': illustration_url
                }
                json_str = json.dumps(text_data, ensure_ascii=False, indent=2)
                st.download_button(
                    "ดาวน์โหลดข้อความ",
                    data=json_str,
                    file_name=f"my_story_{datetime.now().strftime('%Y%m%d_%H%M')}.json",
                    mime="application/json",
                    key="download_stitched_text"
                )
                
    except Exception as e:
        logging.error(f"Error showing stitched story: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการแสดงเรื่องราว กรุณาลองใหม่อีกครั้ง")

def create_bilingual_story_pdf(story_data: Dict[str, str], illustration_url: str = None) -> bytes:
    """Create a PDF with both English and Thai versions of the story and illustration"""
    try:
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(
            buffer,
            pagesize=A4,
            rightMargin=72,
            leftMargin=72,
            topMargin=72,
            bottomMargin=72
        )

        # Create styles
        styles = getSampleStyleSheet()
        title_style = ParagraphStyle(
            'CustomTitle',
            parent=styles['Title'],
            fontSize=24,
            spaceAfter=30,
            alignment=1  # Center alignment
        )
        heading_style = ParagraphStyle(
            'CustomHeading',
            parent=styles['Heading1'],
            fontSize=18,
            spaceAfter=12,
            textColor=colors.blue
        )
        body_style = ParagraphStyle(
            'CustomBody',
            parent=styles['Normal'],
            fontSize=12,
            spaceBefore=6,
            spaceAfter=6,
            leading=16
        )
        
        # Create story elements
        elements = []
        
        # Add title
        elements.append(Paragraph("My Story - JoyStory", title_style))
        elements.append(Spacer(1, 20))
        
        # Add illustration if available
        if illustration_url:
            try:
                # Download and add illustration
                response = requests.get(illustration_url)
                img_data = io.BytesIO(response.content)
                img = Image(img_data, width=400, height=400)
                elements.append(img)
                elements.append(Spacer(1, 20))
            except Exception as e:
                logging.error(f"Error adding illustration to PDF: {str(e)}")
        
        # Add English version
        elements.append(Paragraph("English Version", heading_style))
        elements.append(Paragraph(story_data['polished_english'], body_style))
        elements.append(Spacer(1, 20))
        
        # Add Thai version
        elements.append(Paragraph("ฉบับภาษาไทย", heading_style))
        elements.append(Paragraph(story_data['thai_translation'], body_style))
        
        # Build PDF
        doc.build(elements)
        pdf = buffer.getvalue()
        buffer.close()
        
        return pdf
        
    except Exception as e:
        logging.error(f"Error creating bilingual PDF: {str(e)}")
        raise

def show_story_stats():
    """Display comprehensive story statistics"""
    try:
        st.markdown("""
            <div style="
                background-color: #f3e5f5;
                padding: 20px;
                border-radius: 10px;
                margin: 20px 0;
            ">
                <h3 style="color: #6a1b9a; margin-bottom: 15px; text-align: center;">
                    📊 สถิติการเขียนเรื่องราว
                </h3>
            """, unsafe_allow_html=True)

        # Basic statistics in 2 columns
        col1, col2 = st.columns(2)
        
        with col1:
            st.metric(
                "จำนวนประโยคทั้งหมด",
                len(st.session_state.story),
                help="จำนวนประโยคในเรื่องทั้งหมด"
            )
            st.metric(
                "คำศัพท์ที่ใช้",
                len(st.session_state.stats['vocabulary_used']),
                help="จำนวนคำศัพท์ที่ไม่ซ้ำกัน"
            )
            st.metric(
                "ประโยคที่ถูกต้องแล้ว",
                st.session_state.stats['correct_first_try'],
                help="จำนวนประโยคที่เขียนถูกต้องตั้งแต่ครั้งแรก"
            )
        
        with col2:
            st.metric(
                "คะแนนรวม",
                st.session_state.points['total'],
                help="คะแนนรวมที่ได้รับ"
            )
            st.metric(
                "ความแม่นยำ",
                f"{st.session_state.stats['accuracy_rate']:.1f}%",
                help="อัตราการเขียนถูกต้อง"
            )
            st.metric(
                "Streak สูงสุด",
                st.session_state.points['max_streak'],
                help="จำนวนประโยคถูกต้องติดต่อกันมากที่สุด"
            )

        # Show achievements
        if st.session_state.achievements:
            st.markdown("""
                <h4 style="color: #6a1b9a; margin: 20px 0 10px 0;">
                    🏆 ความสำเร็จที่ได้รับ
                </h4>
            """, unsafe_allow_html=True)
            
            for achievement in st.session_state.achievements:
                st.success(achievement)
        
        # Show writing progress details
        st.markdown("""
            <h4 style="color: #6a1b9a; margin: 20px 0 10px 0;">
                📝 รายละเอียดการเขียน
            </h4>
        """, unsafe_allow_html=True)
        
        details_col1, details_col2 = st.columns(2)
        
        with details_col1:
            # Story composition metrics
            st.markdown("##### 📖 องค์ประกอบเรื่อง")
            avg_sentence_length = st.session_state.stats.get('average_sentence_length', 0)
            st.markdown(f"""
                - ความยาวประโยคเฉลี่ย: {avg_sentence_length:.1f} คำ
                - จำนวนคำทั้งหมด: {st.session_state.stats.get('total_words', 0)} คำ
                - การแก้ไข: {st.session_state.stats.get('corrections_made', 0)} ครั้ง
            """)
        
        with details_col2:
            # Time and progress metrics
            st.markdown("##### ⏱️ เวลาและความก้าวหน้า")
            session_duration = st.session_state.stats.get('session_duration', 0)
            duration_minutes = session_duration / 60
            st.markdown(f"""
                - เวลาที่ใช้: {duration_minutes:.1f} นาที
                - ความก้าวหน้า: {min(len(st.session_state.story) * 5, 100)}%
                - สถานะ: {"จบเรื่องแล้ว" if st.session_state.story_completed else "กำลังเขียน"}
            """)
        
        # Show vocabulary usage
        if st.session_state.stats['vocabulary_used']:
            with st.expander("📚 คำศัพท์ที่ใช้"):
                vocab_list = sorted(list(st.session_state.stats['vocabulary_used']))
                st.markdown(f"""
                    <div style="
                        background-color: white;
                        padding: 10px;
                        border-radius: 5px;
                        max-height: 200px;
                        overflow-y: auto;
                    ">
                        {', '.join(vocab_list)}
                    </div>
                """, unsafe_allow_html=True)
        
        st.markdown("</div>", unsafe_allow_html=True)
        
    except Exception as e:
        logging.error(f"Error showing story stats: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการแสดงสถิติ")

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("### 🎭 ต้องการจบเรื่องหรือไม่?")
        
        # Display ending type options
        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()  # Using st.rerun() instead of st.experimental_rerun()

def handle_ending_mode(text: str):
    """Handle story submission during ending mode"""
    try:
        # Get current remaining sentences
        remaining = st.session_state.sentences_to_end

        # 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)
        
        # 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"""
    try:
        # เล่นเสียงจบเรื่อง
        st.markdown(f"<script>{st.session_state.sound_commands['complete']}</script>", unsafe_allow_html=True)
        st.balloons()

        # Generate 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()

    except Exception as e:
        logging.error(f"Error in complete_story: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการแสดงผลการจบเรื่อง")

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="จำนวนประโยคทั้งหมดที่คุณได้เขียน"
            )
            st.metric(
                "ถูกต้องตั้งแต่แรก", 
                st.session_state.stats['correct_first_try'],
                help="จำนวนประโยคที่ถูกต้องโดยไม่ต้องแก้ไข"
            )
        with col2:
            st.metric(
                "ความแม่นยำ", 
                f"{st.session_state.stats['accuracy_rate']:.1f}%",
                help="เปอร์เซ็นต์ของประโยคที่ถูกต้องตั้งแต่แรก"
            )
            st.metric(
                "คำศัพท์ที่ใช้", 
                len(st.session_state.stats['vocabulary_used']),
                help="จำนวนคำศัพท์ที่ไม่ซ้ำกันที่คุณได้ใช้"
            )

        # Show Achievements
        st.markdown("""
            <div style="margin: 25px 0 15px 0;">
                <h3>🏆 ความสำเร็จ</h3>
            </div>
        """, unsafe_allow_html=True)
        
        if st.session_state.achievements:
            for achievement in st.session_state.achievements:
                st.success(achievement)
        else:
            st.info("ยังไม่มีความสำเร็จ - เขียนต่อไปเพื่อปลดล็อกรางวัล!")
            
            # Show available achievements
            st.markdown("""
                <div style="
                    margin-top: 15px;
                    padding: 15px;
                    background-color: #f5f5f5;
                    border-radius: 8px;
                ">
                    <p style="color: #666; margin-bottom: 10px;">
                        รางวัลที่รอคุณอยู่:
                    </p>
                    <ul style="list-style-type: none; padding-left: 0;">
                        <li>🌟 นักเขียนไร้ที่ติ: เขียนถูกต้อง 5 ประโยคติดต่อกัน</li>
                        <li>📚 ราชาคำศัพท์: ใช้คำศัพท์ไม่ซ้ำกัน 50 คำ</li>
                        <li>📖 นักแต่งนิทาน: เขียนเรื่องยาว 10 ประโยค</li>
                        <li>👑 ราชาความแม่นยำ: มีอัตราความถูกต้อง 80% ขึ้นไป</li>
                    </ul>
                </div>
            """, unsafe_allow_html=True)

def show_save_options():
    """Display save and export options"""
    st.markdown("### 💾 บันทึกเรื่องราว")
    
    if not st.session_state.story:
        st.info("เริ่มเขียนเรื่องราวก่อนเพื่อบันทึกผลงานของคุณ")
        return
        
    col1, col2 = st.columns(2)
    with col1:
        # PDF Export
        if st.button("📑 บันทึกเป็น PDF", 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"
                )
            except Exception as e:
                logging.error(f"Error creating PDF: {str(e)}")
                st.error("เกิดข้อผิดพลาดในการสร้างไฟล์ PDF กรุณาลองใหม่อีกครั้ง")
    
    with col2:
        # JSON Save
        if st.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': convert_sets_to_lists(st.session_state.stats),
                    '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"
                )
            except Exception as e:
                logging.error(f"Error saving progress: {str(e)}")
                st.error("เกิดข้อผิดพลาดในการบันทึกความก้าวหน้า กรุณาลองใหม่อีกครั้ง")

def show_sidebar():
    """Display sidebar content"""
    st.sidebar.markdown("### ⚙️ การตั้งค่า")
    
    # Progress Upload
    st.sidebar.markdown("#### 📂 โหลดความก้าวหน้า")
    uploaded_file = st.sidebar.file_uploader(
        "เลือกไฟล์ .json",
        type=['json'],
        help="เลือกไฟล์ความก้าวหน้าที่บันทึกไว้"
    )
    if uploaded_file:
        if st.sidebar.button("โหลดความก้าวหน้า", use_container_width=True):
            try:
                data = json.loads(uploaded_file.getvalue())
                load_progress(data)
                st.sidebar.success("โหลดความก้าวหน้าเรียบร้อย!")
                st.rerun()
            except Exception as e:
                logging.error(f"Error loading progress: {str(e)}")
                st.sidebar.error("เกิดข้อผิดพลาดในการโหลดไฟล์")
    
    # Level Selection
    st.sidebar.markdown("#### 🎯 ระดับการเรียน")
    level = show_level_selection()
    st.session_state.level = level
    
    # Theme Change Option
    if st.session_state.current_theme:
        st.sidebar.markdown("#### 🎨 ธีมเรื่องราว")
        if st.sidebar.button("🔄 เปลี่ยนธีม", use_container_width=True):
            st.session_state.current_theme = None
            st.session_state.theme_story_starter = None
            st.rerun()
    
    # Reset Option
    st.sidebar.markdown("#### 🔄 รีเซ็ตเรื่องราว")
    if st.sidebar.button("เริ่มเรื่องใหม่", use_container_width=True):
        if st.session_state.story:
            if st.sidebar.checkbox("ยืนยันการเริ่มใหม่"):
                st.session_state.should_reset = True
                st.rerun()
        else:
            st.sidebar.info("ยังไม่มีเรื่องราวที่จะรีเซ็ต")

def clear_input_state():
    """Clear all input-related session states"""
    if 'text_input' in st.session_state:
        st.session_state.text_input = ""
    if 'story_input_area' in st.session_state:
        st.session_state.story_input_area = ""
    if 'last_submission' in st.session_state:
        del st.session_state.last_submission
    st.session_state.clear_input = True
    # Increment clear count to force new input field
    st.session_state.clear_count = st.session_state.get('clear_count', 0) + 1

def show_story_input():
    """Display story input section with improved clear functionality"""
    st.markdown("""
        <div class="thai-eng">
            <div class="thai">✏️ ถึงตาคุณแล้ว</div>
            <div class="eng">Your Turn</div>
        </div>
    """, unsafe_allow_html=True)
    
    # Initialize states if not exists
    if 'clear_input' not in st.session_state:
        st.session_state.clear_input = False
    if 'text_input' not in st.session_state:
        st.session_state.text_input = ""
    if 'last_submission' not in st.session_state:
        st.session_state.last_submission = None

    # 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 key based on clear state
    text_input = st.text_area(
        "เขียนต่อจากเรื่องราว | Continue the story:",
        value=st.session_state.text_input,
        height=100,
        key=f"story_input_area_{st.session_state.get('clear_count', 0)}",
        help="พิมพ์ประโยคภาษาอังกฤษเพื่อต่อเรื่อง",
        label_visibility="collapsed"
    )
    
    # Keep track of current input
    st.session_state.text_input = text_input
    
    # Submit and Clear buttons
    col1, col2, col3 = st.columns([3, 1, 1])

    with col1:
        if st.button("📝 ส่งคำตอบ | Submit", use_container_width=True):
            if not text_input.strip():
                st.warning("กรุณาเขียนข้อความก่อนส่ง")
                return
            try:
                # เล่นเสียงตอนกดปุ่ม Submit
                st.markdown(f"<script>{st.session_state.sound_commands['submit']}</script>", unsafe_allow_html=True)
                with st.spinner("กำลังวิเคราะห์ประโยค..."):
                    handle_story_submission(text_input.strip())
            except Exception as e:
                logging.error(f"Error submitting story: {str(e)}")
                st.error("เกิดข้อผิดพลาดในการส่งคำตอบ กรุณาลองใหม่อีกครั้ง")
    
    with col2:
        # Clear button with unique key
        if st.button(
            "🗑️ ล้างข้อความ | Clear",
            key=f"clear_button_{st.session_state.get('clear_count', 0)}",
            use_container_width=True,
            disabled=not text_input.strip()  # Disable if already empty
        ):
            # Increment clear count to force new key generation
            st.session_state.clear_count = st.session_state.get('clear_count', 0) + 1
            clear_input_state()
            st.rerun()
    
    with col3:
        # Character count
        char_count = len(text_input)
        st.markdown(f"""
            <div style="text-align: right; color: {'red' if char_count > 200 else '#666'}; padding: 8px;">
                {char_count}/200 ตัวอักษร
            </div>
        """, unsafe_allow_html=True)

    # Warning for long text
    if char_count > 200:
        st.warning("⚠️ ข้อความยาวเกิน 200 ตัวอักษร")

def handle_story_submission(text: str):
    """Handle story submission with improved state management"""
    if not st.session_state.story:
        st.error("กรุณาเลือกธีมเรื่องราวก่อนเริ่มเขียน")
        return

    # Check if in ending mode
    if st.session_state.get('ending_mode'):
        handle_ending_mode(text)
        # Clear input after handling ending mode
        clear_input_state()
        return

    # Regular processing
    try:
        # 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)
        
        # Add user's sentence
        st.session_state.story.append({
            "role": "You",
            "content": text,
            "is_corrected": False,
            "is_correct": is_correct,
            "timestamp": datetime.now().isoformat()
        })
        
        # 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()
        
        # Generate AI continuation
        try:
            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:
                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,
                    "timestamp": datetime.now().isoformat()
                })
        except Exception as e:
            logging.error(f"Error generating AI continuation: {str(e)}")
            # 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()
            })
        
        # Update session stats
        update_session_stats()
        
        # Clear input state completely before rerun
        clear_input_state()
        
        # Make sure we don't have any lingering text
        if 'story_input_area' in st.session_state:
            st.session_state.story_input_area = ""
        
        # Increment clear count to force new input field
        st.session_state.clear_count = st.session_state.get('clear_count', 0) + 1
        
        # Rerun to update UI
        st.rerun()
        
    except Exception as e:
        logging.error(f"Error in story submission: {str(e)}")
        # Clear input even on error to prevent stuck state
        clear_input_state()
        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_main_interface():
    """Display main story interface"""
    col1, col2 = st.columns([3, 1])
    
    with col1:
        # Story display
        st.markdown("""
            <div class="thai-eng">
                <div class="thai">📖 เรื่องราวของคุณ</div>
                <div class="eng">Your Story</div>
            </div>
        """, unsafe_allow_html=True)
        
        show_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
        show_feedback_section()
        
        # Writing tools
        show_writing_tools()
        
        # Achievements
        with st.expander("🏆 ความสำเร็จ | Achievements"):
            show_achievements()
        
        # 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 with proper ending type display
        ending_type_display = {
            "Happy Ending": "จบแบบมีความสุข",
            "Mysterious Ending": "จบแบบทิ้งท้ายให้คิดต่อ",
            "Lesson Learned": "จบแบบได้ข้อคิด",
            "Surprise Ending": "จบแบบพลิกความคาดหมาย"
        }
        
        current_ending = ending_type_display.get(
            st.session_state.ending_type,
            st.session_state.ending_type
        )
        
        st.markdown(f"""
            <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;">
                    เรื่องราวของคุณจบลงด้วย: {current_ending}
                </p>
            </div>
        """, unsafe_allow_html=True)
        
        # Show statistics
        show_story_stats()
        
        # Offer story stitching option
        st.markdown("### ✨ ต้องการให้ AI ช่วยเรียบเรียงเรื่องราวให้สมบูรณ์ขึ้นไหม?")
        
        if st.button("🎨 เรียบเรียงเรื่องราว", 
                    key="stitch_story_button", 
                    use_container_width=True):
            with st.spinner("กำลังเรียบเรียงเรื่องราว..."):
                try:
                    stitched_story = generate_stitched_story(
                        story=st.session_state.story,
                        style="Classic Fairytale",  # Default style
                        detail_level="ปานกลาง",    # Default detail level
                        theme=st.session_state.current_theme,
                        level=st.session_state.level
                    )
                    st.session_state.stitched_story = stitched_story
                    show_stitched_story(stitched_story)
                except Exception as e:
                    logging.error(f"Error in story stitching: {str(e)}")
                    st.error("เกิดข้อผิดพลาดในการเรียบเรียงเรื่องราว กรุณาลองใหม่อีกครั้ง")
        
        # Original save options
        st.markdown("### 💾 บันทึกเรื่องราวต้นฉบับ")
        save_col1, save_col2 = st.columns(2)
        
        with save_col1:
            if st.button("📑 บันทึกเป็น PDF", 
                        key="save_original_pdf", 
                        use_container_width=True):
                pdf_data = create_story_pdf()
                st.download_button(
                    "ดาวน์โหลด PDF",
                    data=pdf_data,
                    file_name=f"original_story_{datetime.now().strftime('%Y%m%d_%H%M')}.pdf",
                    mime="application/pdf",
                    key="download_original_pdf"
                )
        
        with save_col2:
            if st.button("🔄 เริ่มเรื่องใหม่", 
                        key="new_story_button", 
                        use_container_width=True):
                if st.checkbox("✅ ยืนยันการเริ่มใหม่", 
                             key="confirm_new_story"):
                    reset_story()
                    st.rerun()

    except Exception as e:
        logging.error(f"Error showing completion options: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการแสดงตัวเลือกหลังจบเรื่อง")

def show_story_stitching_options():
    """Display story stitching options and handle the flow"""
    try:
        st.markdown("""
            <div style="
                background-color: #e3f2fd;
                padding: 20px;
                border-radius: 10px;
                margin: 20px 0;
            ">
                <h3 style="color: #1565c0; margin-bottom: 15px;">
                    ✨ เรียบเรียงเรื่องราวให้สมบูรณ์
                </h3>
                <p style="color: #333;">
                    AI จะช่วย:
                    - ปรับแต่งการเชื่อมประโยคให้ลื่นไหล
                    - เพิ่มรายละเอียดที่น่าสนใจ
                    - แปลเป็นภาษาไทย
                    - รักษาเนื้อเรื่องและความคิดสร้างสรรค์ของคุณไว้
                </p>
            </div>
        """, unsafe_allow_html=True)
        
        col1, col2 = st.columns(2)
        
        with col1:
            style_option = st.selectbox(
                "เลือกรูปแบบการเล่าเรื่อง:",
                options=[
                    "Classic Fairytale - นิทานแบบคลาสสิก",
                    "Modern Adventure - การผจญภัยสมัยใหม่",
                    "Poetic Style - แบบกวีนิพนธ์",
                    "Simple and Clear - เรียบง่ายและชัดเจน"
                ],
                key="story_style_selector"
            )
        
        with col2:
            detail_level = st.select_slider(
                "ระดับรายละเอียด:",
                options=["น้อย", "ปานกลาง", "มาก"],
                value="ปานกลาง",
                key="detail_level_selector"
            )
        
        if st.button("🎨 เริ่มเรียบเรียงเรื่องราว", 
                    key="start_stitching", 
                    use_container_width=True):
            with st.spinner("กำลังเรียบเรียงเรื่องราว..."):
                try:
                    stitched_story = generate_stitched_story(
                        story=st.session_state.story,
                        style=style_option,
                        detail_level=detail_level,
                        theme=st.session_state.current_theme,
                        level=st.session_state.level
                    )
                    
                    # Save stitched story to session state
                    st.session_state.stitched_story = stitched_story
                    
                    # Show the result
                    show_stitched_result(stitched_story)
                    
                except Exception as e:
                    logging.error(f"Error in story stitching: {str(e)}")
                    st.error("เกิดข้อผิดพลาดในการเรียบเรียงเรื่องราว กรุณาลองใหม่อีกครั้ง")
    
    except Exception as e:
        logging.error(f"Error showing stitching options: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการแสดงตัวเลือกการเรียบเรียง")

def generate_stitched_story(story: List[dict], style: str, detail_level: str, theme: str, level: str) -> Dict[str, str]:
    """Generate a polished version of the story with minimal enhancements"""
    try:
        # แยกประโยคต่างๆ ตาม role
        story_parts = {
            'starter': next((entry['content'] for entry in story if entry.get('is_starter')), ''),
            'user_sentences': [entry['content'] for entry in story if entry['role'] == 'You'],
            'ai_responses': [entry['content'] for entry in story if entry['role'] == 'AI' and not entry.get('is_starter')],
            'ending': next((entry['content'] for entry in story if entry.get('is_final')), '')
        }

        # สร้าง prompt ที่เน้นการรักษาประโยคเดิมและเชื่อมต่อแบบกระชับ
        prompt_template = f"""
        Rewrite this story into a flowing narrative.

        CRITICAL RULES:
        1. Keep original sentences almost exactly as they are
        2. Add only minimal connecting words or phrases
        3. Do not add new scenes or events
        4. Do not add lengthy descriptions
        5. Focus on making the story flow naturally
        6. Keep any additions very brief and simple

        Examples of good enhancement:
        Original: "The project started. We made robots."
        Good: "The project started with great excitement. We made robots, our first real creation."
        BAD: "In the sunny classroom, with students gathered around shiny tables, the project started. We made incredible robots with flashing lights and complex circuits."

        Original Story Parts:
        - Beginning: {story_parts['starter']}
        - Main Story: {' '.join(sum(zip(story_parts['user_sentences'], story_parts['ai_responses']), ()))}
        - Ending: {story_parts['ending']}
        
        Create a coherent story that stays very close to the original while making it flow smoothly.
        Remember: Less is more - add only what's necessary for flow.
        """

        # Generate English version with minimal enhancements
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": """You are a story editor who specializes in minimal enhancement.
                    Your goal is to make the story flow while keeping it as close as possible to the original.
                    Remember: Only add what's absolutely necessary for coherence."""
                },
                {
                    "role": "user",
                    "content": prompt_template
                }
            ],
            temperature=0.5  # Lower temperature for more consistent, conservative output
        )
        
        polished_english = response.choices[0].message.content.strip()

        # Generate Thai translation with same minimal approach
        translation_prompt = f"""
        Translate this story to Thai. 

        IMPORTANT:
        1. Keep the translation concise and close to the English version
        2. Do not add extra details or explanations
        3. Maintain the same level of simplicity
        4. Focus on natural Thai flow while keeping original content
        
        Story to translate:
        {polished_english}
        """

        translation_response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": """You are a Thai translator who specializes in concise, accurate translations.
                    Keep the same level of detail as the English version without adding extra elements."""
                },
                {
                    "role": "user",
                    "content": translation_prompt
                }
            ],
            temperature=0.5
        )
        
        thai_translation = translation_response.choices[0].message.content.strip()

        return {
            'polished_english': polished_english,
            'thai_translation': thai_translation,
            'style': style,
            'level': level
        }

    except Exception as e:
        logging.error(f"Error generating stitched story: {str(e)}")
        raise

def show_stitched_result(story_data: Dict[str, str]):
    """Display the stitched story result with enhanced formatting"""
    try:
        # Check if story data exists
        if not story_data or not isinstance(story_data, dict):
            st.error("ไม่พบข้อมูลเรื่องราว กรุณาลองใหม่อีกครั้ง")
            return

        # Header
        st.markdown("### 📖 Your Polished Story | เรื่องราวฉบับสมบูรณ์")
        
        # English Version
        st.subheader("🇬🇧 English Version")
        st.markdown(
            f'<div style="background-color: white; padding: 15px; border-radius: 8px; margin-bottom: 15px;">'
            f'{story_data.get("polished_english", "No content available")}'
            f'</div>',
            unsafe_allow_html=True
        )
        
        # Thai Version
        st.subheader("🇹🇭 ฉบับภาษาไทย")
        st.markdown(
            f'<div style="background-color: white; padding: 15px; border-radius: 8px;">'
            f'{story_data.get("thai_translation", "ไม่พบเนื้อหา")}'
            f'</div>',
            unsafe_allow_html=True
        )

    except Exception as e:
        logging.error(f"Error showing stitched result: {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

def check_audio_system():
    """Check if audio system is working properly"""
    try:
        if 'audio_manager' in st.session_state:
            # ตรวจสอบว่า audio elements ถูกโหลดสำเร็จ
            js_check = """
                <script>
                function checkAudioElements() {
                    const elements = ['bgm', 'sound_submit', 'sound_achievement', 'sound_complete'];
                    for (let id of elements) {
                        if (!document.getElementById(id)) {
                            console.error(`Audio element ${id} not found`);
                            return false;
                        }
                    }
                    return true;
                }
                window.audioSystemWorking = checkAudioElements();
                </script>
            """
            st.markdown(js_check, unsafe_allow_html=True)
            return True
    except Exception as e:
        logging.error(f"Error checking audio system: {str(e)}")
        return False

# === 6. MAIN APPLICATION LOGIC ===
def main():
    try:
        # Initialize states
        init_session_state()
        init_theme_state()

        # Check if reset is needed
        if st.session_state.get('should_reset'):
            reset_story()
            st.rerun()
            return

        # Initialize audio system
        initialize_audio()

        # Check audio system
        if not check_audio_system():
            st.warning("ระบบเสียงอาจทำงานไม่สมบูรณ์ แต่คุณยังสามารถใช้งานแอพได้")
        
        # Initialize ending system state if not exists
        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  # 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("""
            <div style='position: fixed; bottom: 10px; right: 10px; z-index: 1000; 
                        opacity: 0.7; font-size: 0.8em; color: #666;'>
                Powered by JoyStory AI
            </div>
        """, unsafe_allow_html=True)

        # Show header
        st.markdown("# 📖 JoyStory")
        show_welcome_section()
        show_parent_guide()

        # Sidebar
        with st.sidebar:
            show_sidebar()
            show_audio_controls()
            
        # Session Status Check
        check_session_status()

        # Main content area
        main_container = st.container()
        with main_container:
            if not st.session_state.current_theme:
                show_theme_selection()
            else:
                # Show story progress if story exists
                if st.session_state.story:
                    # 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 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 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()
                
                # 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.get('should_reset'):
            reset_story()
            # Reset ending states
            st.session_state.ending_mode = False
            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()

    except Exception as e:
        handle_application_error(e)


def check_session_status():
    """Check and maintain session status"""
    try:
        # Update session duration
        if 'session_start' not in st.session_state:
            st.session_state.session_start = datetime.now()
        
        # 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("เซสชันหมดอายุ กรุณาบันทึกความก้าวหน้าและรีเฟรชหน้าเว็บ")
        
        # Check for inactivity (30 minutes)
        last_interaction = datetime.fromisoformat(st.session_state.last_interaction)
        inactivity_duration = (datetime.now() - last_interaction).total_seconds()
        if inactivity_duration > 1800:  # 30 minutes
            st.info("ไม่มีกิจกรรมเป็นเวลานาน กรุณาบันทึกความก้าวหน้าเพื่อความปลอดภัย")
        
        # Update stats if story exists
        if st.session_state.story:
            update_session_stats()
            
    except Exception as e:
        logging.error(f"Error checking session status: {str(e)}")

def auto_save_progress():
    """Automatically save progress to session state"""
    try:
        current_progress = {
            'timestamp': datetime.now().isoformat(),
            'story': st.session_state.story,
            'stats': st.session_state.stats,
            'points': st.session_state.points,
            'achievements': st.session_state.achievements
        }
        
        # Save to session state
        if 'auto_save' not in st.session_state:
            st.session_state.auto_save = {}
        
        st.session_state.auto_save = current_progress
        
        # Add timestamp for last auto-save
        st.session_state.last_auto_save = datetime.now().isoformat()
        
    except Exception as e:
        logging.error(f"Error in auto-save: {str(e)}")

def handle_application_error(error: Exception):
    """Handle application-wide errors"""
    logging.error(f"Application error: {str(error)}")
    
    error_message = """
    <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>
    """
    
    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 show_debug_info():
    """Show debug information (development only)"""
    if st.session_state.get('debug_mode'):
        with st.expander("🔧 Debug Information"):
            st.json({
                'session_state': {
                    key: str(value) if isinstance(value, (set, datetime)) else value
                    for key, value in st.session_state.items()
                    if key not in ['client', '_client']
                }
            })

# Add CSS for loading animation
st.markdown("""
    <style>
    @keyframes pulse {
        0% { opacity: 0.6; }
        50% { opacity: 1; }
        100% { opacity: 0.6; }
    }
    
    .loading-animation {
        animation: pulse 1.5s infinite;
        background-color: #f0f2f5;
        border-radius: 8px;
        padding: 20px;
        text-align: center;
        margin: 20px 0;
    }
    
    /* Custom Scrollbar */
    ::-webkit-scrollbar {
        width: 8px;
        height: 8px;
    }
    
    ::-webkit-scrollbar-track {
        background: #f1f1f1;
        border-radius: 4px;
    }
    
    ::-webkit-scrollbar-thumb {
        background: #888;
        border-radius: 4px;
    }
    
    ::-webkit-scrollbar-thumb:hover {
        background: #666;
    }
    
    /* Toast Notifications */
    .toast-notification {
        position: fixed;
        bottom: 20px;
        right: 20px;
        padding: 15px 25px;
        background-color: #333;
        color: white;
        border-radius: 8px;
        z-index: 1000;
        animation: slideIn 0.3s ease-out;
    }
    
    @keyframes slideIn {
        from { transform: translateX(100%); }
        to { transform: translateX(0); }
    }
    
    /* Progress Indicator */
    .progress-bar {
        width: 100%;
        height: 4px;
        background-color: #e0e0e0;
        border-radius: 2px;
        overflow: hidden;
    }
    
    .progress-bar-fill {
        height: 100%;
        background-color: #1e88e5;
        transition: width 0.3s ease;
    }
    </style>
""", unsafe_allow_html=True)

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        handle_application_error(e)