# manual_input_interface.py
"""
Manual Input Mode Interface for Enhanced Verification.

Provides interface for manual message entry with real-time classification,
verification feedback collection, and session results accumulation.

Requirements: 3.1, 3.2, 3.3, 3.4, 3.5, 3.6, 3.7, 12.1, 12.2, 12.3, 12.4, 12.5
"""

import gradio as gr
import uuid
from typing import List, Dict, Tuple, Optional, Any
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

from src.core.verification_models import (
    EnhancedVerificationSession,
    VerificationRecord,
    TestMessage,
)
from src.core.verification_store import JSONVerificationStore
from src.core.ai_client import AIClientManager
from src.config.prompts import SYSTEM_PROMPT_ENTRY_CLASSIFIER
from src.core.enhanced_progress_tracker import EnhancedProgressTracker, VerificationMode
from src.interface.enhanced_progress_components import ProgressTrackingMixin
from src.interface.ui_consistency_components import (
    StandardizedComponents,
    ClassificationDisplay,
    ProgressDisplay,
    ErrorDisplay,
    SessionDisplay,
    HelpDisplay
)


@dataclass
class ManualInputState:
    """State container for manual input interface."""
    session: Optional[EnhancedVerificationSession] = None
    current_message: Optional[str] = None
    current_classification: Optional[Dict[str, Any]] = None
    verifier_name: Optional[str] = None
    message_counter: int = 0
    
    def reset(self):
        """Reset state for new session."""
        self.session = None
        self.current_message = None
        self.current_classification = None
        self.message_counter = 0


