import streamlit as st
import json
import datetime
from openai import OpenAI
from typing import Dict, List, Set
import io
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.pdfgen import canvas
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from datetime import datetime 
from typing import Dict, List, Optional, Tuple
import random

# 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'
    }
}

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

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

# Initialize OpenAI client
client = OpenAI()

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

# Add custom CSS including new styles for level selection
st.markdown("""
    <style>
    @import url('https://fonts.googleapis.com/css2?family=Sarabun:wght@400;700&display=swap');
    
    /* ส่วน CSS เดิม */
    .thai-eng {
        font-size: 1.1em;
        padding: 10px;
        background-color: #f8f9fa;
        border-radius: 8px;
        margin: 10px 0;
    }
    .thai {
        color: #1e88e5;
        font-family: 'Sarabun', sans-serif;
    }
    .eng {
        color: #333;
    }
    .level-info {
        background-color: #e3f2fd;
        padding: 10px;
        border-radius: 8px;
        margin: 10px 0;
        font-size: 0.9em;
    }
    .parent-guide {
        background-color: #fff3e0;
        padding: 15px;
        border-radius: 8px;
        margin: 15px 0;
        border-left: 4px solid #ff9800;
    }

    /* เพิ่ม CSS สำหรับ Theme Cards */
    .theme-header {
        text-align: center;
        margin-bottom: 20px;
    }
    
    .theme-header h3 {
        color: #1e88e5;
        font-family: 'Sarabun', sans-serif;
    }
    
    .theme-card {
        position: relative;
        overflow: hidden;
    }
    
    .theme-card:hover {
        transform: translateY(-5px);
        box-shadow: 0 5px 15px rgba(0,0,0,0.2);
    }
    
    .theme-card:before {
        content: '';
        position: absolute;
        top: 0;
        left: 0;
        width: 100%;
        height: 100%;
        background: linear-gradient(45deg, rgba(255,255,255,0.1), rgba(255,255,255,0.5));
        opacity: 0;
        transition: opacity 0.3s ease;
    }
    
    .theme-card:hover:before {
        opacity: 1;
    }
    
    .theme-card.selected {
        transform: scale(0.98);
        box-shadow: 0 2px 5px rgba(0,0,0,0.2);
    }

    /* CSS สำหรับปุ่ม Change Theme */
    .stButton button {
        padding: 0.5rem 1rem;
        border-radius: 0.5rem;
        border: 1px solid rgba(49, 51, 63, 0.2);
        background-color: #ffffff;
        color: #31333f;
        transition: all 0.2s ease;
    }
    
    .stButton button:hover {
        border-color: #1e88e5;
        color: #1e88e5;
        transform: translateY(-1px);
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }

    /* ปรับแต่ง tooltip */
    .stTooltipIcon {
        color: #1e88e5;
    }

    /* Responsive Design */
    @media (max-width: 768px) {
        .theme-card {
            min-height: 120px;
        }
        
        .theme-card h4 {
            font-size: 0.9em;
        }
        
        .theme-card p {
            font-size: 0.8em;
        }
    }

    /* Story Display Styling */
    .story-message {
        padding: 10px;
        margin: 5px 0;
        border-radius: 5px;
    }
    .story-message.ai {
        background-color: #f0f7ff;
    }
    .story-message.user {
        background-color: #fff;
    }
    .story-message.corrected {
        background-color: #f0f9ff;
    }
    </style>
""", unsafe_allow_html=True)

# Initialize session state variables
def init_session_state():
    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
    if 'story' not in st.session_state:
        st.session_state.story = []
    if 'feedback' not in st.session_state:
        st.session_state.feedback = None
    if 'level' not in st.session_state:
        st.session_state.level = 'Beginner'
    if 'unique_words' not in st.session_state:
        st.session_state.unique_words = set()
    if 'total_words' not in st.session_state:
        st.session_state.total_words = 0
    if 'should_reset' not in st.session_state:
        st.session_state.should_reset = False
    if 'user_input' not in st.session_state:
        st.session_state.user_input = ""
    if 'points' not in st.session_state:
        st.session_state.points = {
            'total': 0,
            'perfect_sentences': 0,
            'corrections_made': 0,
            'streak': 0,
            'max_streak': 0
        }
    if 'corrections' not in st.session_state:
        st.session_state.corrections = {}
    if 'stats' not in st.session_state:
        st.session_state.stats = {
            'total_sentences': 0,
            'correct_first_try': 0,
            'accuracy_rate': 0.0,
            'vocabulary_used': set(),
            'corrections_made': 0
        }
    if 'achievements' not in st.session_state:
        st.session_state.achievements = []

init_session_state()

# เพิ่มฟังก์ชันสำหรับจัดการ input
def clear_input():
    st.session_state.user_input = ""

# Callback function for submit button
def submit_story():
    """Handle story submission and processing."""
    if st.session_state.text_input.strip():
        user_text = st.session_state.text_input.strip()
        
        # ตรวจสอบว่ามีเรื่องเริ่มต้นหรือยัง
        if not st.session_state.story:
            st.error("กรุณาเลือกธีมเรื่องราวก่อนเริ่มเขียน")
            return

        try:
            # รับ feedback และตรวจสอบความถูกต้อง
            feedback_data = provide_feedback(user_text, st.session_state.level)
            st.session_state.feedback = feedback_data
            is_correct = not feedback_data.get('has_errors', False)
            
            # เพิ่มประโยคของผู้ใช้ (เพิ่มครั้งเดียว)
            st.session_state.story.append({
                "role": "You",
                "content": user_text,
                "is_corrected": False,
                "is_correct": is_correct
            })
            
            # เพิ่มคำศัพท์ที่ใช้
            words = set(user_text.lower().split())
            st.session_state.stats['vocabulary_used'].update(words)
            
            # อัพเดตคะแนนและ achievements
            update_points(is_correct)
            update_achievements()
            
            # Generate AI continuation
            text_for_continuation = feedback_data['corrected'] if feedback_data['has_errors'] else user_text
            ai_response = generate_story_continuation(text_for_continuation, st.session_state.level)
            
            # เพิ่มการตอบกลับของ AI
            st.session_state.story.append({
                "role": "AI", 
                "content": ai_response
            })
            
            # Clear input
            st.session_state.text_input = ""
            
        except Exception as e:
            st.error("เกิดข้อผิดพลาดในการวิเคราะห์ประโยค กรุณาลองใหม่อีกครั้ง")
            st.error(f"Debug - Error in submit_story: {str(e)}")


