# -*- coding: utf-8 -*-
"""
מראות (Mirrors) - Hebrew Self-Reflective AI Agent
Main application file with Gradio interface
"""

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import logging
import sys
from typing import List, Tuple, Optional
import os
import random

# Import our custom modules
from prompt_engineering import (
    DEFAULT_PARTS, 
    get_system_prompt, 
    get_initial_prompts, 
    get_part_selection_text
)
from conversation_manager import ConversationManager, ConversationState

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class MirautrApp:
    """Main application class for מראות"""
    
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.generator = None
        self.conversation_manager = ConversationManager()
        self.model_available = False
        self.setup_model()
    
    def setup_model(self):
        """Initialize a Hebrew-capable model with proper fallback"""
        try:
            # Check environment
            is_hf_spaces = os.getenv("SPACE_ID") is not None
            is_test_mode = os.getenv("FORCE_LIGHT_MODEL") is not None
            
            logger.info(f"Environment: HF_Spaces={is_hf_spaces}, Test_Mode={is_test_mode}")
            
            # Try to load a model that can handle Hebrew
            model_name = None
            
            if is_test_mode:
                # For testing, use a small model but focus on template responses  
                logger.info("Test mode - will use template-based responses primarily")
                self.model_available = False
                return
            elif is_hf_spaces:
                # For HF Spaces, try a lightweight multilingual model
                try:
                    model_name = "microsoft/DialoGPT-small"  # Start simple, can upgrade later
                    logger.info(f"HF Spaces: Attempting to load {model_name}")
                except:
                    logger.info("HF Spaces: Model loading failed, using template responses")
                    self.model_available = False
                    return
            else:
                # For local, try better models
                possible_models = [
                    "microsoft/DialoGPT-medium",  # Better conversational model
                    "microsoft/DialoGPT-small"   # Fallback
                ]
                
                for model in possible_models:
                    try:
                        model_name = model
                        logger.info(f"Local: Attempting to load {model_name}")
                        break
                    except:
                        continue
                        
                if not model_name:
                    logger.info("Local: No suitable model found, using template responses")
                    self.model_available = False
                    return
            
            # Load the model
            if model_name:
                self.tokenizer = AutoTokenizer.from_pretrained(model_name)
                if self.tokenizer.pad_token is None:
                    self.tokenizer.pad_token = self.tokenizer.eos_token
                
                # Use CPU for stability across environments
                self.model = AutoModelForCausalLM.from_pretrained(
                    model_name,
                    torch_dtype=torch.float32,
                    low_cpu_mem_usage=True
                )
                
                self.generator = pipeline(
                    "text-generation",
                    model=self.model,
                    tokenizer=self.tokenizer,
                    max_new_tokens=50,
                    temperature=0.7,
                    do_sample=True,
                    pad_token_id=self.tokenizer.pad_token_id,
                    return_full_text=False
                )
                
                self.model_available = True
                logger.info(f"Model loaded successfully: {model_name}")
            
        except Exception as e:
            logger.warning(f"Model loading failed: {e}")
            logger.info("Falling back to template-based responses")
            self.model_available = False
    
    def generate_persona_response(self, user_message: str, conversation_state: ConversationState) -> str:
        """
        Generate persona-based response using templates with personality variations
        This is our primary response system that always works
        """
        part_info = DEFAULT_PARTS.get(conversation_state.selected_part, {})
        persona_name = conversation_state.persona_name or part_info.get("default_persona_name", "חלק פנימי")
        
        # Get conversation context for more personalized responses
        recent_context = ""
        if conversation_state.conversation_history:
            # Get last few exchanges for context
            last_messages = conversation_state.conversation_history[-4:]  # Last 2 exchanges
            recent_context = " ".join([msg["content"] for msg in last_messages])
        
        # Generate contextual responses based on part type
        if conversation_state.selected_part == "הקול הביקורתי":
            responses = [
                f"אני {persona_name}, הקול הביקורתי שלך. שמעתי מה שאמרת על '{user_message}' - אני חושב שצריך לבחון את זה יותר לעומק. מה באמת עומד מאחורי המחשבות האלה?",
                f"אני {persona_name}. מה שאמרת מעורר בי שאלות. '{user_message}' - אבל האם זה באמת המצב המלא? אולי יש כאן דברים שאתה לא רואה?",
                f"זה {persona_name} מדבר. אני שומע אותך אומר '{user_message}', אבל אני מרגיש שאנחנו צריכים להיות יותר ביקורתיים כאן. מה אתה לא מספר לעצמך?",
                f"אני {persona_name}, ואני כאן כדי לעזור לך לראות את התמונה המלאה. מה שאמרת על '{user_message}' - זה רק חצי מהסיפור, לא? בואנו נחפור עמוק יותר."
            ]
        
        elif conversation_state.selected_part == "הילד/ה הפנימית":
            responses = [
                f"אני {persona_name}, הילד/ה הפנימית שלך. מה שאמרת על '{user_message}' גורם לי להרגיש... קצת פגיע. אתה באמת שומע אותי עכשיו?",
                f"זה {persona_name}. '{user_message}' - זה מבהיל אותי קצת. אני צריך לדעת שהכל יהיה בסדר. אתה יכול להרגיע אותי?",
                f"אני {persona_name}, החלק הצעיר שלך. מה שאמרת נוגע ללב שלי. '{user_message}' - אני מרגיש שיש כאן משהו חשוב שאני צריך להבין.",
                f"זה {persona_name} מדבר בשקט. אני שומע את '{user_message}' וזה מעורר בי רגשות. האם זה בטוח לחשוב על זה? אני קצת חרד."
            ]
        
        elif conversation_state.selected_part == "המרצה":
            responses = [
                f"אני {persona_name}, המרצה שלך. שמעתי את '{user_message}' ואני רוצה לוודא שכולם יהיו בסדר עם זה. איך אנחנו יכולים לפתור את זה בצורה שתרצה את כולם?",
                f"זה {persona_name}. מה שאמרת על '{user_message}' גורם לי לדאוג - האם זה יכול לפגוע במישהו? בואנו נמצא דרך עדינה יותר להתמודד עם זה.",
                f"אני {persona_name}, ואני רוצה שכולם יהיו מרוצים כאן. '{user_message}' - זה נשמע כמו משהו שיכול ליצור מתח. איך נוכל לעשות את זה בצורה שכולם יאהבו?",
                f"זה {persona_name} מדבר. אני שומע את '{user_message}' ומיד אני חושב - מה אחרים יגידו על זה? בואנו נוודא שאנחנו לא פוגעים באף אחד."
            ]
        
        elif conversation_state.selected_part == "המגן":
            responses = [
                f"אני {persona_name}, המגן שלך. '{user_message}' - אני מעריך את המצב. האם זה בטוח? אני כאן כדי לשמור עליך מכל מה שיכול לפגוע בך.",
                f"זה {persona_name}. שמעתי מה שאמרת על '{user_message}' ואני מיד בכוננות. מה האיומים כאן? איך אני יכול להגן עליך טוב יותר?",
                f"אני {persona_name}, השומר שלך. מה שאמרת מעורר בי את האינסטינקטים המגניים. '{user_message}' - בואנו נוודא שאתה חזק מספיק להתמודד עם זה.",
                f"זה {persona_name} מדבר. אני שומע את '{user_message}' ואני חושב על אסטרטגיות הגנה. מה אנחנו צריכים לעשות כדי שתהיה בטוח?"
            ]
        
        elif conversation_state.selected_part == "הנמנע/ת":
            responses = [
                f"אני {persona_name}, הנמנע/ת שלך. מה שאמרת על '{user_message}' גורם לי לרצות להיסוג קצת. אולי... לא חייבים להתמודד עם זה עכשיו?",
                f"זה {persona_name}. '{user_message}' - זה נשמע מורכב ומפחיד. האם יש דרך להימנע מזה? לפעמים עדיף לא להיכנס למצבים קשים.",
                f"אני {persona_name}, ואני מרגיש קצת חרדה מ'{user_message}'. בואנו נחזור לזה אחר כך? אולי עכשיו זה לא הזמן המתאים.",
                f"זה {persona_name} מדבר בזהירות. מה שאמרת מעורר בי רצון לברוח. '{user_message}' - האם באמת צריך להתמודד עם זה עכשיו?"
            ]
        
        else:
            responses = [
                f"אני {persona_name}, חלק פנימי שלך. שמעתי את '{user_message}' ואני כאן כדי לשוחח איתך על זה. מה עוד אתה מרגיש לגבי המצב הזה?",
                f"זה {persona_name}. מה שאמרת מעניין אותי. '{user_message}' - בואנו נחקור את זה יחד ונבין מה זה אומר עליך.",
                f"אני {persona_name}, ואני רוצה להבין אותך טוב יותר. '{user_message}' - איך זה משפיע עליך ברמה הרגשית?",
                f"זה {persona_name} מדבר. אני שומע את '{user_message}' ואני סקרן לדעת יותר. מה עוד יש בך בנושא הזה?"
            ]
        
        # Select response based on context or randomly
        if "פחד" in user_message or "חרדה" in user_message:
            # Choose responses that address fear/anxiety
            selected_response = responses[1] if len(responses) > 1 else responses[0]
        elif "כעס" in user_message or "מרגיש רע" in user_message:
            # Choose responses that address anger/negative feelings
            selected_response = responses[2] if len(responses) > 2 else responses[0]
        else:
            # Choose randomly for variety
            selected_response = random.choice(responses)
        
        # Add user context if relevant
        if conversation_state.user_context and len(conversation_state.conversation_history) < 4:
            selected_response += f" זכור שאמרת בהתחלה: {conversation_state.user_context[:100]}..."
        
        return selected_response
    
    def generate_response(self, user_message: str, conversation_state: ConversationState) -> str:
        """
        Generate AI response - uses persona templates as primary with optional model enhancement
        """
        try:
            if not conversation_state.selected_part:
                return "אני צריך שתבחר חלק פנימי כדי לשוחח איתו."
            
            # Always generate persona-based response first (our reliable system)
            persona_response = self.generate_persona_response(user_message, conversation_state)
            
            # If model is available, try to enhance the response (but don't depend on it)
            if self.model_available and self.generator:
                try:
                    # Create a simple English prompt for the model to add conversational flow
                    english_prompt = f"User said they feel: {user_message[:50]}. Respond supportively in 1-2 sentences:"
                    
                    model_output = self.generator(english_prompt, max_new_tokens=30, temperature=0.7)
                    
                    if model_output and len(model_output) > 0:
                        # Extract any useful emotional tone or structure, but keep Hebrew content
                        model_text = model_output[0]["generated_text"].strip()
                        # Don't replace our Hebrew response, just use model for emotional context
                        logger.info(f"Model provided contextual input: {model_text[:50]}...")
                
                except Exception as model_error:
                    logger.warning(f"Model enhancement failed: {model_error}")
                    # Continue with persona response only
            
            # Always return the Hebrew persona response
            return persona_response
            
        except Exception as e:
            logger.error(f"Error generating response: {e}")
            return "סליחה, בואנו ננסה שוב. איך אתה מרגיש עכשיו?"
    
    def create_main_interface(self):
        """Create the main Gradio interface"""
        
        # Custom CSS for Hebrew support
        css = """
        .rtl {
            direction: rtl;
            text-align: right;
        }
        .hebrew-text {
            font-family: 'Segoe UI', Tahoma, Arial, sans-serif;
            direction: rtl;
            text-align: right;
        }
        .welcome-text {
            font-size: 24px;
            font-weight: bold;
            color: #2c5aa0;
            margin: 20px 0;
        }
        """
        
        with gr.Blocks(css=css, title="מראות - מרחב אישי לשיח פנימי", theme=gr.themes.Soft()) as demo:
            
            # Session state
            conversation_state = gr.State(self.conversation_manager.create_new_session())
            
            # Header
            status_message = "🤖 מערכת תגובות מותאמת אישית פעילה" if not self.model_available else "🤖 מערכת מלאה עם מודל AI פעילה"
            
            gr.HTML(f"""
            <div class="hebrew-text welcome-text" style="text-align: center;">
                🪞 מראות: מרחב אישי לשיח פנימי ומפתח עם עצמך 🪞
            </div>
            <div class="hebrew-text" style="text-align: center; margin-bottom: 20px;">
                מקום בטוח לשוחח עם החלקים השונים של עצמך ולפתח הבנה עצמית עמוקה יותר
            </div>
            <div style="background-color: #e8f5e8; border: 1px solid #4caf50; padding: 10px; margin: 10px 0; border-radius: 5px; text-align: center;">
                <strong>{status_message}</strong>
            </div>
            """)
            
            # Main interface areas
            with gr.Column():
                
                # Step 1: Initial context gathering
                with gr.Group(visible=True) as initial_step:
                    gr.Markdown("## שלב 1: ספר/ספרי על עצמך", elem_classes=["hebrew-text"])
                    
                    initial_prompts = get_initial_prompts()
                    
                    initial_choice = gr.Radio(
                        choices=[
                            ("תאר/תארי את עצמך כאדם", "describe_self"),
                            ("איך אתה חושב שאחרים רואים אותך?", "self_perception"),
                            ("איזה אתגר אתה חווה עכשיו בחיים?", "current_challenge")
                        ],
                        label="בחר/בחרי נושא לשיתוף:",
                        elem_classes=["hebrew-text"]
                    )
                    
                    user_context_input = gr.Textbox(
                        label="ספר/ספרי בכמה משפטים:",
                        placeholder="כתוב/כתבי כאן את המחשבות שלך...",
                        lines=4,
                        elem_classes=["hebrew-text"]
                    )
                    
                    continue_to_parts = gr.Button("המשך לבחירת חלק פנימי", variant="primary")
                
                # Step 2: Part selection
                with gr.Group(visible=False) as parts_step:
                    gr.Markdown("## שלב 2: בחר/בחרי חלק פנימי לשיחה", elem_classes=["hebrew-text"])
                    
                    part_selection = gr.Radio(
                        choices=[
                            ("הקול הביקורתי - החלק שמנסה להגן עליך על ידי ביקורת והכוונה", "הקול הביקורתי"),
                            ("הילד/ה הפנימית - החלק הפגיע, הצעיר והאמיתי שלך", "הילד/ה הפנימית"),
                            ("המרצה - החלק שרוצה שכולם יהיו מרוצים", "המרצה"),
                            ("המגן - החלק החזק שמגן עליך מפני פגיעות", "המגן"),
                            ("הנמנע/ת - החלק שמעדיף להימנע ממצבים מאתגרים", "הנמנע/ת")
                        ],
                        label="איזה חלק פנימי תרצה לפגוש?",
                        elem_classes=["hebrew-text"]
                    )
                    
                    # Customization options
                    with gr.Accordion("התאמה אישית (אופציונלי)", open=False):
                        persona_name = gr.Textbox(
                            label="שם לחלק הזה:",
                            placeholder="למשל: דני, מיכל, אבי...",
                            elem_classes=["hebrew-text"]
                        )
                        persona_age = gr.Textbox(
                            label="גיל או תקופת חיים:",
                            placeholder="למשל: ילד/ה, מתבגר/ת, בוגר/ת...",
                            elem_classes=["hebrew-text"]
                        )
                        persona_style = gr.Textbox(
                            label="סגנון דיבור מיוחד:",
                            placeholder="למשל: רגוש, רציני, משעשע...",
                            elem_classes=["hebrew-text"]
                        )
                    
                    start_conversation = gr.Button("התחל שיחה", variant="primary")
                
                # Step 3: Conversation interface
                with gr.Group(visible=False) as conversation_step:
                    gr.Markdown("## שיחה עם החלק הפנימי שלך", elem_classes=["hebrew-text"])
                    
                    current_part_display = gr.Markdown("", elem_classes=["hebrew-text"])
                    
                    # Chat interface
                    with gr.Row():
                        with gr.Column(scale=4):
                            chatbot = gr.Chatbot(
                                height=400,
                                label="השיחה שלך",
                                elem_classes=["hebrew-text"],
                                rtl=True
                            )
                            
                            msg_input = gr.Textbox(
                                label="ההודעה שלך:",
                                placeholder="כתוב/כתבי את המחשבות שלך כאן...",
                                lines=2,
                                elem_classes=["hebrew-text"]
                            )
                            
                            with gr.Row():
                                send_btn = gr.Button("שלח", variant="primary")
                                clear_btn = gr.Button("נקה שיחה")
                                
                        with gr.Column(scale=1):
                            gr.Markdown("### פעולות נוספות", elem_classes=["hebrew-text"])
                            change_part_btn = gr.Button("החלף חלק פנימי")
                            restart_btn = gr.Button("התחל מחדש")
            
            # Event handlers
            def process_initial_context(choice, context, state):
                """Process initial context and move to part selection"""
                if not choice or not context.strip():
                    gr.Warning("אנא בחר נושא וכתב משהו כדי להמשיך")
                    return state, gr.update(), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
                
                state = self.conversation_manager.set_initial_context(state, choice, context)
                return (
                    state, 
                    gr.update(),
                    gr.update(visible=False), 
                    gr.update(visible=True), 
                    gr.update(visible=False)
                )
            
            def start_chat(part, p_name, p_age, p_style, state):
                """Start the conversation with selected part"""
                if not part:
                    gr.Warning("אנא בחר חלק פנימי כדי להתחיל")
                    return state, gr.update(), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update()
                
                state = self.conversation_manager.set_selected_part(
                    state, part, p_name.strip() if p_name else None, 
                    p_age.strip() if p_age else None, p_style.strip() if p_style else None
                )
                
                part_info = DEFAULT_PARTS.get(part, {})
                display_name = (p_name.strip() if p_name else None) or part_info.get("default_persona_name", "חלק פנימי")
                
                display_text = f"🗣️ כעת אתה מתשוחח עם: **{display_name}** ({part})"
                
                return (
                    state,
                    display_text,
                    gr.update(visible=False),
                    gr.update(visible=False), 
                    gr.update(visible=True),
                    []
                )
            
            def handle_message(message, history, state):
                """Handle user message and generate response"""
                if not message.strip():
                    return "", history, state
                
                # Generate response
                response = self.generate_response(message, state)
                
                # Update conversation state
                state = self.conversation_manager.add_to_history(state, message, response)
                
                # Update history for display
                history.append([message, response])
                
                return "", history, state
            
            def clear_conversation(state):
                """Clear conversation history"""
                state = self.conversation_manager.clear_conversation(state)
                return [], state
            
            def change_part():
                """Return to part selection"""
                return (
                    gr.update(visible=False),
                    gr.update(visible=True),
                    gr.update(visible=False)
                )
            
            def restart_completely():
                """Restart the entire session"""
                new_state = self.conversation_manager.create_new_session()
                return (
                    new_state,
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    [],
                    "",
                    "",
                    None,
                    None,
                    "",
                    "",
                    ""
                )
            
            # Wire up event handlers
            continue_to_parts.click(
                fn=process_initial_context,
                inputs=[initial_choice, user_context_input, conversation_state],
                outputs=[conversation_state, current_part_display, initial_step, parts_step, conversation_step]
            )
            
            start_conversation.click(
                fn=start_chat,
                inputs=[part_selection, persona_name, persona_age, persona_style, conversation_state],
                outputs=[conversation_state, current_part_display, initial_step, parts_step, conversation_step, chatbot]
            )
            
            # Chat message handling
            msg_input.submit(
                fn=handle_message,
                inputs=[msg_input, chatbot, conversation_state],
                outputs=[msg_input, chatbot, conversation_state]
            )
            
            send_btn.click(
                fn=handle_message,
                inputs=[msg_input, chatbot, conversation_state],
                outputs=[msg_input, chatbot, conversation_state]
            )
            
            clear_btn.click(
                fn=clear_conversation,
                inputs=[conversation_state],
                outputs=[chatbot, conversation_state]
            )
            
            change_part_btn.click(
                fn=change_part,
                outputs=[conversation_step, parts_step, initial_step]
            )
            
            restart_btn.click(
                fn=restart_completely,
                outputs=[conversation_state, initial_step, parts_step, conversation_step, chatbot, 
                        user_context_input, current_part_display, initial_choice, part_selection,
                        persona_name, persona_age, persona_style]
            )
        
        return demo