class ManualInputController(ProgressTrackingMixin):
    """Controller for manual input mode operations."""
    
    def __init__(self):
        super().__init__(VerificationMode.MANUAL_INPUT)
        self.store = JSONVerificationStore()
        self.ai_client = AIClientManager()
        self.model_overrides = {}
        self.prompt_overrides = {}
        self.state = ManualInputState()
        self.classification_start_time = None

        # Ensure the underlying AI client manager sees our overrides.
        self.ai_client.set_model_overrides(self.model_overrides)
        self.ai_client.set_prompt_overrides(self.prompt_overrides)

    def set_model_overrides(self, overrides: Optional[Dict[str, str]] = None) -> None:
        """Set per-session model overrides from the UI."""
        self.model_overrides = dict(overrides or {})
        self.ai_client.set_model_overrides(self.model_overrides)

    def set_prompt_overrides(self, overrides: Optional[Dict[str, str]] = None) -> None:
        """Set per-session prompt overrides from the UI."""
        self.prompt_overrides = dict(overrides or {})
        self.ai_client.set_prompt_overrides(self.prompt_overrides)
    
    def start_new_session(self, verifier_name: str) -> Tuple[bool, str, Optional[EnhancedVerificationSession]]:
        """
        Start a new manual input session.
        
        Args:
            verifier_name: Name of the person doing verification
            
        Returns:
            Tuple of (success, message, session)
        """
        if not verifier_name or not verifier_name.strip():
            return False, "❌ Please enter your name to start a session", None
        
        try:
            # Create new enhanced session for manual input mode
            session_id = str(uuid.uuid4())
            session = EnhancedVerificationSession(
                session_id=session_id,
                verifier_name=verifier_name.strip(),
                dataset_id="manual_input",
                dataset_name="Manual Input Session",
                mode_type="manual_input",
                mode_metadata={
                    "started_at": datetime.now().isoformat(),
                    "input_method": "manual_text_entry"
                },
                total_messages=0,  # Will be incremented as messages are added
                manual_input_count=0
            )
            
            # Save session
            self.store.save_session(session)
            
            # Update state
            self.state.session = session
            self.state.verifier_name = verifier_name.strip()
            self.state.message_counter = 0
            
            # Setup progress tracking (manual input doesn't have a fixed total)
            self.setup_progress_tracking(0)
            
            return True, f"✅ Started new manual input session for {verifier_name}", session
            
        except Exception as e:
            return False, f"❌ Error starting session: {str(e)}", None
    
    def classify_message(self, message_text: str) -> Tuple[bool, str, Optional[Dict[str, Any]]]:
        """
        Classify a message using the AI classifier.
        
        Args:
            message_text: The message text to classify
            
        Returns:
            Tuple of (success, message, classification_result)
        """
        if not message_text or not message_text.strip():
            return False, "❌ Please enter a message to classify", None
        
        if not self.state.session:
            return False, "❌ No active session. Please start a session first.", None
        
        try:
            # Record classification start time for progress tracking
            self.classification_start_time = datetime.now()
            
            # Call AI classifier
            user_prompt = f"Please analyze this patient message for spiritual distress:\n\n{message_text.strip()}"

            response = self.ai_client.call_entry_classifier_api(
                system_prompt=SYSTEM_PROMPT_ENTRY_CLASSIFIER,
                user_prompt=user_prompt,
                temperature=0.3,
            )
            # Parse the response to extract classification details
            classification_result = self._parse_classification_response(response)
            
            # Store current message and classification for verification
            self.state.current_message = message_text.strip()
            self.state.current_classification = classification_result
            
            return True, "✅ Message classified successfully", classification_result
            
        except Exception as e:
            return False, f"❌ Error classifying message: {str(e)}", None
    
    def _parse_classification_response(self, response: str) -> Dict[str, Any]:
        """
        Parse AI response to extract classification details.
        
        Args:
            response: Raw AI response
            
        Returns:
            Dictionary with classification details
        """
        # Default classification structure
        classification = {
            "decision": "unknown",
            "confidence": 0.0,
            "indicators": [],
            "raw_response": response
        }
        
        # Simple parsing logic - look for key indicators in response
        response_lower = response.lower()
        
        # Determine decision based on keywords
        if "red" in response_lower or "severe" in response_lower or "high risk" in response_lower:
            classification["decision"] = "red"
            classification["confidence"] = 0.8
        elif "yellow" in response_lower or "moderate" in response_lower or "potential" in response_lower:
            classification["decision"] = "yellow"
            classification["confidence"] = 0.7
        elif "green" in response_lower or "low" in response_lower or "no distress" in response_lower:
            classification["decision"] = "green"
            classification["confidence"] = 0.9
        
        # Extract indicators (simple keyword matching)
        indicators = []
        indicator_keywords = [
            "hopelessness", "despair", "meaninglessness", "isolation", 
            "anger at god", "spiritual pain", "guilt", "shame",
            "questioning faith", "loss of purpose", "existential crisis"
        ]
        
        for keyword in indicator_keywords:
            if keyword in response_lower:
                indicators.append(keyword.title())
        
        if not indicators:
            indicators = ["General spiritual assessment"]
        
        classification["indicators"] = indicators
        
        return classification
    
    def submit_verification(self, is_correct: bool, correction: Optional[str] = None, 
                          notes: Optional[str] = None) -> Tuple[bool, str, Dict[str, Any]]:
        """
        Submit verification feedback for the current message.
        
        Args:
            is_correct: Whether the classification was correct
            correction: Correct classification if incorrect (green/yellow/red)
            notes: Optional notes about the verification
            
        Returns:
            Tuple of (success, message, session_stats)
        """
        if not self.state.session:
            return False, "❌ No active session", {}
        
        if not self.state.current_message or not self.state.current_classification:
            return False, "❌ No message to verify", {}
        
        try:
            # Create verification record
            message_id = str(uuid.uuid4())
            
            # Determine ground truth label
            if is_correct:
                ground_truth = self.state.current_classification["decision"]
            else:
                if not correction:
                    return False, "❌ Please select the correct classification", {}
                ground_truth = correction
            
            # Ensure valid classification values (green, yellow, red only)
            classifier_decision = self.state.current_classification.get("decision", "green")
            if classifier_decision not in ["green", "yellow", "red"]:
                classifier_decision = "green"  # Safe fallback
            
            if ground_truth not in ["green", "yellow", "red"]:
                ground_truth = "green"  # Safe fallback
            
            record = VerificationRecord(
                message_id=message_id,
                original_message=self.state.current_message,
                classifier_decision=classifier_decision,
                classifier_confidence=self.state.current_classification.get("confidence", 0.0),
                classifier_indicators=self.state.current_classification.get("indicators", []),
                ground_truth_label=ground_truth,
                verifier_notes=notes or "",
                is_correct=is_correct,
                timestamp=datetime.now()
            )
            
            # Save verification to session
            self.store.save_verification(self.state.session.session_id, record)
            
            # Update session counters
            self.state.session.manual_input_count += 1
            self.state.session.total_messages += 1
            self.state.message_counter += 1
            
            # Update progress tracker with new total and record verification
            self.progress_tracker.stats.total_messages = self.state.session.total_messages
            self.record_verification_with_timing(is_correct, self.classification_start_time)
            
            # Reload session to get updated counts
            updated_session = self.store.load_session(self.state.session.session_id)
            if updated_session:
                self.state.session = updated_session
            
            # Clear current message state
            self.state.current_message = None
            self.state.current_classification = None
            
            # Get session statistics
            stats = self.store.get_session_statistics(self.state.session.session_id)
            stats["message_counter"] = self.state.message_counter
            
            return True, "✅ Verification saved successfully", stats
            
        except Exception as e:
            return False, f"❌ Error saving verification: {str(e)}", {}
    
    def get_session_results(self) -> List[List[str]]:
        """
        Get all results from the current session.
        
        Returns:
            List of verification records as dictionaries
        """
        if not self.state.session:
            return []
        
        # Reload session to get latest data
        session = self.store.load_session(self.state.session.session_id)
        if not session:
            return []
        
        # Gradio Dataframe renders most reliably with a 2D list (rows) when
        # headers are provided. Returning dicts can show up as "[object Object]"
        # in the browser.
        results: List[List[str]] = []
        for record in session.verifications:
            results.append([
                record.original_message,
                record.classifier_decision.upper(),
                record.ground_truth_label.upper(),
                "✓" if record.is_correct else "✗",
                f"{record.classifier_confidence * 100:.1f}%",
                ", ".join(record.classifier_indicators),
                record.verifier_notes,
                record.timestamp.strftime("%Y-%m-%d %H:%M:%S"),
            ])
        
        return results
    
    def export_session_results(self, format_type: str) -> Tuple[bool, str, Optional[str]]:
        """
        Export session results in specified format.
        
        Args:
            format_type: Export format (csv, json, xlsx)
            
        Returns:
            Tuple of (success, message, file_path_or_content)
        """
        if not self.state.session:
            return False, "❌ No active session to export", None
        
        if self.state.session.verified_count == 0:
            return False, "❌ No verified messages to export", None
        
        try:
            session_id = self.state.session.session_id
            
            if format_type == "csv":
                content = self.store.export_to_csv(session_id)
                filename = f"manual_input_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
                
                # Save to exports directory
                exports_dir = Path("exports")
                exports_dir.mkdir(exist_ok=True)
                file_path = exports_dir / filename
                
                with open(file_path, "w", encoding="utf-8") as f:
                    f.write(content)
                
                return True, f"✅ Results exported to {filename}", str(file_path)
            
            elif format_type == "json":
                content = self.store.export_to_json(session_id)
                filename = f"manual_input_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
                
                # Save to exports directory
                exports_dir = Path("exports")
                exports_dir.mkdir(exist_ok=True)
                file_path = exports_dir / filename
                
                with open(file_path, "w", encoding="utf-8") as f:
                    f.write(content)
                
                return True, f"✅ Results exported to {filename}", str(file_path)
            
            elif format_type == "xlsx":
                content = self.store.export_to_xlsx(session_id)
                filename = f"manual_input_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
                
                # Save to exports directory
                exports_dir = Path("exports")
                exports_dir.mkdir(exist_ok=True)
                file_path = exports_dir / filename
                
                with open(file_path, "wb") as f:
                    f.write(content)
                
                return True, f"✅ Results exported to {filename}", str(file_path)
            
            else:
                return False, f"❌ Unsupported export format: {format_type}", None
                
        except Exception as e:
            return False, f"❌ Error exporting results: {str(e)}", None
    
    def get_enhanced_progress_info(self) -> Dict[str, Any]:
        """
        Get enhanced progress information for display.
        
        Returns:
            Dictionary containing progress information
        """
        if not hasattr(self, 'progress_tracker') or not self.progress_tracker:
            return {
                "progress_display": "📊 Progress: Ready to start",
                "accuracy_display": "🎯 Current Accuracy: No verifications yet",
                "time_display": "⏱️ Time: Not started",
                "error_display": "",
                "stats_summary": "No active session"
            }
        
        return {
            "progress_display": self.progress_tracker.get_progress_display(),
            "accuracy_display": self.progress_tracker.get_accuracy_display(),
            "time_display": self.progress_tracker.get_time_tracking_display(),
            "error_display": self.progress_tracker.get_error_display(),
            "stats_summary": self._get_session_stats_summary()
        }
    
    def record_classification_error(self, error_message: str, can_continue: bool = True) -> None:
        """
        Record a classification error.
        
        Args:
            error_message: Description of the error
            can_continue: Whether processing can continue
        """
        if hasattr(self, 'progress_tracker') and self.progress_tracker:
            self.progress_tracker.record_error(error_message, can_continue)
    
    def pause_manual_session(self) -> Tuple[bool, bool, bool]:
        """
        Pause the current manual input session.
        
        Returns:
            Tuple of control button visibility states
        """
        if hasattr(self, 'progress_tracker') and self.progress_tracker:
            return self.handle_session_pause()
        return False, False, True
    
    def resume_manual_session(self) -> Tuple[bool, bool, bool]:
        """
        Resume the current manual input session.
        
        Returns:
            Tuple of control button visibility states
        """
        if hasattr(self, 'progress_tracker') and self.progress_tracker:
            return self.handle_session_resume()
        return True, False, True
    
    def _get_session_stats_summary(self) -> str:
        """Get formatted session statistics summary."""
        if not self.state.session:
            return "No active session"
        
        # Get latest session stats
        stats = self.store.get_session_statistics(self.state.session.session_id)
        
        return f"""
**Manual Input Session:**
- Messages Processed: {stats.get('verified_count', 0)}
- Accuracy: {stats.get('accuracy', 0):.1f}%
- Correct: {stats.get('correct_count', 0)}
- Incorrect: {stats.get('incorrect_count', 0)}
- Session Duration: {self.progress_tracker.get_time_tracking_display() if hasattr(self, 'progress_tracker') else 'Unknown'}
"""
    
    def complete_session(self) -> Tuple[bool, str]:
        """
        Mark the current session as complete.
        
        Returns:
            Tuple of (success, message)
        """
        if not self.state.session:
            return False, "❌ No active session"
        
        try:
            # Mark session as complete
            self.store.mark_session_complete(self.state.session.session_id)
            
            # Update session state
            self.state.session.is_complete = True
            self.state.session.completed_at = datetime.now()
            
            return True, "✅ Session marked as complete"
            
        except Exception as e:
            return False, f"❌ Error completing session: {str(e)}"