# Theme Helper Functions
def get_available_themes(level: str) -> List[Dict]:
    """Get list of themes available for the current level."""
    return [
        theme for theme in story_themes.values()
        if level in theme['level_range']
    ]

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.
    """
    theme = story_themes.get(theme_id, {})
    
    # Theme-specific elements for dynamic generation
    theme_elements = {
        'fantasy': {
            'characters': {
                'Beginner': [
                    ('young wizard', 'พ่อมดน้อย'),
                    ('fairy', 'นางฟ้า'),
                    ('friendly dragon', 'มังกรใจดี'),
                    ('magical cat', 'แมวเวทมนตร์'),
                ],
                'Intermediate': [
                    ('mysterious wizard', 'พ่อมดลึกลับ'),
                    ('ancient dragon', 'มังกรโบราณ'),
                    ('enchanted princess', 'เจ้าหญิงต้องมนตร์'),
                ],
                'Advanced': [
                    ('legendary sorcerer', 'จอมเวทในตำนาน'),
                    ('mythical creature', 'สัตว์ในตำนาน'),
                    ('wise elder wizard', 'พ่อมดผู้เฒ่าผู้เป็นปราชญ์'),
                ]
            },
            'locations': {
                'Beginner': [
                    ('magical garden', 'สวนเวทมนตร์'),
                    ('enchanted forest', 'ป่าวิเศษ'),
                ],
                'Intermediate': [
                    ('ancient castle', 'ปราสาทโบราณ'),
                    ('wizard school', 'โรงเรียนเวทมนตร์'),
                ],
                'Advanced': [
                    ('mystical realm', 'อาณาจักรเวทมนตร์'),
                    ('floating islands', 'เกาะลอยฟ้า'),
                ]
            },
            'objects': {
                'Beginner': [
                    ('magic wand', 'ไม้กายสิทธิ์'),
                    ('glowing crystal', 'คริสตัลเรืองแสง'),
                ],
                'Intermediate': [
                    ('ancient spellbook', 'ตำราเวทโบราณ'),
                    ('magical artifact', 'วัตถุวิเศษ'),
                ],
                'Advanced': [
                    ('legendary sword', 'ดาบในตำนาน'),
                    ('mystical orb', 'ลูกแก้ววิเศษ'),
                ]
            }
        },
        'nature': {
            'characters': {
                'Beginner': [
                    ('little bird', 'นกน้อย'),
                    ('friendly squirrel', 'กระรอกน้อยใจดี'),
                    ('playful rabbit', 'กระต่ายซน'),
                ],
                'Intermediate': [
                    ('wise owl', 'นกฮูกผู้รอบรู้'),
                    ('busy beaver', 'บีเวอร์ขยัน'),
                    ('clever fox', 'จิ้งจอกเจ้าปัญญา'),
                ],
                'Advanced': [
                    ('majestic eagle', 'อินทรีผู้สง่างาม'),
                    ('mysterious panther', 'เสือดำลึกลับ'),
                ]
            },
            'locations': {
                'Beginner': [
                    ('flower garden', 'สวนดอกไม้'),
                    ('small pond', 'สระน้ำเล็กๆ'),
                ],
                'Intermediate': [
                    ('deep forest', 'ป่าลึก'),
                    ('mountain peak', 'ยอดเขาสูง'),
                ],
                'Advanced': [
                    ('ancient rainforest', 'ป่าดึกดำบรรพ์'),
                    ('hidden valley', 'หุบเขาซ่อนเร้น'),
                ]
            },
            'objects': {
                'Beginner': [
                    ('colorful flower', 'ดอกไม้สีสวย'),
                    ('fallen leaf', 'ใบไม้ร่วง'),
                ],
                'Intermediate': [
                    ('old tree', 'ต้นไม้เก่าแก่'),
                    ('clear stream', 'ลำธารใส'),
                ],
                'Advanced': [
                    ('rare plant', 'พืชหายาก'),
                    ('ancient tree', 'ต้นไม้โบราณ'),
                ]
            }
        },
        'space': {
            'characters': {
                'Beginner': [
                    ('friendly alien', 'มนุษย์ต่างดาวใจดี'),
                    ('little astronaut', 'นักบินอวกาศตัวน้อย'),
                    ('space robot', 'หุ่นยนต์อวกาศ'),
                ],
                'Intermediate': [
                    ('space explorer', 'นักสำรวจอวกาศ'),
                    ('alien scientist', 'นักวิทยาศาสตร์ต่างดาว'),
                    ('space pilot', 'นักบินยานอวกาศ'),
                ],
                'Advanced': [
                    ('galactic commander', 'ผู้บัญชาการกาแล็กซี่'),
                    ('space archaeologist', 'นักโบราณคดีอวกาศ'),
                ]
            },
            'locations': {
                'Beginner': [
                    ('small planet', 'ดาวเคราะห์เล็กๆ'),
                    ('space station', 'สถานีอวกาศ'),
                ],
                'Intermediate': [
                    ('mysterious galaxy', 'กาแล็กซี่ลึกลับ'),
                    ('alien world', 'โลกต่างดาว'),
                ],
                'Advanced': [
                    ('black hole', 'หลุมดำ'),
                    ('distant nebula', 'เนบิวลาไกลโพ้น'),
                ]
            },
            'objects': {
                'Beginner': [
                    ('shiny meteor', 'ดาวตกเปล่งประกาย'),
                    ('space telescope', 'กล้องดูดาว'),
                ],
                'Intermediate': [
                    ('advanced spaceship', 'ยานอวกาศลำใหม่'),
                    ('mysterious signal', 'สัญญาณลึกลับ'),
                ],
                'Advanced': [
                    ('alien artifact', 'วัตถุโบราณจากต่างดาว'),
                    ('space anomaly', 'ความผิดปกติในอวกาศ'),
                ]
            }
        },
        'adventure': {
            'characters': {
                'Beginner': [
                    ('young explorer', 'นักผจญภัยตัวน้อย'),
                    ('treasure hunter', 'นักล่าสมบัติ'),
                    ('brave sailor', 'กะลาสีผู้กล้า'),
                ],
                'Intermediate': [
                    ('experienced guide', 'ไกด์ผู้ชำนาญ'),
                    ('mysterious traveler', 'นักเดินทางลึกลับ'),
                ],
                'Advanced': [
                    ('legendary explorer', 'นักสำรวจในตำนาน'),
                    ('ancient guardian', 'ผู้พิทักษ์โบราณ'),
                ]
            },
            'locations': {
                'Beginner': [
                    ('hidden cave', 'ถ้ำซ่อนเร้น'),
                    ('treasure island', 'เกาะสมบัติ'),
                ],
                'Intermediate': [
                    ('ancient ruins', 'ซากปรักหักพัง'),
                    ('mysterious temple', 'วัดลึกลับ'),
                ],
                'Advanced': [
                    ('lost city', 'เมืองที่สาบสูญ'),
                    ('forbidden valley', 'หุบเขาต้องห้าม'),
                ]
            },
            'objects': {
                'Beginner': [
                    ('old map', 'แผนที่เก่า'),
                    ('golden compass', 'เข็มทิศทองคำ'),
                ],
                'Intermediate': [
                    ('ancient scroll', 'ม้วนกระดาษโบราณ'),
                    ('magical key', 'กุญแจวิเศษ'),
                ],
                'Advanced': [
                    ('legendary artifact', 'วัตถุโบราณในตำนาน'),
                    ('sacred relic', 'วัตถุศักดิ์สิทธิ์'),
                ]
            }
        },
        'school': {
            'characters': {
                'Beginner': [
                    ('new student', 'นักเรียนใหม่'),
                    ('kind teacher', 'คุณครูใจดี'),
                    ('best friend', 'เพื่อนรัก'),
                ],
                'Intermediate': [
                    ('class president', 'หัวหน้าห้อง'),
                    ('school librarian', 'บรรณารักษ์'),
                ],
                'Advanced': [
                    ('exchange student', 'นักเรียนแลกเปลี่ยน'),
                    ('science genius', 'อัจฉริยะวิทยาศาสตร์'),
                ]
            },
            'locations': {
                'Beginner': [
                    ('classroom', 'ห้องเรียน'),
                    ('school playground', 'สนามเด็กเล่น'),
                ],
                'Intermediate': [
                    ('science lab', 'ห้องทดลองวิทยาศาสตร์'),
                    ('school library', 'ห้องสมุด'),
                ],
                'Advanced': [
                    ('school theater', 'โรงละครโรงเรียน'),
                    ('computer room', 'ห้องคอมพิวเตอร์'),
                ]
            },
            'objects': {
                'Beginner': [
                    ('new book', 'หนังสือเล่มใหม่'),
                    ('special pencil', 'ดินสอวิเศษ'),
                ],
                'Intermediate': [
                    ('science project', 'โครงงานวิทยาศาสตร์'),
                    ('old diary', 'ไดอารี่เก่า'),
                ],
                'Advanced': [
                    ('robotics kit', 'ชุดประดิษฐ์หุ่นยนต์'),
                    ('ancient textbook', 'ตำราเรียนโบราณ'),
                ]
            }
        },
        'superhero': {
            'characters': {
                'Beginner': [
                    ('young hero', 'ฮีโร่ตัวน้อย'),
                    ('super pet', 'สัตว์เลี้ยงพลังพิเศษ'),
                    ('friendly sidekick', 'ผู้ช่วยฮีโร่'),
                ],
                'Intermediate': [
                    ('masked vigilante', 'ฮีโร่ปริศนา'),
                    ('super teacher', 'คุณครูพลังพิเศษ'),
                ],
                'Advanced': [
                    ('legendary hero', 'ฮีโร่ในตำนาน'),
                    ('protector of justice', 'ผู้พิทักษ์ความยุติธรรม'),
                ]
            },
            'locations': {
                'Beginner': [
                    ('hero school', 'โรงเรียนฮีโร่'),
                    ('secret hideout', 'ที่ซ่อนลับ'),
                ],
                'Intermediate': [
                    ('hero headquarters', 'ศูนย์บัญชาการฮีโร่'),
                    ('training ground', 'สนามฝึกซ้อม'),
                ],
                'Advanced': [
                    ('hero academy', 'สถาบันฝึกฮีโร่'),
                    ('crisis center', 'ศูนย์รับมือวิกฤต'),
                ]
            },
            'objects': {
                'Beginner': [
                    ('hero mask', 'หน้ากากฮีโร่'),
                    ('power crystal', 'คริสตัลพลัง'),
                ],
                'Intermediate': [
                    ('super gadget', 'อุปกรณ์พิเศษ'),
                    ('power suit', 'ชุดพลังพิเศษ'),
                ],
                'Advanced': [
                    ('ultimate weapon', 'อาวุธสุดยอด'),
                    ('ancient power source', 'แหล่งพลังโบราณ'),
                ]
            }
        },
        'mystery': {
            'characters': {
                'Beginner': [
                    ('young detective', 'นักสืบตัวน้อย'),
                    ('mysterious neighbor', 'เพื่อนบ้านปริศนา'),
                    ('helpful friend', 'เพื่อนผู้ช่วย'),
                ],
                'Intermediate': [
                    ('famous detective', 'นักสืบชื่อดัง'),
                    ('mystery writer', 'นักเขียนนิยายปริศนา'),
                ],
                'Advanced': [
                    ('master detective', 'ยอดนักสืบ'),
                    ('crime expert', 'ผู้เชี่ยวชาญคดี'),
                ]
            },
            'locations': {
                'Beginner': [
                    ('old house', 'บ้านเก่า'),
                    ('mystery shop', 'ร้านค้าปริศนา'),
                ],
                'Intermediate': [
                    ('abandoned building', 'ตึกร้าง'),
                    ('detective office', 'สำนักงานนักสืบ'),
                ],
                'Advanced': [
                    ('secret laboratory', 'ห้องทดลองลับ'),
                    ('mystery mansion', 'คฤหาสน์ปริศนา'),
                ]
            },
            'objects': {
                'Beginner': [
                    ('mysterious letter', 'จดหมายปริศนา'),
                    ('strange footprint', 'รอยเท้าประหลาด'),
                ],
                'Intermediate': [
                    ('ancient cipher', 'รหัสลับโบราณ'),
                    ('hidden clue', 'เบาะแสซ่อนเร้น'),
                ],
                'Advanced': [
                    ('secret document', 'เอกสารลับ'),
                    ('mysterious artifact', 'วัตถุปริศนา'),
                ]
            }
        },
        'science': {
            'characters': {
                'Beginner': [
                    ('young scientist', 'นักวิทยาศาสตร์ตัวน้อย'),
                    ('robot friend', 'หุ่นยนต์เพื่อนรัก'),
                    ('curious student', 'นักเรียนช่างสงสัย'),
                ],
                'Intermediate': [
                    ('brilliant inventor', 'นักประดิษฐ์อัจฉริยะ'),
                    ('science teacher', 'คุณครูวิทยาศาสตร์'),
                ],
                'Advanced': [
                    ('famous researcher', 'นักวิจัยชื่อดัง'),
                    ('genius professor', 'ศาสตราจารย์อัจฉริยะ'),
                ]
            },
            'locations': {
                'Beginner': [
                    ('science lab', 'ห้องทดลองวิทยาศาสตร์'),
                    ('invention workshop', 'ห้องประดิษฐ์'),
                ],
                'Intermediate': [
                    ('research center', 'ศูนย์วิจัย'),
                    ('robotics lab', 'ห้องปฏิบัติการหุ่นยนต์'),
                ],
                'Advanced': [
                    ('quantum laboratory', 'ห้องปฏิบัติการควอนตัม'),
                    ('innovation center', 'ศูนย์นวัตกรรม'),
                ]
            },
            'objects': {
                'Beginner': [
                    ('colorful chemical', 'สารเคมีสีสันสดใส'),
                    ('strange invention', 'สิ่งประดิษฐ์แปลกใหม่'),
                ],
                'Intermediate': [
                    ('advanced machine', 'เครื่องจักรทันสมัย'),
                    ('experimental device', 'อุปกรณ์ทดลอง'),
                ],
                'Advanced': [
                    ('groundbreaking discovery', 'การค้นพบที่ยิ่งใหญ่'),
                    ('revolutionary invention', 'สิ่งประดิษฐ์ปฏิวัติโลก'),
                ]
            }
        }
    }

    def get_random_element(category: str) -> Tuple[str, str]:
        """Get random element from the specified category for current theme and level."""
        elements = theme_elements.get(theme_id, {}).get(category, {}).get(level, [])
        if not elements:
            elements = theme_elements.get(theme_id, {}).get(category, {}).get('Beginner', [])
        return random.choice(elements) if elements else ('', '')

    # Get random elements
    character_en, character_th = get_random_element('characters')
    location_en, location_th = get_random_element('locations')
    object_en, object_th = get_random_element('objects')

    # Templates for different levels
    templates = {
        'Beginner': {
            'en': [
                f"One day, {character_en} found {object_en} in {location_en}...",
                f"In {location_en}, {character_en} saw something special...",
                f"{character_en} was walking in {location_en} when suddenly...",
                f"The story begins when {character_en} discovered {object_en}...",
            ],
            'th': [
                f"วันหนึ่ง {character_th}เจอ{object_th}ใน{location_th}...",
                f"ที่{location_th} {character_th}เห็นบางสิ่งที่พิเศษ...",
                f"{character_th}กำลังเดินอยู่ใน{location_th} เมื่อจู่ๆ...",
                f"เรื่องราวเริ่มต้นเมื่อ{character_th}ค้นพบ{object_th}...",
            ]
        },
        'Intermediate': {
            'en': [
                f"While exploring {location_en}, {character_en} discovered {object_en} that seemed magical...",
                f"The story begins when {character_en} encountered a mysterious {object_en} in {location_en}...",
                f"Something strange was happening in {location_en}, and {character_en} knew it had to do with {object_en}...",
                f"Nobody knew why {character_en} found {object_en} in {location_en}, but...",
            ],
            'th': [
                f"ขณะที่สำรวจ{location_th} {character_th}ได้ค้นพบ{object_th}ที่ดูมีเวทมนตร์...",
                f"เรื่องราวเริ่มต้นเมื่อ{character_th}พบกับ{object_th}ที่ลึกลับใน{location_th}...",
                f"มีบางสิ่งประหลาดเกิดขึ้นใน{location_th} และ{character_th}รู้ว่ามันเกี่ยวข้องกับ{object_th}...",
                f"ไม่มีใครรู้ว่าทำไม{character_th}ถึงพบ{object_th}ใน{location_th} แต่...",
            ]
        },
        'Advanced': {
            'en': [
                f"Legend speaks of {object_en} hidden within {location_en}, and {character_en} was destined to find it...",
                f"In the depths of {location_en}, {character_en} uncovered an ancient mystery surrounding {object_en}...",
                f"As {character_en} ventured through {location_en}, the secrets of {object_en} began to unfold...",
                f"The discovery of {object_en} by {character_en} in {location_en} would change everything...",
            ],
            'th': [
                f"ตำนานเล่าขานถึง{object_th}ที่ซ่อนอยู่ใน{location_th} และ{character_th}ถูกลิขิตให้เป็นผู้ค้นพบ...",
                f"ในส่วนลึกของ{location_th} {character_th}ได้พบกับปริศนาโบราณเกี่ยวกับ{object_th}...",
                f"ขณะที่{character_th}ผจญภัยใน{location_th} ความลับของ{object_th}ก็เริ่มเผยออกมา...",
                f"การค้นพบ{object_th}โดย{character_th}ใน{location_th}จะเปลี่ยนแปลงทุกสิ่ง...",
            ]
        }
    }

    # Select random template for the current level
    level_templates = templates.get(level, templates['Beginner'])
    starter_en = random.choice(level_templates['en'])
    starter_th = random.choice(level_templates['th'])

    return {
        'en': starter_en,
        'th': starter_th
    }

def convert_sets_to_lists(obj):
    """แปลง set เป็น list สำหรับการ serialize เป็น JSON"""
    if isinstance(obj, set):
        return list(obj)
    elif isinstance(obj, dict):
        return {key: convert_sets_to_lists(value) for key, value in obj.items()}
    elif isinstance(obj, list):
        return [convert_sets_to_lists(item) for item in obj]
    return obj
            
def get_theme_story_starter(theme_id: str, level: str) -> Dict[str, str]:
    """Get a dynamically generated story starter for the selected theme and level."""
    return generate_dynamic_story_starter(theme_id, level)

def get_theme_vocabulary(theme_id: str, level: str) -> List[Dict[str, str]]:
    """
    Get theme-specific vocabulary with Thai translations and example sentences.
    Returns list of dictionaries containing word, type, thai meaning, and example.
    """
    vocab_data = {
        'fantasy': {
            'Beginner': [
                {
                    'word': 'magic',
                    'type': 'noun',
                    'thai': 'เวทมนตร์',
                    'example_en': 'The magic made flowers appear.',
                    'example_th': 'เวทมนตร์ทำให้ดอกไม้ปรากฏขึ้น'
                },
                {
                    'word': 'wizard',
                    'type': 'noun',
                    'thai': 'พ่อมด',
                    'example_en': 'The wizard waved his wand.',
                    'example_th': 'พ่อมดโบกไม้กายสิทธิ์'
                }
            ],
            'Intermediate': [
                {
                    'word': 'enchanted',
                    'type': 'adjective',
                    'thai': 'ที่ถูกสะกดด้วยเวทมนตร์',
                    'example_en': 'They found an enchanted forest.',
                    'example_th': 'พวกเขาพบป่าที่ถูกสะกดด้วยเวทมนตร์'
                }
            ]
        },
        'mystery': {
            'Beginner': [
                {
                    'word': 'clue',
                    'type': 'noun',
                    'thai': 'เบาะแส',
                    'example_en': 'The detective found an important clue.',
                    'example_th': 'นักสืบพบเบาะแสสำคัญ'
                }
            ]
        }
    }

    # แสดงคำศัพท์พร้อมตัวอย่างประโยคในหน้า UI
    def show_vocabulary(word_list: List[Dict[str, str]]):
        st.markdown("### 📚 คำศัพท์น่ารู้ประจำธีม")
        for word_data in word_list:
            st.markdown(f"""
                <div style="background-color: #f5f5f5; padding: 10px; border-radius: 5px; margin: 5px 0;">
                    <strong>{word_data['word']}</strong> ({word_data['type']}) - {word_data['thai']}
                    <br/>
                    <span style="color: #666;">
                        🇬🇧 {word_data['example_en']}<br/>
                        🇹🇭 {word_data['example_th']}
                    </span>
                </div>
            """, unsafe_allow_html=True)

    # ปรับปรุงฟังก์ชัน show_theme_vocabulary
    def show_theme_vocabulary():
        if st.session_state.current_theme:
            vocab_list = vocab_data.get(st.session_state.current_theme, {}).get(st.session_state.level, [])
            if vocab_list:
                show_vocabulary(vocab_list)
                
                # เพิ่มปุ่มสำหรับฝึกคำศัพท์
                if st.button("🎯 ฝึกคำศัพท์"):
                    practice_vocabulary(vocab_list)

    # เพิ่มฟังก์ชันฝึกคำศัพท์
    def practice_vocabulary(vocab_list: List[Dict[str, str]]):
        st.markdown("### 🎮 ฝึกคำศัพท์")
        
        # สุ่มคำศัพท์มาทำแบบฝึกหัด
        word = random.choice(vocab_list)
        
        # สร้างตัวเลือก
        choices = [word['thai']]  # คำตอบที่ถูก
        other_words = [w['thai'] for w in vocab_list if w['thai'] != word['thai']]
        choices.extend(random.sample(other_words, min(3, len(other_words))))
        random.shuffle(choices)
        
        st.write(f"คำว่า '{word['word']}' แปลว่าอะไร?")
        
        for choice in choices:
            if st.button(choice):
                if choice == word['thai']:
                    st.success("🎉 ถูกต้อง!")
                    st.write(f"ตัวอย่างประโยค: {word['example_en']}")
                    st.write(f"แปล: {word['example_th']}")
                else:
                    st.error("❌ ลองใหม่อีกครั้ง")

    return vocab_data.get(theme_id, {}).get(level, [])

# Add to session state initialization
def init_theme_state():
    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

# Theme Selection UI
def show_theme_selection():
    st.markdown("""
        <div class="theme-header">
            <h3>🎨 เลือกธีมเรื่องราว | Choose Story Theme</h3>
        </div>
    """, unsafe_allow_html=True)
    
    # สร้าง 2 rows, แต่ละ row มี 4 themes
    themes_per_row = 4
    available_themes = get_available_themes(st.session_state.level)
    
    # แบ่งธีมเป็น 2 rows
    for i in range(0, len(available_themes), themes_per_row):
        row_themes = available_themes[i:i + themes_per_row]
        cols = st.columns(themes_per_row)
        
        for col, theme in zip(cols, row_themes):
            with col:
                # สร้าง theme card แบบ interactive
                theme_card = f"""
                <div class="theme-card" 
                     onclick="this.classList.toggle('selected')"
                     style="background: linear-gradient(135deg, {theme['background_color']}, white);
                            border: 2px solid {theme['accent_color']};
                            border-radius: 12px;
                            padding: 15px;
                            margin: 5px;
                            cursor: pointer;
                            transition: all 0.3s ease;
                            box-shadow: 0 2px 5px rgba(0,0,0,0.1);
                            text-align: center;
                            min-height: 150px;">
                    <div style="font-size: 2em; margin-bottom: 8px;">
                        {theme['icon']}
                    </div>
                    <h4 style="color: {theme['accent_color']}; 
                               margin: 5px 0; 
                               font-family: 'Sarabun', sans-serif;">
                        {theme['name_th']}
                    </h4>
                    <p style="font-size: 0.85em; 
                             color: #666; 
                             margin-top: 5px;
                             display: -webkit-box;
                             -webkit-line-clamp: 2;
                             -webkit-box-orient: vertical;
                             overflow: hidden;">
                        {theme['description_th']}
                    </p>
                </div>
                """
                
                if st.markdown(theme_card, unsafe_allow_html=True):
                    st.session_state.current_theme = theme['id']
                    starter = generate_dynamic_story_starter(theme['id'], st.session_state.level)
                    st.session_state.theme_story_starter = starter
                    st.rerun()

# เพิ่ม CSS สำหรับ theme cards
st.markdown("""
    <style>
    .theme-header {
        text-align: center;
        margin-bottom: 20px;
    }
    
    .theme-header h3 {
        color: #1e88e5;
        font-family: 'Sarabun', sans-serif;
    }
    
    .theme-card {
        position: relative;
        overflow: hidden;
    }
    
    .theme-card:hover {
        transform: translateY(-5px);
        box-shadow: 0 5px 15px rgba(0,0,0,0.2);
    }
    
    .theme-card:before {
        content: '';
        position: absolute;
        top: 0;
        left: 0;
        width: 100%;
        height: 100%;
        background: linear-gradient(45deg, rgba(255,255,255,0.1), rgba(255,255,255,0.5));
        opacity: 0;
        transition: opacity 0.3s ease;
    }
    
    .theme-card:hover:before {
        opacity: 1;
    }
    
    .theme-card.selected {
        transform: scale(0.98);
        box-shadow: 0 2px 5px rgba(0,0,0,0.2);
    }
    
    /* Responsive adjustments */
    @media (max-width: 768px) {
        .theme-card {
            min-height: 120px;
        }
        
        .theme-card h4 {
            font-size: 0.9em;
        }
        
        .theme-card p {
            font-size: 0.8em;
        }
    }
    </style>
