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

# === 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>
    /* 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"""
    try:
        # Reset story-related states
        st.session_state.story = []
        st.session_state.feedback = None
        st.session_state.theme_story_starter = None
        
        # Reset points
        st.session_state.points = {
            'total': 0,
            'perfect_sentences': 0,
            'corrections_made': 0,
            'streak': 0,
            'max_streak': 0
        }
        
        # Reset stats
        st.session_state.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 progress
        st.session_state.current_milestone = 0
        st.session_state.next_milestone = 5
        
        # Reset flag
        st.session_state.should_reset = False
        
        # 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 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
        story_context = '\n'.join([
            entry['content'] for entry in st.session_state.story[-3:]
        ]) 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("🎉 ได้รับความสำเร็จใหม่: นักเขียนไร้ที่ติ!")
        
        if (len(st.session_state.stats['vocabulary_used']) >= 50 and 
            "📚 ราชาคำศัพท์" not in current_achievements):
            current_achievements.append("📚 ราชาคำศัพท์")
            st.success("🎉 ได้รับความสำเร็จใหม่: ราชาคำศัพท์!")
        
        if (len(st.session_state.story) >= 10 and 
            "📖 นักแต่งนิทาน" not in current_achievements):
            current_achievements.append("📖 นักแต่งนิทาน")
            st.success("🎉 ได้รับความสำเร็จใหม่: นักแต่งนิทาน!")
        
        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.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 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.
        """,
        'Intermediate': """
        Focus on:
        - Sentence variety
        - Past Tense usage
        - Vocabulary appropriateness
        - Basic punctuation
        - Simple conjunctions
        Provide moderately detailed feedback in Thai.
        """,
        'Advanced': """
        Focus on:
        - Complex sentence structures
        - Various tense usage
        - Advanced vocabulary
        - All punctuation
        - Style and flow
        Provide comprehensive feedback in Thai.
        """
    }

    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=200,
                    temperature=0.3
                )
                
                # 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("กำลังเตรียมเรื่องราว..."):
            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:
        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 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

        # Add user's sentence
        st.session_state.story.append({
            "role": "You",
            "content": text,
            "is_corrected": False,
            "is_correct": True,
            "timestamp": datetime.now().isoformat(),
            "remaining_sentences": remaining
        })
        
        # Decrease remaining count for user's sentence
        st.session_state.sentences_to_end -= 1
        remaining = st.session_state.sentences_to_end
        
        if remaining > 0:
            # Generate AI's ending continuation
            ai_response = generate_ending_continuation(
                text,
                ending_type=st.session_state.ending_type,
                remaining_sentences=remaining
            )
            
            # 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
            )
            
            # 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()  # Using st.rerun() instead of st.experimental_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) -> str:
    """Generate the final ending sentence 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 sentence that:
                - Brings joy and satisfaction
                - Resolves the main story elements
                - Ends on a positive note
                Maximum 2 sentences.
            """,
            "Mysterious Ending": """
                Create a final mysterious ending sentence that:
                - Leaves an intriguing question
                - Creates a sense of wonder
                - Maintains some mystery
                Maximum 2 sentences.
            """,
            "Lesson Learned": """
                Create a final sentence that:
                - Shows what was learned
                - Provides a moral lesson
                - Connects to the story's events
                Maximum 2 sentences.
            """,
            "Surprise Ending": """
                Create a final twist ending that:
                - Provides an unexpected but logical conclusion
                - Connects to previous story elements
                - Creates a satisfying surprise
                Maximum 2 sentences.
            """
        }

        response = client.chat.completions.create(
            model="gpt-4",
            messages=[
                {
                    "role": "system",
                    "content": f"""
                    You are a master storyteller creating the final ending.
                    {ending_prompts[ending_type]}
                    This must be the absolute final sentence(s) of the story.
                    Make it conclusive and satisfying.
                    """
                },
                {
                    "role": "user",
                    "content": f"Story context:\n{story_summary}\n\nCreate the final ending:"
                }
            ],
            max_tokens=100,
            temperature=0.7
        )
        
        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) -> 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
            """
        }

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": f"""
                    {ending_prompts[ending_type]}
                    
                    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.
                    - Indicate awareness of the remaining sentences.
                    """
                },
                {
                    "role": "user",
                    "content": f"Continue and work towards ending this story: {text}"
                }
            ],
            max_tokens=100,
            temperature=0.7
        )

        return response.choices[0].message.content.strip()
        
    except Exception as e:
        logging.error(f"Error generating ending continuation: {str(e)}")
        return "The story moved towards its conclusion..."


def complete_story():
    """Handle story completion and celebration"""
    st.balloons()  # แสดงเอฟเฟคฉลอง
    
    # สร้าง Story Summary
    story_summary = generate_story_summary(st.session_state.story)
    
    # แสดงหน้าจบเรื่อง
    st.markdown(f"""
        <div style="
            background-color: #e8f5e9;
            padding: 20px;
            border-radius: 10px;
            text-align: center;
            margin: 20px 0;
        ">
            <h2>🎉 ยินดีด้วย! คุณเขียนเรื่องราวจบสมบูรณ์แล้ว</h2>
            
            <div style="margin: 20px 0;">
                <h3>📝 สรุปเรื่องราว</h3>
                <p>{story_summary}</p>
            </div>
            
            <div style="margin: 20px 0;">
                <h3>🏆 ความสำเร็จ</h3>
                <p>จำนวนประโยค: {len(st.session_state.story)}</p>
                <p>คำศัพท์ที่ใช้: {len(st.session_state.stats['vocabulary_used'])}</p>
                <p>ความแม่นยำ: {st.session_state.stats['accuracy_rate']:.1f}%</p>
            </div>
        </div>
    """, unsafe_allow_html=True)
    
    # แสดงตัวเลือกหลังจบเรื่อง
    col1, col2 = st.columns(2)
    with col1:
        if st.button("💾 บันทึกเรื่องราว", use_container_width=True):
            save_completed_story()
    with col2:
        if st.button("🔄 เริ่มเรื่องใหม่", use_container_width=True):
            reset_story()
            st.rerun()

def show_feedback_section():
    """Display writing feedback section"""
    if not st.session_state.feedback:
        return

    st.markdown("""
        <div style="margin-bottom: 20px;">
            <div class="thai-eng">
                <div class="thai">📝 คำแนะนำจากครู</div>
                <div class="eng">Writing Feedback</div>
            </div>
        </div>
    """, unsafe_allow_html=True)
    
    feedback_data = st.session_state.feedback
    if isinstance(feedback_data, dict) and feedback_data.get('has_errors'):
        # Show error feedback
        st.markdown(f"""
            <div style="
                background-color: #fff3e0;
                padding: 20px;
                border-radius: 10px;
                border-left: 4px solid #ff9800;
                margin: 10px 0;
            ">
                <p style="color: #e65100; margin-bottom: 15px;">
                    {feedback_data['feedback']}
                </p>
                <div style="
                    background-color: #fff;
                    padding: 15px;
                    border-radius: 8px;
                    margin-top: 10px;
                ">
                    <p style="color: #666; margin-bottom: 5px;">
                        ประโยคที่ถูกต้อง:
                    </p>
                    <p style="
                        color: #2e7d32;
                        font-weight: 500;
                        font-size: 1.1em;
                        margin: 0;
                    ">
                        {feedback_data['corrected']}
                    </p>
                </div>
            </div>
        """, unsafe_allow_html=True)
        
        # Correction button
        if st.button(
            "✍️ แก้ไขประโยคให้ถูกต้อง",
            key="correct_button",
            help="คลิกเพื่อแก้ไขประโยคให้ถูกต้องตามคำแนะนำ"
        ):
            last_user_entry_idx = next(
                (i for i, entry in reversed(list(enumerate(st.session_state.story))) 
                 if entry['role'] == 'You'),
                None
            )
            if last_user_entry_idx is not None:
                apply_correction(last_user_entry_idx, feedback_data['corrected'])
                st.rerun()
    else:
        # Show success feedback
        st.markdown(f"""
            <div style="
                background-color: #e8f5e9;
                padding: 20px;
                border-radius: 10px;
                border-left: 4px solid #4caf50;
                margin: 10px 0;
            ">
                <p style="color: #2e7d32; margin: 0;">
                    {feedback_data.get('feedback', '✨ เขียนได้ถูกต้องแล้วค่ะ!')}
                </p>
            </div>
        """, unsafe_allow_html=True)

def show_writing_tools():
    """Display writing tools section"""
    with st.expander("✨ เครื่องมือช่วยเขียน | Writing Tools"):
        col1, col2 = st.columns(2)
        
        with col1:
            if st.button("🎯 ขอคำใบ้", use_container_width=True):
                with st.spinner("กำลังสร้างคำใบ้..."):
                    prompt = get_creative_prompt()
                    st.markdown(f"""
                        <div style="
                            background-color: #f3e5f5;
                            padding: 15px;
                            border-radius: 8px;
                            margin-top: 10px;
                        ">
                            <div style="color: #6a1b9a;">💭 {prompt['thai']}</div>
                            <div style="color: #666; font-style: italic;">
                                💭 {prompt['eng']}
                            </div>
                        </div>
                    """, unsafe_allow_html=True)
        
        with col2:
            if st.button("📚 คำศัพท์แนะนำ", use_container_width=True):
                with st.spinner("กำลังค้นหาคำศัพท์..."):
                    vocab_suggestions = get_vocabulary_suggestions()
                    st.markdown("#### 📚 คำศัพท์น่ารู้")
                    for word in vocab_suggestions:
                        st.markdown(f"""
                            <div style="
                                background-color: #f5f5f5;
                                padding: 10px;
                                border-radius: 5px;
                                margin: 5px 0;
                            ">
                                • {word}
                            </div>
                        """, unsafe_allow_html=True)

def show_achievements():
    """Display achievements and stats"""
    with st.container():
        # Points and Streak
        st.markdown(f"""
            <div style="
                background-color: #e3f2fd;
                padding: 20px;
                border-radius: 10px;
                margin-bottom: 20px;
                text-align: center;
            ">
                <h3 style="color: #1e88e5; margin: 0;">
                    🌟 คะแนนรวม: {st.session_state.points['total']}
                </h3>
                <p style="margin: 5px 0;">
                    🔥 Streak ปัจจุบัน: {st.session_state.points['streak']} ประโยค
                </p>
                <p style="margin: 5px 0;">
                    ⭐ Streak สูงสุด: {st.session_state.points['max_streak']} ประโยค
                </p>
            </div>
        """, unsafe_allow_html=True)

        # Writing Stats
        st.markdown("""
            <div style="margin-bottom: 15px;">
                <h3>📊 สถิติการเขียน</h3>
            </div>
        """, unsafe_allow_html=True)
        
        col1, col2 = st.columns(2)
        with col1:
            st.metric(
                "ประโยคที่เขียนทั้งหมด", 
                st.session_state.stats['total_sentences'],
                help="จำนวนประโยคทั้งหมดที่คุณได้เขียน"
            )
            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 show_story_input():
    """Display story input section"""
    st.markdown("""
        <div class="thai-eng">
            <div class="thai">✏️ ถึงตาคุณแล้ว</div>
            <div class="eng">Your Turn</div>
        </div>
    """, unsafe_allow_html=True)
    
    # Initialize clear_input flag if not exists
    if 'clear_input' not in st.session_state:
        st.session_state.clear_input = False
    
    # Default value for text input
    default_value = "" if st.session_state.clear_input else st.session_state.get('text_input', "")
    
    # If clear_input flag is True, reset it
    if st.session_state.clear_input:
        st.session_state.clear_input = False

    # 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
    text_input = st.text_area(
        "เขียนต่อจากเรื่องราว | Continue the story:",
        value=st.session_state.get('text_input', ""),
        height=100,
        key="story_input_area",
        help="พิมพ์ประโยคภาษาอังกฤษเพื่อต่อเรื่อง",
        label_visibility="collapsed"
    )
    
    # Save current input to session state
    st.session_state.text_input = text_input
    
    # Submit button with character count
    col1, col2 = st.columns([3, 1])
    with col1:
        if st.button("📝 ส่งคำตอบ | Submit", use_container_width=True):
            if not text_input.strip():
                st.warning("กรุณาเขียนข้อความก่อนส่ง")
                return
            try:
                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:
        char_count = len(text_input)
        st.markdown(f"""
            <div style="text-align: right; color: {'red' if char_count > 200 else '#666'};">
                {char_count}/200 ตัวอักษร
            </div>
        """, unsafe_allow_html=True)

def handle_story_submission(text: str):
    """Handle story submission and processing"""
    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)
        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()
        
        # Set flag to clear input on next rerun
        st.session_state.clear_input = True
        
        # Rerun to update UI
        st.rerun()
        
    except Exception as e:
        logging.error(f"Error in story submission: {str(e)}")
        raise


def 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 show_completion_options():
    """Display options and summary when story is completed"""
    try:
        # Show completion celebration
        st.balloons()
        
        # Create completion message
        st.markdown("""
            <div style="
                background-color: #e8f5e9;
                padding: 20px;
                border-radius: 10px;
                text-align: center;
                margin: 20px 0;
                border: 2px solid #4caf50;
            ">
                <h2 style="color: #2e7d32; margin-bottom: 15px;">
                    🎉 ยินดีด้วย! คุณเขียนเรื่องราวจบสมบูรณ์แล้ว
                </h2>
                <p style="color: #1b5e20;">
                    เรื่องราวของคุณจบลงด้วย: {st.session_state.ending_type}
                </p>
            </div>
        """, unsafe_allow_html=True)
        
        # Show story statistics
        st.markdown("### 📊 สรุปเรื่องราว")
        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="จำนวนคำศัพท์ที่ไม่ซ้ำกัน"
            )
        
        with col2:
            st.metric(
                "คะแนนรวม",
                st.session_state.points['total'],
                help="คะแนนรวมที่ได้รับ"
            )
            st.metric(
                "ความแม่นยำ",
                f"{st.session_state.stats['accuracy_rate']:.1f}%",
                help="อัตราการเขียนถูกต้อง"
            )
        
        # Show action buttons with unique keys
        st.markdown("### 🎯 ต้องการทำอะไรต่อ?")
        col1, col2 = st.columns(2)
        
        with col1:
            if st.button("💾 บันทึกเรื่องราว", 
                        key="save_story_completion", 
                        use_container_width=True):
                show_save_dialog()
        
        with col2:
            if st.button("🔄 เริ่มเรื่องใหม่", 
                        key="new_story_completion", 
                        use_container_width=True):
                reset_story_with_confirmation()

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

def show_save_dialog():
    """Display save options dialog"""
    try:
        st.markdown("### 💾 บันทึกเรื่องราว")
        
        col1, col2 = st.columns(2)
        
        with col1:
            # PDF Export with unique key
            if st.button("📑 บันทึกเป็น PDF", 
                        key="save_pdf_button", 
                        use_container_width=True):
                try:
                    pdf = create_story_pdf()
                    st.download_button(
                        "📥 ดาวน์โหลด PDF",
                        data=pdf,
                        file_name=f"story_{datetime.now().strftime('%Y%m%d')}.pdf",
                        mime="application/pdf",
                        key="download_pdf_button"
                    )
                except Exception as e:
                    logging.error(f"Error creating PDF: {str(e)}")
                    st.error("เกิดข้อผิดพลาดในการสร้างไฟล์ PDF")
        
        with col2:
            # JSON Save with unique key
            if st.button("💾 บันทึกความก้าวหน้า", 
                        key="save_progress_button", 
                        use_container_width=True):
                try:
                    story_data = {
                        'timestamp': datetime.now().isoformat(),
                        'level': st.session_state.level,
                        'story': st.session_state.story,
                        'achievements': st.session_state.achievements,
                        'stats': {
                            key: list(value) if isinstance(value, set) else value
                            for key, value in st.session_state.stats.items()
                        },
                        'points': st.session_state.points
                    }
                    
                    st.download_button(
                        "📥 ดาวน์โหลดไฟล์บันทึก",
                        data=json.dumps(story_data, ensure_ascii=False, indent=2),
                        file_name=f"story_progress_{datetime.now().strftime('%Y%m%d')}.json",
                        mime="application/json",
                        key="download_json_button"
                    )
                except Exception as e:
                    logging.error(f"Error saving progress: {str(e)}")
                    st.error("เกิดข้อผิดพลาดในการบันทึกความก้าวหน้า")
    except Exception as e:
        logging.error(f"Error showing save dialog: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการแสดงตัวเลือกการบันทึก")

def reset_story_with_confirmation():
    """Reset story with confirmation dialog"""
    try:
        if st.session_state.story:
            if st.checkbox("✅ ยืนยันการเริ่มใหม่", 
                          key="reset_confirmation_checkbox"):
                st.warning("การเริ่มใหม่จะลบเรื่องราวปัจจุบันทั้งหมด")
                if st.button("🔄 ยืนยันการเริ่มใหม่", 
                            key="confirm_reset_button", 
                            use_container_width=True):
                    reset_story()
                    st.success("เริ่มต้นใหม่เรียบร้อยแล้ว!")
                    st.rerun()
    except Exception as e:
        logging.error(f"Error in reset confirmation: {str(e)}")
        st.error("เกิดข้อผิดพลาดในการรีเซ็ตเรื่อง")

def create_story_pdf():
    """Create PDF of the story"""
    try:
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(buffer, pagesize=A4)
        story_elements = []
        
        # Add title
        styles = getSampleStyleSheet()
        title_style = ParagraphStyle(
            'CustomTitle',
            parent=styles['Title'],
            fontSize=24,
            spaceAfter=30
        )
        story_elements.append(Paragraph("My Story", title_style))
        
        # Add content
        content_style = ParagraphStyle(
            'CustomBody',
            parent=styles['Normal'],
            fontSize=12,
            leading=16,
            spaceAfter=12
        )
        
        for entry in st.session_state.story:
            if entry['role'] == 'AI':
                color = colors.blue
            else:
                color = colors.black
            
            p = Paragraph(
                f"<font color={color}>{entry['content']}</font>",
                content_style
            )
            story_elements.append(p)
        
        # Build PDF
        doc.build(story_elements)
        pdf = buffer.getvalue()
        buffer.close()
        
        return pdf
        
    except Exception as e:
        logging.error(f"Error creating PDF: {str(e)}")
        raise

# === 6. MAIN APPLICATION LOGIC ===
def main():
    try:
        # Initialize states
        init_session_state()
        init_theme_state()
        
        # 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()

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