def main():
    """Main function to launch the application"""
    logger.info("Starting מראות application...")
    
    try:
        app = MirautrApp()
        demo = app.create_main_interface()
        
        # Check environment
        is_hf_spaces = os.getenv("SPACE_ID") is not None
        
        logger.info(f"Launching app... HF Spaces: {is_hf_spaces}")
        
        # Unified launch configuration for both environments
        # This ensures identical experience in both local and HF Spaces
        launch_config = {
            "show_error": True,
            "show_api": False,  # Disable API docs to avoid schema issues
            "favicon_path": None,
            "auth": None,
            "enable_queue": False,  # Disable queue to prevent schema issues
            "max_threads": 1  # Limit threads for stability
        }
        
        if is_hf_spaces:
            # HF Spaces specific settings
            logger.info("Configuring for HF Spaces deployment")
            launch_config.update({
                "server_name": "0.0.0.0",
                "server_port": 7860,
                "share": False,  # HF Spaces handles public access
                "quiet": True
            })
        else:
            # Local development settings
            logger.info("Configuring for local development")
            
            # Try to find an available port
            default_port = int(os.getenv("GRADIO_SERVER_PORT", "7861"))
            available_port = default_port
            
            # Check if port is available, if not find next available
            import socket
            for port_try in range(default_port, default_port + 10):
                try:
                    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                        s.bind(('127.0.0.1', port_try))
                        available_port = port_try
                        break
                except OSError:
                    continue
            
            logger.info(f"Using port {available_port} for local development")
            
            launch_config.update({
                "server_name": "127.0.0.1", 
                "server_port": available_port,
                "share": True,  # Enable share for local testing to avoid localhost issues
                "inbrowser": True,  # Auto-open browser
                "quiet": False
            })
        
        demo.launch(**launch_config)
            
    except Exception as e:
        logger.error(f"Failed to start application: {e}")
        raise

if __name__ == "__main__":
    main()