""", unsafe_allow_html=True)


# Add theme-specific CSS
def add_theme_css():
    st.markdown("""
        <style>
        .theme-card {
            transition: transform 0.2s;
        }
        .theme-card:hover {
            transform: translateY(-2px);
        }
        .theme-vocabulary {
            background-color: #f8f9fa;
            padding: 10px;
            border-radius: 8px;
            margin: 5px 0;
        }
        </style>
    """, unsafe_allow_html=True)

# Theme-specific vocabulary suggestions
def show_theme_vocabulary():
    if st.session_state.current_theme:
        vocab_list = get_theme_vocabulary(
            st.session_state.current_theme, 
            st.session_state.level
        )
        if vocab_list:
            st.markdown("### 📚 Theme Vocabulary")
            for word in vocab_list:
                st.markdown(f"- {word}")

def show_welcome_section():
    st.markdown("""
        <div class="welcome-header">
            <div class="thai">
            🌟 ยินดีต้อนรับสู่ JoyStory - แอพฝึกเขียนภาษาอังกฤษแสนสนุก!
            <br>
            เรียนรู้ภาษาอังกฤษผ่านการเขียนเรื่องราวด้วยตัวเอง พร้อมผู้ช่วย AI ที่จะช่วยแนะนำและให้กำลังใจ
            </div>
            <div class="eng">
            Welcome to JoyStory - Fun English Writing Adventure!
            </div>
        </div>
    """, unsafe_allow_html=True)

def show_parent_guide():
    with st.expander("📖 คำแนะนำสำหรับผู้ปกครอง | Parent's Guide"):
        st.markdown("""
            <div class="parent-guide">
                <h4>คำแนะนำในการใช้งาน</h4>
                <ul>
                    <li>แนะนำให้นั่งเขียนเรื่องราวร่วมกับน้องๆ</li>
                    <li>ช่วยอธิบายคำแนะนำและคำศัพท์ที่น้องๆ ไม่เข้าใจ</li>
                    <li>ให้กำลังใจและชื่นชมเมื่อน้องๆ เขียนได้ดี</li>
                    <li>ใช้เวลาในการเขียนแต่ละครั้งไม่เกิน 20-30 นาที</li>
                </ul>
                <p>💡 เคล็ดลับ: ให้น้องๆ พูดเรื่องราวที่อยากเขียนเป็นภาษาไทยก่อน 
                แล้วค่อยๆ ช่วยกันแปลงเป็นภาษาอังกฤษ</p>
            </div>
        """, unsafe_allow_html=True)

def generate_story_continuation(user_input: str, level: str) -> str:
    """Generate AI story continuation using ChatGPT with level-appropriate content."""
    level_context = {
        'Beginner': """
        Role: You are a teaching assistant for Thai students in grades 1-3.
        Rules:
        - Use only 1-2 VERY simple sentences
        - Use Present Simple Tense only
        - Use basic vocabulary (family, school, daily activities)
        - Each sentence should be 5-7 words maximum
        - Focus on clear, basic responses
        Example responses:
        - "The cat sits under the tree."
        - "The boy plays with his toy car."
        - "They walk to school together."
        """,
        'Intermediate': """
        Role: You are a teaching assistant for Thai students in grades 4-6.
        Rules:
        - Use exactly 2 sentences maximum
        - Can use Present or Past Tense
        - Keep each sentence under 12 words
        - Use grade-appropriate vocabulary
        - Add simple descriptions but stay concise
        Example responses:
        - "The brown cat jumped over the tall wooden fence. It landed softly in the garden."
        - "Tom carefully opened his mysterious new book. The colorful pages showed amazing magical creatures."
        """,
        'Advanced': """
        Role: You are a teaching assistant for Thai students in grades 7-9.
        Rules:
        - Use 2-3 sentences maximum (no more!)
        - Various tenses are allowed
        - No strict word limit per sentence, but keep overall response concise
        - Use more sophisticated vocabulary and sentence structures
        - Create engaging responses that encourage creative continuation
        - Focus on quality and natural flow rather than sentence length
        Example responses:
        - "Sarah discovered an ancient-looking letter hidden beneath the creaky floorboards. The yellowed paper contained a mysterious message."
        - "As the storm clouds gathered overhead, James remembered the old legend about the mountain. Lightning illuminated the winding path that led to the cave entrance."
        """
    }
    
    try:
        story_context = '\n'.join([entry['content'] for entry in st.session_state.story[-3:]])  # Only use last 3 entries
        
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": f"""You are a storytelling assistant for Thai students.
                {level_context[level]}
                CRUCIAL GUIDELINES:
                - NEVER exceed the maximum number of 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
                
                Remember: This is interactive storytelling - let the student drive the story forward."""},
                {"role": "user", "content": f"Story context (recent):\n{story_context}\nStudent's input:\n{user_input}\nProvide a brief continuation:"}
            ],
            max_tokens={
                'Beginner': 30,
                'Intermediate': 40,
                'Advanced': 50
            }[level],
            temperature=0.7,
            presence_penalty=0.6,  # Discourage repetitive responses
            frequency_penalty=0.6   # Encourage diversity in responses
        )
        
        # Additional length check and cleanup
        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 with proper punctuation
        response_text = '. '.join(sentences) + '.'
            
        return response_text
        
    except Exception as e:
        st.error(f"Error generating story continuation: {str(e)}")
        return "I'm having trouble continuing the story. Please try again."

# ฟังก์ชันสำหรับแก้ไขประโยค
def apply_correction(story_index: int, corrected_text: str):
    """Apply correction to a specific story entry."""
    if 0 <= story_index < len(st.session_state.story):
        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.success("✅ แก้ไขประโยคเรียบร้อยแล้ว!")
        
        # อัพเดทสถิติ
        if 'stats' in st.session_state:
            st.session_state.stats['corrections_made'] = \
                st.session_state.stats.get('corrections_made', 0) + 1

# เพิ่ม CSS สำหรับการแสดงผล
st.markdown("""
    <style>
    .story-message {
        padding: 10px;
        margin: 5px 0;
        border-radius: 5px;
    }
    .story-message.ai {
        background-color: #f0f7ff;
    }
    .story-message.user {
        background-color: #fff;
    }
    .story-message.corrected {
        background-color: #f0f9ff;
    }
    </style>
