import streamlit as st
import json
import datetime
from openai import OpenAI
from typing import Dict, List, Set

# 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');
    
    .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;
    }
    </style>
""", unsafe_allow_html=True)

# Initialize session state variables
def init_session_state():
    if 'story' not in st.session_state:
        st.session_state.story = []
    if 'feedback' not in st.session_state:
        st.session_state.feedback = None
    if 'corrections' not in st.session_state:
        st.session_state.corrections = {}  # เก็บประวัติการแก้ไข
    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 'badges' not in st.session_state:
        st.session_state.badges = []
    if 'should_reset' not in st.session_state:
        st.session_state.should_reset = False

init_session_state()

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

# Callback function for submit button
def submit_story():
    if st.session_state.text_input.strip():
        user_text = st.session_state.text_input
        
        # เพิ่มประโยคของผู้ใช้
        story_index = len(st.session_state.story)
        st.session_state.story.append({
            "role": "You", 
            "content": user_text,
            "is_corrected": False
        })
        
        try:
            # รับ feedback และประโยคที่ถูกต้อง
            feedback_data = provide_feedback(user_text, st.session_state.level)
            st.session_state.feedback = feedback_data
            
            # ถ้ามีข้อผิดพลาด แสดงคำแนะนำ
            if feedback_data['has_errors']:
                st.markdown(f"""
                    <div style='background-color: #fff3e0; padding: 10px; border-radius: 5px; margin: 10px 0;'>
                        <p style='color: #ff6d00;'>🎯 คำแนะนำ:</p>
                        <p>{feedback_data['feedback']}</p>
                    </div>
                """, unsafe_allow_html=True)
            
            # Generate AI continuation using corrected text if there were errors
            ai_response = generate_story_continuation(
                feedback_data['corrected'] if feedback_data['has_errors'] else user_text, 
                st.session_state.level
            )
            st.session_state.story.append({"role": "AI", "content": ai_response})
            
        except Exception as e:
            st.error("เกิดข้อผิดพลาดในการวิเคราะห์ประโยค กรุณาลองใหม่อีกครั้ง")
            # Log error for debugging
            st.error(f"Debug - Error in submit_story: {str(e)}")
        
        # Clear input
        st.session_state.text_input = ""

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.corrections[story_index] = {
            'original': original_text,
            'corrected': corrected_text,
            'timestamp': datetime.datetime.now().isoformat()
        }
        # แก้ไขประโยคในเรื่อง
        st.session_state.story[story_index]['content'] = corrected_text
        st.session_state.story[story_index]['is_corrected'] = 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:
            # Try to parse the JSON response
            feedback_data = json.loads(response_text)
            
            # Validate the 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 response")
            
            return feedback_data
            
        except json.JSONDecodeError as json_err:
            # If JSON parsing fails, try to extract information manually
            st.error(f"Debug - Response received: {response_text}")
            
            # Return a safe fallback response
            return {
                "feedback": "⚠️ คำแนะนำ: ระบบไม่สามารถวิเคราะห์ประโยคได้อย่างถูกต้อง กรุณาลองใหม่อีกครั้ง",
                "corrected": text,
                "has_errors": False
            }
            
    except Exception as e:
        st.error(f"Debug - Error type: {type(e).__name__}")
        st.error(f"Debug - Error message: {str(e)}")
        
        # Return a safe fallback response
        return {
            "feedback": "⚠️ ขออภัย ระบบไม่สามารถวิเคราะห์ประโยคได้ในขณะนี้",
            "corrected": text,
            "has_errors": False
        }

def update_achievements(text: str):
    """Update user achievements based on their writing."""
    words = set(text.lower().split())
    st.session_state.unique_words.update(words)
    st.session_state.total_words += len(words)
    
    achievements = {
        '50 Words': len(st.session_state.unique_words) >= 50,
        '100 Words': len(st.session_state.unique_words) >= 100,
        'Story Master': len(st.session_state.story) >= 10,
    }
    
    for badge, condition in achievements.items():
        if condition and badge not in st.session_state.badges:
            st.session_state.badges.append(badge)
            st.success(f"🏆 Achievement Unlocked: {badge}!")

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.badges = []
    st.session_state.should_reset = False

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

# Main UI Layout
st.markdown("# 📖 JoyStory")
show_welcome_section()
show_parent_guide()  # Add parent guide right after welcome section

# Sidebar for settings
with st.sidebar:
    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'],
        help="เลือกระดับที่เหมาะสมกับความสามารถของน้องๆ"
    )
    
    # แสดงคำอธิบายระดับ
    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:
    # Story Display Box
    st.markdown("""
        <div class="thai-eng">
            <div class="thai">📖 เรื่องราวของคุณ</div>
            <div class="eng">Your Story</div>
        </div>
    """, unsafe_allow_html=True)
    
    story_display = st.container()

    with story_display:
        if not st.session_state.story:
            st.info("เริ่มต้นผจญภัยด้วยการเขียนประโยคแรกกันเลย! | Start your adventure by writing the first sentence!")
        else:
            for idx, entry in enumerate(st.session_state.story):
                if entry['role'] == 'You':
                    col_text, col_button = st.columns([4, 1])
                    with col_text:
                        prefix = "👤 You (แก้ไขแล้ว):" if entry.get('is_corrected') else "👤 You:"
                        st.write(f"{prefix} {entry['content']}")
                    
                    # ถ้ามี feedback และมีข้อผิด แสดงปุ่มแก้ไข
                    if (idx in st.session_state.corrections and 
                        not entry.get('is_corrected') and 
                        st.session_state.feedback and 
                        st.session_state.feedback.get('has_errors')):
                        with col_button:
                            if st.button(f"✍️ แก้ไขประโยค", key=f"correct_{idx}"):
                                apply_correction(idx, st.session_state.feedback['corrected'])
                                st.rerun()
                elif entry['role'] == 'AI':
                    st.write("🤖 AI:", entry['content'])
    
    # User Input Box
    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"
    )
    
    # Submit button with callback
    st.button(
        "ส่งคำตอบ | Submit",
        on_click=submit_story
    )

with col2:
    # Feedback Display
    if st.session_state.feedback:
        st.markdown("""
            <div class="thai-eng">
                <div class="thai">📝 คำแนะนำ</div>
                <div class="eng">Writing Feedback</div>
            </div>
        """, unsafe_allow_html=True)
        st.markdown(f"*{st.session_state.feedback}*")
    
    # Help and Suggestions Box
    st.markdown("""
        <div class="thai-eng">
            <div class="thai">✨ เครื่องมือช่วยเขียน</div>
            <div class="eng">Writing Tools</div>
        </div>
    """, unsafe_allow_html=True)
    
    if st.button("ดูคำศัพท์แนะนำ | Get Vocabulary Ideas"):
        vocab_suggestions = get_vocabulary_suggestions()
        st.markdown("#### 📚 คำศัพท์น่ารู้ | Useful Words")
        for word in vocab_suggestions:
            st.markdown(f"• {word}")
    
    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)
    
    # Achievements Box
    st.markdown("""
        <div class="thai-eng">
            <div class="thai">🏆 ความสำเร็จ</div>
            <div class="eng">Achievements</div>
        </div>
    """, unsafe_allow_html=True)
    
    if st.session_state.badges:
        for badge in st.session_state.badges:
            st.success(f"🏆 {badge}")
    else:
        st.write("เขียนต่อไปเพื่อรับรางวัล | Keep writing to earn badges!")
    
    # Save Story Button
    if st.session_state.story:
        if st.button("บันทึกเรื่องราว | Save Story"):
            story_data = {
                'level': st.session_state.level,
                'date': datetime.datetime.now().isoformat(),
                'story': st.session_state.story,
                'achievements': st.session_state.badges,
                'total_words': st.session_state.total_words,
                'unique_words': list(st.session_state.unique_words)
            }
            st.download_button(
                label="ดาวน์โหลดเรื่องราว | Download Story",
                data=json.dumps(story_data, indent=2),
                file_name='joystory.json',
                mime='application/json'
            )