src/ui/page_handlers.py

"""Page handling and navigation functionality"""
import uuid

import gradio as gr
from ..config.settings import TUTORIAL_EXAMPLE
from ..session.session_manager import SessionManager


class PageHandlers:
    """Handles page switching, navigation, and page-specific functionality"""
    
    def __init__(self, app):
        self.app = app
    
    def switch_to_page(self, page_name, update_nav=True):
        """Switch to a different page and update navigation if needed."""
        page_updates = {
            "welcome": (True, False, False, False, False, False),
            "tutorial": (False, True, False, False, False, False), 
            "creative": (False, False, True, False, False, False),
            "gallery": (False, False, False, True, False, False),
            "results": (False, False, False, False, True, False),
            "session": (False, False, False, False, False, True)
        }
        
        if page_name not in page_updates:
            page_name = "welcome"
            
        visibility = page_updates[page_name]
        
        nav_updates = []
        if update_nav:
            # Update navigation button styles
            nav_buttons = ["welcome", "tutorial", "creative", "gallery", "results", "session"]
            for nav_page in nav_buttons:
                if nav_page == page_name:
                    nav_updates.append(gr.update(elem_classes=["nav-button", "active"]))
                else:
                    nav_updates.append(gr.update(elem_classes=["nav-button"]))
        else:
            nav_updates = [gr.update()] * 6
            
        return (
            page_name,  # current_page state
            gr.update(visible=visibility[0]),  # landing_page 
            gr.update(visible=visibility[1]),  # tutorial_page
            gr.update(visible=visibility[2]),  # creative_page
            gr.update(visible=visibility[3]),  # gallery_page
            gr.update(visible=visibility[4]),  # results_page
            gr.update(visible=visibility[5]),  # session_page
            *nav_updates  # nav button updates
        )
    
    def show_gallery(self, min_score=0.3):
        """Show gallery with basic UI updates"""
        gallery_html = self.generate_gallery_html(min_score)
        switch_result = self.switch_to_page("gallery", update_nav=False)
        return switch_result + (gallery_html,)
    
    def generate_gallery_html(self, min_score=0.3):
        """Generate gallery HTML content"""
        gallery_responses = self.app.data_manager.get_gallery_responses(min_score=min_score, limit=20)
        
        if not gallery_responses:
            return """
                <div class="gallery-item">
                    <h3>🔍 No responses found at this creativity level yet!</h3>
                    <p>Be the first to create some legendary responses! Lower the filter or come back later.</p>
                </div>
                """
        
        gallery_html = ""
        for i, response in enumerate(gallery_responses, 1):
            # Truncate long responses for gallery display
            user_input = response.get("user_continuation", "").strip()
            full_response = response.get("full_response_from_user", "").strip()
            prompt = response.get("prompt", "").strip()
            
            # Truncate for display
            display_prompt = (prompt[:100] + "...") if len(prompt) > 100 else prompt
            
            achievements = self.app.achievement_system.determine_achievement_titles(
                response["cosine_distance"], 
                response["num_user_tokens"]
            )
            achievement_badges = " ".join([f'<span class="achievement-badge">{title}</span>' for title in achievements[:2]])
            
            # Get original response for comparison
            original_response = response.get("llm_full_response_original", "").strip()
            
            gallery_html += f"""
            <div class="gallery-item">
                <div class="gallery-item-score">
                    🎨 Creativity Score: {response['cosine_distance']:.3f} 
                    ({response['num_user_tokens']} token{'s' if response['num_user_tokens'] != 1 else ''})
                </div>
                {achievement_badges}
                <h4>#{i} - Prompt:</h4>
                <p style="opacity: 0.8; font-style: italic;">{display_prompt}</p>
                <h4>Creative Input: <span style="color: #ffa726;">"{user_input}"</span></h4>
                
                <div style="display: flex; gap: 15px; margin-top: 15px;">
                    <div style="flex: 1;">
                        <h4>🤖 Original AI Response:</h4>
                        <div style="background: rgba(255,255,255,0.05); padding: 15px; border-radius: 8px; border-left: 4px solid #999; max-height: 200px; overflow-y: auto; font-family: monospace; font-size: 14px;">
                            {original_response}
                        </div>
                    </div>
                    <div style="flex: 1;">
                        <h4>✨ Creative Response (after user input):</h4>
                        <div style="background: rgba(255,255,255,0.05); padding: 15px; border-radius: 8px; border-left: 4px solid #ffa726; max-height: 200px; overflow-y: auto; font-family: monospace; font-size: 14px;">
                            {full_response}
                        </div>
                    </div>
                </div>
            </div>
            """
        
        return gallery_html
    
    def show_inspire_me_examples(self, current_prompt):
        """Show new random inspiring examples for current prompt"""
        if not current_prompt:
            return ""
        
        examples = self.app.data_manager.get_inspire_me_examples(
            current_prompt["prompt"],
            current_prompt["llm_partial_response"],
            limit=5
        )
        
        return self.generate_inspire_me_html(examples)
    
    def show_tutorial(self):
        """Show the tutorial page starting at step 1"""
        example = TUTORIAL_EXAMPLE
        step_content = self.get_tutorial_step_content(1)
        switch_result = self.switch_to_page("tutorial", update_nav=False)
        return switch_result + (
            gr.update(value=step_content["explanation"]),  # tutorial_step_explanation
            gr.update(visible=step_content["show_prompt"], value=example["prompt"]),  # tutorial_prompt_display
            gr.update(visible=step_content["show_partial"], value=example["llm_partial_response"]),  # tutorial_partial_response_display
            gr.update(visible=step_content["show_input"]),  # tutorial_input_section
            gr.update(visible=step_content["show_results"]),  # tutorial_results_section
            gr.update(visible=False),  # tutorial_prev_btn (hidden on step 1)
            gr.update(visible=True),   # tutorial_next_btn
            gr.update(visible=False),  # tutorial_play_btn (hidden until last step)
            1  # Reset tutorial step to 1
        )
    
    def get_tutorial_step_content(self, step):
        """Get content for a specific tutorial step using unified example data"""
        example = TUTORIAL_EXAMPLE
        
        step_content = {
            1: {
                "explanation": "## Step 1: The Prompt\n\nFirst, I'll show you what I was asked to write about ✍️<br>This gives you context for helping me be creative!",
                "show_prompt": True,
                "show_partial": False,
                "show_input": False,
                "show_results": False
            },
            2: {
                "explanation": "## Step 2: The start of my response\n\nThen I'll show you how I started responding 🤖😴<br>As you can see... it's pretty boring and predictable!",
                "show_prompt": True,
                "show_partial": True,
                "show_input": False,
                "show_results": False
            },
            3: {
                "explanation": "## Step 3: Your creative input \n\nThis is where the magic happens ✨ You add a few words (**tokens**) to steer me somewhere interesting.<br>Think unexpected, surprising, or completely different from what I started with.<br>I only want a little help - make each word count!<br><br>**🔤 What are tokens?** Your input is measured in \\\"tokens\\\" (how I process text). About 5 tokens ≈ 3-5 words.<br>Don't worry - you'll see your exact token count as you type!",
                "show_prompt": True,
                "show_partial": True,
                "show_input": True,
                "show_results": False,
                "example_input": example["example_input"]
            },
            4: {
                "explanation": "## Step 4: I'll continue the response \n\nAfter you give me your creative direction, I'll finish the response based on your inspiration 🤖🎉<br>You'll see how different the result is from my original boring version - that's the magic of collaboration!",
                "show_prompt": True,
                "show_partial": True,
                "show_input": True,
                "show_results": True,
                "example_input": example["example_input"],
                "example_creative": example["example_creative"]
            },
            5: {
                "explanation": "## Step 5: Get scored on our collaborative response 🎨\n\nAfter I finish our collaborative response, we'll get scored for how different our combined response is from my original boring version!<br><br>**📊 How scoring works:** I use \\\"semantic similarity\\\" to compare the meaning of both responses. The more different the meaning, the higher your creativity score!<br>• 0.7+ = 🌟 Legendary creativity<br>• 0.5+ = 🔥 Exceptional creativity<br>• 0.3+ = ✨ Great creativity<br><br>**💡 Pro tip:** Unexpected words often lead to the most creative directions! 🏆",
                "show_prompt": False,
                "show_partial": False,
                "show_input": False,
                "show_results": True,
                "example_input": example["example_input"],
                "example_creative": example["example_creative"],
                "example_score": example["example_score"]
            }
        }
        
        return step_content.get(step, step_content[1])
    
    def generate_inspire_me_html(self, examples):
        """Generate HTML for inspire me examples"""
        if not examples:
            return self.app.template_renderer.load_template("inspire-me-empty.html")
        else:
            inspire_html = self.app.template_renderer.load_template("inspire-me-header.html")
            for example in examples:
                user_input = example.get("user_continuation", "").strip()
                score = example["cosine_distance"]
                tokens = example["num_user_tokens"]
                plural = "s" if tokens != 1 else ""
                
                inspire_html += self.app.template_renderer.load_template("inspire-me-example.html",
                    user_input=user_input,
                    score=score,
                    tokens=tokens,
                    plural=plural
                )
            
            inspire_html += self.app.template_renderer.load_template("inspire-me-footer.html")
            return inspire_html
    
    def tutorial_next_step(self, current_step):
        """Advance to next tutorial step"""
        next_step = min(current_step + 1, 5)
        example = TUTORIAL_EXAMPLE
        step_content = self.get_tutorial_step_content(next_step)
        
        # Prepare example data for the step
        input_value = step_content.get("example_input", "")
        token_count = self.app.llm_manager.count_tokens(input_value) if input_value else 0
        
        # Create score display for final step
        score_display = ""
        if next_step == 5 and "example_score" in step_content:
            score_display = self.app.scorer.create_enhanced_score_display(
                step_content["example_score"], 1, 95.0, token_count, 10
            )
        
        return (
            gr.update(value=step_content["explanation"]),  # tutorial_step_explanation
            gr.update(visible=step_content["show_prompt"]),  # tutorial_prompt_display
            gr.update(visible=step_content["show_partial"]),  # tutorial_partial_response_display
            gr.update(visible=step_content["show_input"]),  # tutorial_input_section
            gr.update(value=input_value) if step_content["show_input"] else gr.update(),  # tutorial_user_input
            gr.update(value=f"**Tokens:** {token_count}/5") if step_content["show_input"] else gr.update(),  # tutorial_token_counter
            gr.update(visible=step_content["show_results"]),  # tutorial_results_section
            gr.update(value=example["llm_full_response_original"]) if step_content["show_results"] else gr.update(),  # tutorial_original_response
            gr.update(value=step_content.get("example_creative", "")) if step_content["show_results"] else gr.update(),  # tutorial_creative_response
            gr.update(value=score_display) if step_content["show_results"] else gr.update(),  # tutorial_score_display
            gr.update(visible=next_step > 1),  # tutorial_prev_btn
            gr.update(visible=next_step < 5),  # tutorial_next_btn
            gr.update(visible=next_step == 5),  # tutorial_play_btn (show on last step)
            next_step  # Update tutorial step
        )
    
    def tutorial_prev_step(self, current_step):
        """Go back to previous tutorial step"""
        prev_step = max(current_step - 1, 1)
        step_content = self.get_tutorial_step_content(prev_step)
        
        return (
            gr.update(value=step_content["explanation"]),  # tutorial_step_explanation
            gr.update(visible=step_content["show_prompt"]),  # tutorial_prompt_display
            gr.update(visible=step_content["show_partial"]),  # tutorial_partial_response_display
            gr.update(visible=step_content["show_input"]),  # tutorial_input_section
            gr.update(visible=step_content["show_results"]),  # tutorial_results_section
            gr.update(visible=prev_step > 1),  # tutorial_prev_btn
            gr.update(visible=prev_step < 4),  # tutorial_next_btn
            gr.update(visible=prev_step == 4),  # tutorial_play_btn
            prev_step  # Update tutorial step
        )
    
    def start_game(self, current_prompt, session_id):
        """Start the game with state management"""
        
        # Initialize session if needed
        if session_id is None:
            session_id = str(uuid.uuid4())
        
        # Get new prompt if needed
        if current_prompt is None:
            current_prompt = self.app.get_random_prompt()
        
        # Load initial inspire me examples
        initial_examples = self.app.data_manager.get_inspire_me_examples(
            current_prompt["prompt"],
            current_prompt["llm_partial_response"],
            limit=5
        )
        inspire_html = self.generate_inspire_me_html(initial_examples)
        switch_result = self.switch_to_page("creative", update_nav=False)
        return switch_result + (
            current_prompt['prompt'],
            current_prompt['llm_partial_response'],
            inspire_html,
            current_prompt,
            session_id
        )
    
    def back_to_landing(self):
        """Go back to landing page"""
        result = self.switch_to_page("welcome", update_nav=True)
        return result + ("",)  # Add empty gallery_display as the 14th output
    
    def update_token_count(self, text):
        """Update token counter and visual tokens in real-time"""
        if not text.strip():
            return (
                gr.update(value="**Tokens:** 0/5", elem_classes=["token-counter"]),
                gr.update(value="Type to see tokens...", elem_classes=["token-visualization"])
            )
        
        # Count tokens and decode them individually
        tokens, token_texts = self.app.llm_manager.tokenize_for_visualization(text)
        token_count = len(tokens)
        
        # Create visual token display
        token_html_parts = []
        for i, token_text in enumerate(token_texts):
            if i < 5:
                token_class = f"token-{i+1}"
            else:
                token_class = "token-excess"
            
            # Escape HTML and handle special characters
            token_display = token_text.replace(' ', '·').replace('\\n', '↵')
            if not token_display.strip():
                token_display = '·'
            
            token_html_parts.append(f'<span class="token {token_class}">{token_display}</span>')
        
        if token_html_parts:
            token_visualization_html = ' '.join(token_html_parts)
        else:
            token_visualization_html = "Type to see tokens..."
        
        # Update counter with color coding
        if token_count <= 5:
            counter_class = ["token-counter", "valid"]
            counter_text = f"**Tokens:** {token_count}/5"
        else:
            counter_class = ["token-counter", "error"] 
            counter_text = f"**Tokens:** {token_count}/5 (Too many!)"
        
        return (
            gr.update(value=counter_text, elem_classes=counter_class),
            gr.update(value=token_visualization_html, elem_classes=["token-visualization"])
        )
    
    def on_submit_start(self, user_text):
        if not user_text.strip():
            return (
                gr.update(interactive=True),
                gr.update(visible=False),
                gr.update(visible=True, value="Please enter some text to continue the response."),
                gr.update(visible=False),  # Hide landing
                gr.update(visible=False),  # Hide gallery
                gr.update(visible=True),   # Show creative
                gr.update(visible=False),  # Hide results
                gr.update()
            )
        
        return (
            gr.update(interactive=False),
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=False),  # Hide landing
            gr.update(visible=False),  # Hide gallery
            gr.update(visible=True),   # Show creative
            gr.update(visible=False),  # Hide results
            gr.update()
        )
    
    def on_submit_process(self, user_text, current_prompt, session_id):
        result = self.app.process_submission(user_text, current_prompt, session_id)
        generated_response, cosine_distance, rank, percentile, mean_score, violin_plot, prompt_results = result
        
        if "Error:" in generated_response or "Please" in generated_response:
            return (
                gr.update(interactive=True),     # submit_btn
                gr.update(visible=False),        # loading_message
                gr.update(visible=True, value=generated_response),  # error_message
                gr.update(visible=False),        # landing_page
                gr.update(visible=False),        # gallery_page
                gr.update(visible=True),         # creative_page
                gr.update(visible=False),        # results_page
                gr.update(visible=False),        # cosine_distance_display
                gr.update(),                     # results_prompt_display
                gr.update(),                     # results_partial_display
                gr.update(),                     # results_user_input_display
                gr.update(),                     # original_response_display
                gr.update(),                     # results_continuation_display
                gr.update(visible=False),        # violin_plot_display
                gr.update(),                     # current_page (no change)
                gr.update(),                     # has_results (no change)
                gr.update(),                     # nav_results (no change)
                gr.update(),                     # nav_session (no change)
                gr.update(),                     # nav_welcome (no change)
                gr.update(),                     # nav_tutorial (no change)
                gr.update(),                     # nav_creative (no change) 
                gr.update()                      # nav_gallery (no change)
            )
        
        # Create enhanced score display with progress bars and metrics
        user_tokens = self.app.llm_manager.count_tokens(user_text)
        same_category_attempts = len(prompt_results)
        
        score_text = self.app.scorer.create_enhanced_score_display(
            cosine_distance, rank, percentile, user_tokens, same_category_attempts
        )
        
        # Show Results/Progress nav buttons after first submission
        # Switch to results page and make nav buttons visible
        return (
            gr.update(interactive=True),                    # submit_btn
            gr.update(visible=False),                       # loading_message
            gr.update(visible=False),                       # error_message
            gr.update(visible=False),                       # landing_page (hide)
            gr.update(visible=False),                       # gallery_page (hide)  
            gr.update(visible=False),                       # creative_page (hide)
            gr.update(visible=True),                        # results_page (show)
            gr.update(visible=True, value=score_text),      # cosine_distance_display
            gr.update(value=current_prompt['prompt']), # results_prompt_display
            gr.update(value=current_prompt['llm_partial_response']), # results_partial_display
            gr.update(value=user_text),                     # results_user_input_display
            gr.update(value=current_prompt['llm_full_response_original']), # original_response_display
            gr.update(value=generated_response),          # results_continuation_display
            gr.update(visible=True, value=violin_plot) if violin_plot else gr.update(visible=False),  # violin_plot_display
            "results",                                       # current_page state
            True,                                           # has_results state
            gr.update(visible=True, elem_classes=["nav-button", "active"]),  # nav_results (visible + active)
            gr.update(visible=True, elem_classes=["nav-button"]),            # nav_session (visible + inactive)
            gr.update(elem_classes=["nav-button"]),         # nav_welcome (inactive)
            gr.update(elem_classes=["nav-button"]),         # nav_tutorial (inactive)  
            gr.update(elem_classes=["nav-button"]),         # nav_creative (inactive)
            gr.update(elem_classes=["nav-button"])          # nav_gallery (inactive)
        )
    
    def try_same_prompt(self):
        switch_result = self.switch_to_page("creative")
        return switch_result + (
            gr.update(interactive=True),  # submit_btn
            gr.update(value=""),         # user_input (clear)
            gr.update(visible=False),    # loading_message
            gr.update(visible=False),    # error_message
            gr.update(value="**Tokens:** 0/5"),  # token_counter (reset)
            gr.update(value="Once you start typing, I'll show you what you've written, but in tokens!")  # token_visualization (reset)
        )
    
    def try_new_prompt(self):
        new_prompt = self.app.get_random_prompt()
        initial_examples = self.app.data_manager.get_inspire_me_examples(
            new_prompt["prompt"],
            new_prompt["llm_partial_response"],
            limit=5
        )
        inspire_html = self.generate_inspire_me_html(initial_examples)
        switch_result = self.switch_to_page("creative")
        return switch_result + (
            new_prompt['prompt'],       # prompt_display
            new_prompt['llm_partial_response'], # partial_response_display
            gr.update(interactive=True), # submit_btn
            "",                         # user_input (clear)
            gr.update(visible=False),   # loading_message
            gr.update(visible=False),   # error_message
            inspire_html,               # inspire_me_display
            new_prompt                  # current_prompt state update
        )
    
    def load_creative_with_prompt(self, current_prompt, session_id):
        """Load creative page with a new prompt"""
        
        # Initialize session if needed
        if session_id is None:
            session_id = str(uuid.uuid4())
            
        # Get new prompt (if none provided)
        if current_prompt is None:
            current_prompt = self.app.get_random_prompt()
        initial_examples = self.app.data_manager.get_inspire_me_examples(
            current_prompt["prompt"],
            current_prompt["llm_partial_response"],
            limit=5
        )
        inspire_html = self.generate_inspire_me_html(initial_examples)
        switch_result = self.switch_to_page("creative")
        return switch_result + (current_prompt['prompt'], current_prompt['llm_partial_response'], inspire_html, current_prompt, session_id)
    
    def load_gallery_with_content(self):
        """Load gallery page with gallery content"""
        gallery_html = self.generate_gallery_html(0.3)
        switch_result = self.switch_to_page("gallery")
        return switch_result + (gallery_html,)
    
    def load_tutorial_with_content(self):
        """Load tutorial page with tutorial content initialized"""
        example = TUTORIAL_EXAMPLE
        step_content = self.get_tutorial_step_content(1)
        switch_result = self.switch_to_page("tutorial")
        return switch_result + (
            step_content["explanation"],  # tutorial_step_explanation
            gr.update(value=example["prompt"], visible=step_content["show_prompt"]),  # tutorial_prompt_display
            gr.update(value=example["llm_partial_response"], visible=step_content["show_partial"]),  # tutorial_partial_response_display
            gr.update(visible=step_content["show_input"]),  # tutorial_input_section
            gr.update(visible=step_content["show_results"]),  # tutorial_results_section
            gr.update(visible=False),  # tutorial_prev_btn (hidden on step 1)
            gr.update(visible=True),   # tutorial_next_btn
            gr.update(visible=False),  # tutorial_play_btn (hidden until last step)
            1  # Reset tutorial step to 1
        )
    
    def navigate_to_session(self, session_id):
        
        # Create session manager instance with session_id
        session_manager = SessionManager(
            session_id=session_id,
            data_manager=self.app.data_manager,
            achievement_system=self.app.achievement_system
        )
        
        main_stats, history, achievements = session_manager.generate_session_page_html(
            self.app.template_renderer,
            self.app.statistics_calculator,
            self.app.scorer
        )
        switch_result = self.switch_to_page("session")
        return switch_result + (main_stats, history, achievements)