""", unsafe_allow_html=True)

def get_vocabulary_suggestions() -> List[str]:
    """Get contextual vocabulary suggestions with Thai translations."""
    try:
        recent_story = '\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 with their Thai translations and examples.
                Format each suggestion as:
                word (type) - คำแปล | example sentence
                Make sure words match {st.session_state.level} level."""},
                {"role": "user", "content": f"Story context:\n{recent_story}\n\nSuggest 5 relevant words with Thai translations:"}
            ],
            max_tokens=200,
            temperature=0.8
        )
        return response.choices[0].message.content.split('\n')
    except Exception as e:
        st.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"]

# In the main UI section, update how vocabulary suggestions are displayed:
    if st.button("Get Vocabulary Ideas"):
        vocab_suggestions = get_vocabulary_suggestions()
        st.markdown("#### 📚 Suggested Words")
        for word in vocab_suggestions:
            st.markdown(f"• *{word}*")

# And update how feedback is displayed to be more concise:
    if st.session_state.feedback:
        st.markdown("""
            <div style='background-color: #f0f2f6; padding: 8px; border-radius: 4px; margin-bottom: 10px;'>
                📝 <i>{}</i>
            </div>
            """.format(st.session_state.feedback), unsafe_allow_html=True)

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

def provide_feedback(text: str, level: str) -> Dict[str, str]:
    """Provide feedback and corrected sentence."""
    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": f"""You are a Thai English teacher helping {level} students.
                Your task is to review the student's sentence and provide feedback.
                
                Return your response in this EXACT format only (must be valid JSON):
                {{
                    "feedback": "(ข้อเสนอแนะภาษาไทย)",
                    "corrected": "(ประโยคภาษาอังกฤษที่ถูกต้อง)",
                    "has_errors": true/false
                }}
                
                Example 1 - with error:
                {{
                    "feedback": "คำว่า 'go' เมื่อพูดถึงเหตุการณ์ในอดีต ต้องเปลี่ยนเป็น 'went'",
                    "corrected": "I went to school yesterday.",
                    "has_errors": true
                }}
                
                Example 2 - no error:
                {{
                    "feedback": "เขียนได้ถูกต้องแล้วค่ะ ประโยคสื่อความหมายได้ดี",
                    "corrected": "The cat is sleeping.",
                    "has_errors": false
                }}"""},
                {"role": "user", "content": f"Review this sentence and provide feedback in the specified JSON format: {text}"}
            ],
            max_tokens=200,
            temperature=0.3
        )
        
        # Get the response text
        response_text = response.choices[0].message.content.strip()
        
        try:
            # แสดง debug info เพื่อตรวจสอบ response
            st.write("Debug - Response received:", response_text)
            
            # Parse JSON response
            feedback_data = json.loads(response_text)
            
            # Validate required fields
            if all(key in feedback_data for key in ['feedback', 'corrected', 'has_errors']):
                return feedback_data
            else:
                raise ValueError("Missing required fields in response")
                
        except json.JSONDecodeError as json_err:
            st.error(f"Debug - JSON Error: {str(json_err)}")
            raise
            
    except Exception as e:
        st.error(f"Debug - Error: {str(e)}")
        return {
            "feedback": "⚠️ ระบบไม่สามารถวิเคราะห์ประโยคได้ กรุณาลองใหม่อีกครั้ง",
            "corrected": text,
            "has_errors": False
        }

def update_points(is_correct_first_try: bool):
    """อัพเดตคะแนนตามผลการเขียน"""
    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
        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
    )

def show_story():
    """Display the story with proper formatting."""
    story_display = st.container()
    
    with story_display:
        if not st.session_state.story:
            st.info("เลือกธีมเรื่องราวที่ต้องการเพื่อเริ่มต้นการผจญภัย!")
        else:
            # แสดงข้อความทั้งหมดตามลำดับ
            for entry in st.session_state.story:
                if entry['role'] == 'AI':
                    if entry.get('is_starter'):
                        # แสดงประโยคเริ่มต้นพิเศษ
                        st.markdown(f"""
                            <div style="background-color: #f0f7ff; 
                                      padding: 15px; 
                                      border-radius: 8px; 
                                      margin: 5px 0;
                                      border-left: 4px solid #1e88e5;">
                                <p style="color: #1e88e5; margin-bottom: 5px;">🎬 เริ่มเรื่อง:</p>
                                <p style="color: #666;">{entry.get('thai_content', '')}</p>
                                <p>{entry['content']}</p>
                            </div>
                        """, unsafe_allow_html=True)
                    else:
                        # แสดงข้อความ AI ปกติ
                        st.markdown(f"""
                            <div style="padding: 10px; margin: 5px 0;">
                                <p>🤖 AI: {entry['content']}</p>
                            </div>
                        """, unsafe_allow_html=True)
                elif entry['role'] == 'You':
                    # กำหนดสถานะและสไตล์การแสดงผล
                    status_icon = "✅ " if entry.get('is_correct') else "✍️ "
                    bg_color = "#f0f9ff" if entry.get('is_corrected') else "white"
                    
                    st.markdown(f"""
                        <div style="padding: 10px; 
                                  margin: 5px 0; 
                                  background-color: {bg_color}; 
                                  border-radius: 5px;">
                            <p>👤 You: {status_icon}{entry['content']}</p>
                        </div>
                    """, unsafe_allow_html=True)

# Add CSS for consistent styling
st.markdown("""
    <style>
    .story-message {
        padding: 10px;
        margin: 5px 0;
        border-radius: 5px;
    }
    .story-message.ai {
        background-color: #f0f7ff;
    }
    .story-message.user {
        background-color: #fff;
    }
    .story-message.corrected {
        background-color: #f0f9ff;
    }
    </style>