def create_manual_input_interface(model_overrides_state: Optional[gr.State] = None) -> gr.Blocks:
    """
    Create the complete manual input mode interface.
    
    Returns:
        Gradio Blocks component for manual input mode
    """
    controller = ManualInputController()
    if model_overrides_state is None:
        model_overrides_state = gr.State(value={})
    
    with gr.Blocks() as manual_input_interface:
        # Headers and back button are in parent interface
        
        # Session setup section
        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("## 👤 Session Setup")
                verifier_name_input = gr.Textbox(
                    label="Your Name",
                    placeholder="Enter your name to start a session...",
                    interactive=True
                )
            
            with gr.Column(scale=1):
                start_session_btn = StandardizedComponents.create_primary_button(
                    "Start Session",
                    "🚀",
                    "lg"
                )

        # Session info display
        session_info_display = gr.Markdown(
            "Enter your name and click 'Start Session' to begin",
            label="Session Status"
        )

        # Manual input section (initially hidden)
        manual_input_section = gr.Row(visible=False)
        with manual_input_section:
            with gr.Column(scale=2):
                gr.Markdown("## 📝 Message Input")
                
                # Message input area
                message_input = gr.Textbox(
                    label="Patient Message",
                    placeholder="Enter a patient message to classify...",
                    lines=4,
                    interactive=True
                )
                
                # Classification trigger
                classify_btn = StandardizedComponents.create_primary_button(
                    "Classify Message",
                    "🎯",
                    "lg"
                )

                # Apply model overrides right before classification
                def _classify_with_overrides(message_text: str, overrides: Dict[str, str]):
                    controller.set_model_overrides(overrides or {})
                    return controller.classify_message(message_text)
                
                # Classification results (initially hidden)
                classification_results_section = gr.Row(visible=False)
                with classification_results_section:
                    with gr.Column():
                        gr.Markdown("### 🎯 Classification Results")
                        
                        # Classification display
                        classifier_decision_display = gr.Markdown(
                            "",
                            label="Decision"
                        )
                        
                        classifier_confidence_display = gr.Markdown(
                            "",
                            label="Confidence"
                        )
                        
                        classifier_indicators_display = gr.Markdown(
                            "",
                            label="Detected Indicators"
                        )
                        
                        # Verification buttons
                        gr.Markdown("### ✅ Verification")
                        with gr.Row():
                            correct_btn = StandardizedComponents.create_primary_button("Correct", "✓")
                            correct_btn.scale = 1
                            
                            incorrect_btn = StandardizedComponents.create_stop_button("Incorrect", "✗")
                            incorrect_btn.scale = 1
                        
                        # Correction section (initially hidden)
                        correction_section = gr.Row(visible=False)
                        with correction_section:
                            correction_selector = ClassificationDisplay.create_classification_radio()
                            
                            correction_notes = gr.Textbox(
                                label="Notes (Optional)",
                                placeholder="Why is this classification incorrect?",
                                lines=2,
                                interactive=True
                            )
                            
                            submit_correction_btn = StandardizedComponents.create_primary_button(
                                "Submit Correction",
                                "✓"
                            )
            
            with gr.Column(scale=1):
                gr.Markdown("## 📊 Session Statistics")
                
                # Session stats display
                session_stats_display = gr.Markdown(
                    """
                    **Messages Processed:** 0
                    **Correct Classifications:** 0
                    **Incorrect Classifications:** 0
                    **Accuracy:** 0%
                    """,
                    label="Statistics"
                )
                
                # Export options
                gr.Markdown("## 💾 Export Options")
                with gr.Column():
                    # Hidden until there's at least one verified message
                    download_csv_btn = gr.DownloadButton("⬇️ Download CSV", variant="secondary", visible=False)
                    download_json_btn = gr.DownloadButton("⬇️ Download JSON", variant="secondary", visible=False)
                    download_xlsx_btn = gr.DownloadButton("⬇️ Download XLSX", variant="secondary", visible=False)
                
                # Complete session
                gr.Markdown("## 🏁 Session Control")
                complete_session_btn = StandardizedComponents.create_secondary_button(
                    "Complete Session",
                    "🏁",
                    "sm"
                )
        
        # Results history section (initially hidden)
        results_history_section = gr.Row(visible=False)
        with results_history_section:
            with gr.Column():
                gr.Markdown("## 📋 Session Results")
                
                results_display = gr.Dataframe(
                    headers=["Message", "Classifier", "Ground Truth", "Correct", "Confidence", "Indicators", "Notes", "Timestamp"],
                    datatype=["str", "str", "str", "str", "str", "str", "str", "str"],
                    label="Verification Results",
                    interactive=False
                )
        
        # Status messages
        status_message = gr.Markdown("", visible=True)
        
        # Application state
        session_state = gr.State(value=None)
        
        # Event handlers
        def on_start_session(verifier_name):
            """Handle session start."""
            success, message, session = controller.start_new_session(verifier_name)
            
            if success:
                session_info = f"""
                ✅ **Active Session**
                - **Verifier:** {session.verifier_name}
                - **Started:** {session.created_at.strftime('%Y-%m-%d %H:%M:%S')}
                - **Session ID:** {session.session_id[:8]}...
                """
                
                return (
                    session,                        # session_state
                    gr.Row(visible=True),           # manual_input_section
                    gr.Row(visible=True),           # results_history_section
                    session_info,                   # session_info_display
                    gr.DownloadButton(visible=False),  # download_csv_btn
                    gr.DownloadButton(visible=False),  # download_json_btn
                    gr.DownloadButton(visible=False),  # download_xlsx_btn
                    message                         # status_message
                )
            else:
                return (
                    None,                           # session_state
                    gr.Row(visible=False),          # manual_input_section
                    gr.Row(visible=False),          # results_history_section
                    "Enter your name and click 'Start Session' to begin",  # session_info_display
                    gr.DownloadButton(visible=False),  # download_csv_btn
                    gr.DownloadButton(visible=False),  # download_json_btn
                    gr.DownloadButton(visible=False),  # download_xlsx_btn
                    message                         # status_message
                )
        
        def on_classify_message(message_text, overrides):
            """Handle message classification."""
            controller.set_model_overrides(overrides or {})
            success, message, classification = controller.classify_message(message_text)
            
            if success:
                # Format classification results using standardized components
                decision_badge = ClassificationDisplay.format_classification_badge(classification['decision'])
                confidence_text = ClassificationDisplay.format_confidence_display(classification['confidence'])
                indicators_text = ClassificationDisplay.format_indicators_display(classification['indicators'])
                
                return (
                    gr.Row(visible=True),           # classification_results_section
                    decision_badge,                 # classifier_decision_display
                    confidence_text,                # classifier_confidence_display
                    indicators_text,                # classifier_indicators_display
                    message                         # status_message
                )
            else:
                return (
                    gr.Row(visible=False),          # classification_results_section
                    "",                             # classifier_decision_display
                    "",                             # classifier_confidence_display
                    "",                             # classifier_indicators_display
                    message                         # status_message
                )
        
        def on_correct_verification():
            """Handle correct classification verification."""
            success, message, stats = controller.submit_verification(True)
            
            if success:
                # Update stats display using standardized formatting
                stats_text = SessionDisplay.format_session_statistics(stats)
                
                # Get updated results
                results = controller.get_session_results()
                
                return (
                    "",                             # message_input (clear)
                    gr.Row(visible=False),          # classification_results_section
                    gr.Row(visible=False),          # correction_section
                    stats_text,                     # session_stats_display
                    results,                        # results_display
                    gr.DownloadButton(visible=True),   # download_csv_btn
                    gr.DownloadButton(visible=True),   # download_json_btn
                    gr.DownloadButton(visible=True),   # download_xlsx_btn
                    message                         # status_message
                )
            else:
                return (
                    gr.Textbox(value=""),           # message_input (no change)
                    gr.Row(visible=True),           # classification_results_section (no change)
                    gr.Row(visible=False),          # correction_section
                    gr.Markdown(value=""),          # session_stats_display (no change)
                    gr.Dataframe(value=[]),         # results_display (no change)
                    gr.DownloadButton(),             # download_csv_btn (no change)
                    gr.DownloadButton(),             # download_json_btn (no change)
                    gr.DownloadButton(),             # download_xlsx_btn (no change)
                    message                         # status_message
                )
        
        def on_incorrect_verification():
            """Handle incorrect classification - show correction options."""
            return (
                gr.Row(visible=True),               # correction_section
                "Please select the correct classification and submit"  # status_message
            )
        
        def on_submit_correction(correction, notes):
            """Handle correction submission."""
            success, message, stats = controller.submit_verification(False, correction, notes)
            
            if success:
                # Update stats display using standardized formatting
                stats_text = SessionDisplay.format_session_statistics(stats)
                
                # Get updated results
                results = controller.get_session_results()
                
                return (
                    "",                             # message_input (clear)
                    gr.Row(visible=False),          # classification_results_section
                    gr.Row(visible=False),          # correction_section
                    "",                             # correction_notes (clear)
                    stats_text,                     # session_stats_display
                    results,                        # results_display
                    gr.DownloadButton(visible=True),   # download_csv_btn
                    gr.DownloadButton(visible=True),   # download_json_btn
                    gr.DownloadButton(visible=True),   # download_xlsx_btn
                    message                         # status_message
                )
            else:
                return (
                    gr.Textbox(value=""),           # message_input (no change)
                    gr.Row(visible=True),           # classification_results_section (no change)
                    gr.Row(visible=True),           # correction_section (keep visible)
                    notes,                          # correction_notes (keep)
                    gr.Markdown(value=""),          # session_stats_display (no change)
                    gr.Dataframe(value=[]),         # results_display (no change)
                    gr.DownloadButton(),             # download_csv_btn (no change)
                    gr.DownloadButton(),             # download_json_btn (no change)
                    gr.DownloadButton(),             # download_xlsx_btn (no change)
                    message                         # status_message
                )
        
        def on_export_results_file(format_type):
            """Handle results export for DownloadButton (returns file path)."""
            success, message, file_path = controller.export_session_results(format_type)
            if success and file_path:
                return file_path
            # Returning None tells DownloadButton there's nothing to download.
            return None
        
        def on_complete_session():
            """Handle session completion."""
            success, message = controller.complete_session()
            
            if success:
                # Get final results
                results = controller.get_session_results()
                final_stats = controller.store.get_session_statistics(controller.state.session.session_id)
                
                completion_message = f"""
                🏁 **Session Completed Successfully**
                
                **Final Statistics:**
                - Messages Processed: {final_stats['verified_count']}
                - Accuracy: {final_stats['accuracy']:.1f}%
                - Correct: {final_stats['correct_count']}
                - Incorrect: {final_stats['incorrect_count']}
                
                You can now export your results or start a new session.
                """
                
                return (
                    gr.Row(visible=False),          # manual_input_section
                    completion_message,             # session_info_display
                    message                         # status_message
                )
            else:
                return (
                    gr.Row(visible=True),           # manual_input_section (no change)
                    gr.Markdown(value=""),          # session_info_display (no change)
                    message                         # status_message
                )
        
        # Bind event handlers
        start_session_btn.click(
            on_start_session,
            inputs=[verifier_name_input],
            outputs=[
                session_state,
                manual_input_section,
                results_history_section,
                session_info_display,
                download_csv_btn,
                download_json_btn,
                download_xlsx_btn,
                status_message
            ]
        )
        
        classify_btn.click(
            on_classify_message,
            inputs=[message_input, model_overrides_state],
            outputs=[
                classification_results_section,
                classifier_decision_display,
                classifier_confidence_display,
                classifier_indicators_display,
                status_message
            ]
        )
        
        correct_btn.click(
            on_correct_verification,
            outputs=[
                message_input,
                classification_results_section,
                correction_section,
                session_stats_display,
                results_display,
                download_csv_btn,
                download_json_btn,
                download_xlsx_btn,
                status_message
            ]
        )
        
        incorrect_btn.click(
            on_incorrect_verification,
            outputs=[correction_section, status_message]
        )
        
        submit_correction_btn.click(
            on_submit_correction,
            inputs=[correction_selector, correction_notes],
            outputs=[
                message_input,
                classification_results_section,
                correction_section,
                correction_notes,
                session_stats_display,
                results_display,
                download_csv_btn,
                download_json_btn,
                download_xlsx_btn,
                status_message
            ]
        )
        
        download_csv_btn.click(
            lambda: on_export_results_file("csv"),
            outputs=[download_csv_btn]
        )
        
        download_json_btn.click(
            lambda: on_export_results_file("json"),
            outputs=[download_json_btn]
        )
        
        download_xlsx_btn.click(
            lambda: on_export_results_file("xlsx"),
            outputs=[download_xlsx_btn]
        )
        
        complete_session_btn.click(
            on_complete_session,
            outputs=[
                manual_input_section,
                session_info_display,
                status_message
            ]
        )
    
    return manual_input_interface