""", unsafe_allow_html=True)

# แสดงผลคะแนนและความสำเร็จ
def show_achievements():
    """แสดงความสำเร็จและสถิติ"""
    with st.container():
        # 1. แสดงคะแนนรวมและ Streak
        st.markdown(f"""
            <div style='background-color: #f0f8ff; padding: 15px; border-radius: 10px; margin-bottom: 15px;'>
                <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)

        # 2. แสดงสถิติการเขียน
        st.markdown("""
            <div style='margin-bottom: 10px;'>
                <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="จำนวนคำศัพท์ที่ไม่ซ้ำกันที่คุณได้ใช้"
            )

        # 3. แสดงความสำเร็จ
        st.markdown("""
            <div style='margin: 20px 0 10px 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("ยังไม่มีความสำเร็จ - เขียนต่อไปเพื่อปลดล็อกรางวัล!")

        # 4. แสดงความสำเร็จที่ยังไม่ได้รับ (ตัวเลือก)
        if not st.session_state.achievements:  # ถ้ายังไม่มีความสำเร็จ
            st.markdown("""
                <div style='margin-top: 15px; font-size: 0.9em; color: #666;'>
                    เป้าหมายที่จะได้รับความสำเร็จ:
                </div>
            """, unsafe_allow_html=True)
            st.markdown("""
            - 🌟 นักเขียนไร้ที่ติ: เขียนถูกต้อง 5 ประโยคติดต่อกัน
            - 📚 ราชาคำศัพท์: ใช้คำศัพท์ไม่ซ้ำกัน 50 คำ
            - 📖 นักแต่งนิทาน: เขียนเรื่องยาว 10 ประโยค
            - 👑 ราชาความแม่นยำ: มีอัตราความถูกต้อง 80% ขึ้นไป
            """)

def update_achievements():
    """ตรวจสอบและอัพเดตความสำเร็จ"""
    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("🎉 ได้รับความสำเร็จใหม่: ราชาความแม่นยำ!")

def create_story_pdf():
    """สร้าง PDF จากเรื่องราวที่เขียน"""
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(
        buffer,
        pagesize=A4,
        rightMargin=72,
        leftMargin=72,
        topMargin=72,
        bottomMargin=72
    )

    # สร้าง styles สำหรับ PDF
    styles = getSampleStyleSheet()
    title_style = ParagraphStyle(
        'CustomTitle',
        parent=styles['Heading1'],
        fontSize=24,
        spaceAfter=30,
        alignment=1  # center
    )
    
    story_style = ParagraphStyle(
        'StoryText',
        parent=styles['Normal'],
        fontSize=12,
        spaceBefore=12,
        spaceAfter=12,
        leading=16
    )
    
    stats_style = ParagraphStyle(
        'Stats',
        parent=styles['Normal'],
        fontSize=10,
        textColor=colors.gray,
        alignment=1
    )

    # สร้างเนื้อหา PDF
    elements = []
    
    # หน้าปก
    elements.append(Paragraph("My English Story Adventure", title_style))
    elements.append(Paragraph(
        f"Created by: {st.session_state.get('player_name', 'Young Writer')}<br/>"
        f"Date: {datetime.now().strftime('%B %d, %Y')}<br/>"
        f"Level: {st.session_state.level}",
        stats_style
    ))
    elements.append(Spacer(1, 30))

    # เพิ่มสถิติการเขียน
    stats_text = (
        f"Total Points: {st.session_state.points['total']}<br/>"
        f"Perfect Sentences: {st.session_state.points['perfect_sentences']}<br/>"
        f"Accuracy Rate: {st.session_state.stats['accuracy_rate']:.1f}%<br/>"
        f"Unique Words Used: {len(st.session_state.stats['vocabulary_used'])}"
    )
    elements.append(Paragraph(stats_text, stats_style))
    elements.append(Spacer(1, 30))

    # เพิ่มเนื้อเรื่อง
    elements.append(Paragraph("The Story", styles['Heading2']))
    for entry in st.session_state.story:
        if entry['role'] == 'You':
            text = f"👤 {entry['content']}"
            if entry.get('is_correct'):
                style = ParagraphStyle(
                    'Correct',
                    parent=story_style,
                    textColor=colors.darkgreen
                )
            else:
                style = story_style
        else:  # AI response
            text = f"🤖 {entry['content']}"
            style = ParagraphStyle(
                'AI',
                parent=story_style,
                textColor=colors.navy
            )
        elements.append(Paragraph(text, style))

    # เพิ่ม achievements ที่ได้รับ
    if hasattr(st.session_state, 'achievements') and st.session_state.achievements:
        elements.append(Spacer(1, 20))
        elements.append(Paragraph("Achievements Earned", styles['Heading2']))
        for achievement in st.session_state.achievements:
            elements.append(Paragraph(f"🏆 {achievement}", stats_style))

    # สร้าง PDF
    doc.build(elements)
    pdf = buffer.getvalue()
    buffer.close()
    return pdf

# เพิ่มฟังก์ชันโหลดความก้าวหน้า
def load_progress(uploaded_file):
    """โหลดความก้าวหน้าจากไฟล์ JSON"""
    try:
        data = json.loads(uploaded_file.getvalue())
        st.session_state.level = data['level']
        st.session_state.story = data['story']
        st.session_state.achievements = data['achievements']
        st.session_state.points = data['points']
        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'])  # แปลง list กลับเป็น set
        }
        st.success("โหลดความก้าวหน้าเรียบร้อย!")
        st.rerun()
    except Exception as e:
        st.error(f"เกิดข้อผิดพลาดในการโหลดไฟล์: {str(e)}")

def reset_story():
    """Reset the story and related state variables."""
    st.session_state.story = []
    st.session_state.feedback = None
    st.session_state.unique_words = set()
    st.session_state.total_words = 0
    st.session_state.achievements = []  # เปลี่ยนจาก badges เป็น achievements
    st.session_state.should_reset = False
    
    # เพิ่มการ 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()
    }

# Handle story reset if needed
if st.session_state.should_reset:
    reset_story()

# Main UI Layout
# In your main UI section, after init_session_state()
init_theme_state()
add_theme_css()
st.markdown("# 📖 JoyStory")
show_welcome_section()
show_parent_guide()  # Add parent guide right after welcome section

# After showing the welcome section and before the main story area
if not st.session_state.current_theme:
    show_theme_selection()
else:
    # Show current theme and option to change
    if st.button("🔄 Change Theme"):
        st.session_state.current_theme = None
        st.rerun()
    
    # Show theme-specific story starter if no story has been started
    if not st.session_state.story and st.session_state.theme_story_starter:
        st.markdown(f"""
            <div class="story-starter">
                <p class="thai">{st.session_state.theme_story_starter['th']}</p>
                <p class="eng">{st.session_state.theme_story_starter['en']}</p>
            </div>
        """, unsafe_allow_html=True)

# Sidebar for settings
with st.sidebar:
    st.markdown("### 📂 โหลดความก้าวหน้า")
    uploaded_file = st.file_uploader(
        "เลือกไฟล์ .json",
        type=['json'],
        help="เลือกไฟล์ความก้าวหน้าที่บันทึกไว้"
    )
    if uploaded_file:
        if st.button("โหลดความก้าวหน้า"):
            load_progress(uploaded_file)
            
    st.markdown("""
        <div class="thai">
        🎯 เลือกระดับการเรียนรู้
        </div>
    """, unsafe_allow_html=True)

    level = st.radio(
        "ระดับการเรียน",  # เพิ่ม label
        options=list(level_options.keys()),
        format_func=lambda x: level_options[x]['thai_name'],
        label_visibility="collapsed"  # ซ่อน label แต่ยังคงมีสำหรับ accessibility
    )
    
    # แสดงคำอธิบายระดับ
    st.markdown(f"""
        <div class="level-info thai">
            🎓 {level_options[level]['description']}
            <br>📚 เหมาะสำหรับอายุ: {level_options[level]['age_range']}
        </div>
    """, unsafe_allow_html=True)

    st.session_state.level = level
    
    if st.button("เริ่มเรื่องใหม่ | Start New Story"):
        st.session_state.should_reset = True
        st.rerun()
        

# Main content area
col1, col2 = st.columns([3, 1])

with col1:
    # แสดงส่วนบนของหน้า
    st.markdown("""
        <div class="thai-eng">
            <div class="thai">📖 เรื่องราวของคุณ</div>
            <div class="eng">Your Story</div>
        </div>
    """, unsafe_allow_html=True)

    # ตรวจสอบว่ามีธีมที่เลือกแล้วหรือไม่
    if not st.session_state.current_theme:
        # ถ้ายังไม่มีธีม แสดงส่วนเลือกธีม
        show_theme_selection()
    else:
        # ถ้ามีธีมแล้ว แสดงปุ่ม Change Theme
        if st.button("🔄 Change Theme", key="change_theme_button", 
                    help="คลิกเพื่อเปลี่ยนธีมเรื่องราว"):
            # รีเซ็ตค่าเมื่อต้องการเปลี่ยนธีม
            st.session_state.current_theme = None
            st.session_state.theme_story_starter = None
            st.session_state.story = []
            st.rerun()
        
        # แสดงเนื้อเรื่อง
        show_story()

    # User Input Box (แสดงเฉพาะเมื่อมีธีมแล้ว)
    if st.session_state.current_theme:
        st.markdown("""
            <div class="thai-eng">
                <div class="thai">✏️ ถึงตาคุณแล้ว</div>
                <div class="eng">Your Turn</div>
            </div>
        """, unsafe_allow_html=True)
        
        # Text input with callback
        st.text_area(
            "เขียนต่อจากเรื่องราว | Continue the story:",
            height=100,
            key="text_input",
            label_visibility="collapsed" 
        )
        
        # Submit button with callback
        st.button(
            "ส่งคำตอบ | Submit",
            on_click=submit_story
        )

with col2:
    # 1. Feedback Section (Most Important)
    if st.session_state.feedback:
        st.markdown("""
            <div class="feedback-box">
                <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'):
            st.markdown(f"""
                <div style='background-color: #f0f2f6; 
                       padding: 15px; 
                       border-radius: 8px; 
                       border-left: 4px solid #FF9800;
                       margin: 5px 0;
                       font-family: "Sarabun", sans-serif;'>
                    <p style='color: #1e88e5; margin-bottom: 10px;'>
                        {feedback_data['feedback']}
                    </p>
                    <p style='color: #666; font-size: 0.9em;'>
                        ประโยคที่ถูกต้อง:<br/>
                        <span style='color: #4CAF50; font-weight: bold;'>
                            {feedback_data['corrected']}
                        </span>
                    </p>
                </div>
            """, unsafe_allow_html=True)
            
            # แสดงปุ่มแก้ไข
            if st.button("✍️ แก้ไขประโยคให้ถูกต้อง", key="correct_button"):
                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:
            st.markdown(f"""
                <div style='background-color: #f0f2f6; 
                       padding: 15px; 
                       border-radius: 8px; 
                       border-left: 4px solid #4CAF50;
                       margin: 5px 0;
                       font-family: "Sarabun", sans-serif;'>
                    <p style='color: #1e88e5;'>
                        {feedback_data.get('feedback', '✨ เขียนได้ถูกต้องแล้วค่ะ!')}
                    </p>
                </div>
            """, unsafe_allow_html=True)

    # 2. Writing Tools (Expandable)
    with st.expander("✨ เครื่องมือช่วยเขียน | Writing Tools"):
        if st.button("🎯 ขอคำใบ้ | Get Creative Prompt"):
            prompt = get_creative_prompt()
            st.markdown(f"""
                <div class="thai-eng">
                    <div class="thai">💭 {prompt['thai']}</div>
                    <div class="eng">💭 {prompt['eng']}</div>
                </div>
            """, unsafe_allow_html=True)
        
        # Vocabulary ย้ายมาอยู่ในนี้
        if st.button("📚 คำศัพท์แนะนำ | Vocabulary Tips"):
            vocab_suggestions = get_vocabulary_suggestions()
            st.markdown("#### 📚 คำศัพท์น่ารู้ | Useful Words")
            for word in vocab_suggestions:
                st.markdown(f"• {word}")

    # 3. Achievements (Expandable)
    with st.expander("🏆 ความสำเร็จ | Achievements"):
        show_achievements()

    # 4. Save Options (At the bottom)
    st.markdown("### 💾 บันทึกเรื่องราว")
    if st.session_state.story:
        col1, col2 = st.columns(2)
        
        with col1:
            # บันทึกเป็น PDF
            pdf = create_story_pdf()
            st.download_button(
                label="📑 PDF",
                data=pdf,
                file_name=f"my_story_{datetime.now().strftime('%Y%m%d')}.pdf",
                mime="application/pdf"
            )
        
        with col2:
            # บันทึกความก้าวหน้า (แก้ไขส่วนนี้)
            story_data = {
                'level': st.session_state.level,
                'date': datetime.now().isoformat(),
                'story': st.session_state.story,
                'achievements': st.session_state.achievements,
                'points': st.session_state.points,
                'stats': {
                    'total_sentences': st.session_state.stats['total_sentences'],
                    'correct_first_try': st.session_state.stats['correct_first_try'],
                    'accuracy_rate': st.session_state.stats['accuracy_rate'],
                    'vocabulary_used': list(st.session_state.stats['vocabulary_used'])  # แปลง set เป็น list
                }
            }
            
            # แปลงข้อมูลทั้งหมดที่อาจมี set
            serializable_data = convert_sets_to_lists(story_data)
            
            st.download_button(
                label="💾 Save",
                data=json.dumps(serializable_data, indent=2),
                file_name=f"story_progress_{datetime.now().strftime('%Y%m%d')}.json",
                mime="application/json"
            )


# Add new CSS for improved layout
st.markdown("""
    <style>
    .feedback-box {
        margin-bottom: 20px;
    }
    .expander-box {
        margin-top: 10px;
    }
    .save-options {
        position: sticky;
        bottom: 0;
        background: white;
        padding: 10px 0;
        border-top: 1px solid #eee;
    }
    </style>
""", unsafe_allow_html=True)