# simplified_gradio_app.py
"""
Simplified Gradio Interface based on customer feedback from Or_4.txt

Key changes:
1. Provider Summary is the final exchange in Conversation Verification tab
2. Streamlined interface focusing on essential features
3. CSV export for verification results

Requirements: Customer feedback Or_4.txt
"""

import os
import sys

# Ensure project root is in Python path
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if project_root not in sys.path:
    sys.path.insert(0, project_root)

from dotenv import load_dotenv

# Load environment variables
load_dotenv()

import gradio as gr

# Import modularized components
from src.interface.session_manager import SimplifiedSessionData
from src.interface import stats_handlers
from src.interface import simplified_chat_handlers  # Use simplified handlers
from src.interface import verification_handlers
from src.interface import prompt_handlers
from src.interface import model_handlers
from src.interface import profile_handlers
from src.interface.simplified_help_content import SIMPLIFIED_HELP_CONTENT

try:
    from app_config import GRADIO_CONFIG
except ImportError:
    GRADIO_CONFIG = {"theme": "soft", "show_api": False}


def create_simplified_interface():
    """
    Create simplified Gradio interface based on customer feedback.
    
    Focus on:
    - Chat interface
    - Conversation Verification with Provider Summary as final step
    - Simplified data export
    """
    
    debug_mode = os.getenv("LOG_PROMPTS", "false").lower() == "true"
    
    # Theme setup
    theme_name = GRADIO_CONFIG.get("theme", "soft")
    if theme_name.lower() == "soft":
        theme = gr.themes.Soft()
    else:
        theme = gr.themes.Default()
    
    demo = gr.Blocks(
        title="🏥 Medical Assistant - Simplified",
        analytics_enabled=False
    )
    demo.theme = theme
    
    with demo:
        # Session state
        session_data = gr.State(value=None)
        
        # Header
        gr.Markdown("# 🏥 Medical Assistant with Spiritual Support")
        gr.Markdown("Simplified interface for conversation verification and testing")
        
        if debug_mode:
            gr.Markdown("⚠️ **DEBUG MODE:** Prompts and responses are logged")
        
        # Session info
        session_info = gr.Markdown("🔄 **Initializing session...**")
        
        # Initialize session
        def initialize_session():
            """Initialize new user session."""
            new_session = SimplifiedSessionData()
            session_info_text = f"""
✅ **Session Ready**
🆔 Session: `{new_session.session_id[:8]}...`
🕒 Started: {new_session.created_at[:19]}
"""
            return new_session, session_info_text
        
        # Main interface tabs
        main_tabs = gr.Tabs(elem_id="main_tabs")
        with main_tabs:
            
            # Chat tab
            with gr.TabItem("💬 Chat", id="chat"):
                with gr.Row():
                    with gr.Column(scale=2):
                        chatbot = gr.Chatbot(
                            label="💬 Conversation",
                            height=450
                        )
                        
                        with gr.Row():
                            msg = gr.Textbox(
                                label="Your message",
                                placeholder="Type your health question...",
                                scale=4
                            )
                            send_btn = gr.Button("📤 Send", scale=1, variant="primary")
                        
                        with gr.Row():
                            clear_btn = gr.Button("🗑️ Clear Chat", scale=1)
                        
                        # Quick examples
                        gr.Markdown("### ⚡ Quick Start:")
                        with gr.Row():
                            example_medical = gr.Button("🟢 I am fine", size="sm")
                            example_wellness = gr.Button("🟡 I'm feeling stressed", size="sm")
                            example_help = gr.Button("🔴 Emotional crisis", size="sm")
                    
                    with gr.Column(scale=1):
                        status_box = gr.Markdown(
                            value="🔄 Loading...",
                            label="📊 Status"
                        )
                        
                        refresh_btn = gr.Button("🔄 Check Status", size="sm")
                        
                        # Conversation statistics
                        gr.Markdown("### 📈 Conversation Stats")
                        conversation_stats = gr.Markdown(
                            value="No conversation yet",
                            label="Statistics"
                        )

            # Conversation Verification tab - MAIN FOCUS
            with gr.TabItem("🧾 Conversation Verification", id="conversation_verification"):
                gr.Markdown("## 🧾 Conversation Verification")
                gr.Markdown(
                    "Review each exchange from the chat conversation. "
                    "**The Provider Summary will be shown as the final step** for your review and feedback."
                )

                conv_verify_state = gr.State(value=None)
                conv_verify_index = gr.State(value=0)
                conv_verify_records = gr.State(value=[])

                with gr.Row():
                    generate_conv_verification_btn = gr.Button(
                        "🛠 Generate Verification from Chat", 
                        variant="primary"
                    )

                conv_verify_status = gr.Markdown(value="", visible=True)
                conv_verify_exchange = gr.HTML(value="", label="Current Exchange")

                with gr.Row():
                    conv_correct_btn = gr.Button("✅ Correct", variant="primary")
                    conv_incorrect_btn = gr.Button("❌ Incorrect")
                    conv_prev_btn = gr.Button("⬅️ Previous")
                    conv_next_btn = gr.Button("Next ➡️")

                # Shown only when marking Incorrect
                with gr.Row(visible=False) as conv_incorrect_comment_row:
                    with gr.Column(scale=3):
                        gr.Markdown("### Select Correct Classification:")
                        conv_correct_classification = gr.Radio(
                            choices=[
                                "🟢 Should be GREEN - No distress",
                                "🟡 Should be YELLOW - Needs clarification", 
                                "🔴 Should be RED - Spiritual distress"
                            ],
                            label="Correct Classification",
                            interactive=True
                        )
                        conv_incorrect_comment = gr.Textbox(
                            label="Comment (why incorrect / what to fix)",
                            placeholder="Add your feedback here...",
                            lines=3,
                        )
                    with gr.Column(scale=1):
                        conv_save_comment_btn = gr.Button("💾 Save comment", variant="secondary")

                with gr.Row():
                    with gr.Column(scale=1):
                        conv_position = gr.Markdown(value="")
                    with gr.Column(scale=1):
                        conv_stats = gr.HTML(value="")
                
                # Manual download options (backup)
                gr.Markdown("### 📥 Manual Export")
                with gr.Row():
                    conv_verify_download_csv_btn = gr.DownloadButton("📄 Download CSV", variant="secondary")

            # Model Selection tab
            with gr.TabItem("⚙️ Model Settings", id="model_settings"):
                gr.Markdown("## ⚙️ AI Model Configuration")
                gr.Markdown("Select which AI models to use for different tasks. Changes apply to your current session.")
                
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### 🔍 Spiritual Monitor (Classifier)")
                        spiritual_model = gr.Dropdown(
                            choices=[
                                "gemini-2.5-flash",
                                "gemini-2.0-flash",
                                "gemini-3-flash-preview",
                                "claude-sonnet-4-5-20250929",
                                "claude-sonnet-4-20250514",
                                "claude-3-7-sonnet-20250219"
                            ],
                            value="gemini-2.5-flash",
                            label="Spiritual Distress Analyzer",
                            interactive=True
                        )
                        
                        gr.Markdown("### 🟡 Soft Spiritual Triage")
                        soft_spiritual_triage_model = gr.Dropdown(
                            choices=[
                                "claude-sonnet-4-5-20250929",
                                "claude-sonnet-4-20250514",
                                "claude-3-7-sonnet-20250219",
                                "gemini-2.5-flash",
                                "gemini-2.0-flash",
                                "gemini-3-flash-preview"
                            ],
                            value="claude-sonnet-4-5-20250929",
                            label="Soft Spiritual Triage",
                            interactive=True
                        )
                    
                    with gr.Column():
                        gr.Markdown("### 📊 Triage Response Evaluator")
                        triage_evaluate_model = gr.Dropdown(
                            choices=[
                                "gemini-2.5-flash",
                                "gemini-2.0-flash",
                                "gemini-3-flash-preview",
                                "claude-sonnet-4-5-20250929",
                                "claude-sonnet-4-20250514",
                                "claude-3-7-sonnet-20250219"
                            ],
                            value="gemini-2.5-flash",
                            label="Triage Response Evaluator",
                            interactive=True
                        )
                        
                        gr.Markdown("### 🏥 Medical Assistant")
                        medical_model = gr.Dropdown(
                            choices=[
                                "claude-sonnet-4-5-20250929",
                                "claude-sonnet-4-20250514",
                                "claude-3-7-sonnet-20250219",
                                "gemini-2.5-flash",
                                "gemini-2.0-flash",
                                "gemini-3-flash-preview"
                            ],
                            value="claude-sonnet-4-5-20250929",
                            label="Medical Assistant",
                            interactive=True
                        )
                    
                    with gr.Column():
                        gr.Markdown("### 🩺 Soft Medical Triage")
                        soft_triage_model = gr.Dropdown(
                            choices=[
                                "claude-sonnet-4-5-20250929",
                                "claude-sonnet-4-20250514",
                                "claude-3-7-sonnet-20250219",
                                "gemini-2.5-flash",
                                "gemini-2.0-flash",
                                "gemini-3-flash-preview"
                            ],
                            value="claude-sonnet-4-5-20250929",
                            label="Soft Medical Triage",
                            interactive=True
                        )
                        
                        gr.Markdown("### 💬 Medical Brain Summary Generator")
                        spiritual_care_message_model = gr.Dropdown(
                            choices=[
                                "claude-sonnet-4-5-20250929",
                                "claude-sonnet-4-20250514",
                                "claude-3-7-sonnet-20250219",
                                "gemini-2.5-flash",
                                "gemini-2.0-flash",
                                "gemini-3-flash-preview"
                            ],
                            value="claude-sonnet-4-5-20250929",
                            label="Medical Brain Summary Generator (uses Spiritual Care Message prompt)",
                            interactive=True
                        )
                
                with gr.Row():
                    apply_models_btn = gr.Button("✅ Apply Model Settings", variant="primary", scale=2)
                    reset_models_btn = gr.Button("🔄 Reset to Defaults", scale=1)
                
                model_status = gr.HTML(value="", visible=True)
            
            # Edit Prompts tab
            with gr.TabItem("🔧 Edit Prompts", id="edit_prompts"):
                gr.Markdown("## 🔧 Customize AI Prompts")
                gr.Markdown("⚠️ **Note:** Changes apply only to your current session.")
                
                # Prompt selector
                with gr.Row():
                    prompt_selector = gr.Dropdown(
                        choices=[
                            "🔍 Spiritual Monitor (Classifier)",
                            "🟡 Soft Spiritual Triage",
                            "📊 Triage Response Evaluator",
                            "🏥 Medical Assistant",
                            "🩺 Soft Medical Triage",
                            "💬 Spiritual Care Message"
                        ],
                        value="🔍 Spiritual Monitor (Classifier)",
                        label="Select Prompt to Edit",
                        interactive=True
                    )
                
                with gr.Row():
                    with gr.Column(scale=3):
                        # Prompt editor
                        prompt_editor = gr.Code(
                            label="System Prompt",
                            value="",
                            language="markdown",
                            lines=25,
                            interactive=True
                        )
                        
                        with gr.Row():
                            load_prompt_btn = gr.Button("📥 Load Prompt", scale=1)
                            apply_prompt_btn = gr.Button("✅ Apply Changes", variant="primary", scale=2)
                            reset_prompt_btn = gr.Button("🔄 Reset to Default", variant="secondary", scale=1)
                        
                        with gr.Row():
                            promote_prompt_btn = gr.Button("📤 Promote to File", variant="stop", scale=1)
                            validate_prompt_btn = gr.Button("🔍 Validate", variant="secondary", scale=1)
                        
                        prompt_status = gr.HTML(
                            value="", 
                            visible=True,
                            elem_classes=["prompt-status-container"]
                        )
                    
                    with gr.Column(scale=1):
                        gr.Markdown("### 📋 Prompt Info")
                        prompt_info_display = gr.HTML(value="""
<div style="font-family: system-ui; padding: 1em; background-color: #f9fafb; border-radius: 8px;">
<p><strong>Select a prompt to edit</strong></p>

<p><strong>Available prompts:</strong></p>
<ul style="margin-left: 1em;">
<li>🔍 Spiritual Monitor</li>
<li>🟡 Triage Questions</li>
<li>📊 Triage Evaluation</li>
<li>🏥 Medical Assistant</li>
<li>🩺 Soft Triage</li>
<li>💬 Spiritual Care Message (used for Medical Brain Summary)</li>
</ul>

<p><strong>Tips:</strong></p>
<ul style="margin-left: 1em;">
<li>Load prompt first</li>
<li>Edit carefully</li>
<li>Test changes</li>
<li>Reset if needed</li>
</ul>
</div>
                        """)
            
            # Patient Profiles tab
            with gr.TabItem("👥 Patient Profiles", id="profiles"):
                gr.Markdown("## 👥 Patient Profile Management")
                gr.Markdown("Select a predefined profile or customize the current patient settings.")
                
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### 📋 Predefined Profiles")
                        profile_selector = gr.Dropdown(
                            choices=[
                                "👤 Default Profile (Serhii)",
                                "🟢 GREEN - Healthy Coping",
                                "🟡 YELLOW - Mild Distress",
                                "🟡 YELLOW - Grief & Loss",
                                "🟡 YELLOW - Existential Questions",
                                "🟡 YELLOW - Spiritual Disconnection",
                                "🔴 RED - Crisis (Suicidal Risk)",
                                "🔴 RED - Severe Hopelessness",
                                "🔴 RED - Spiritual Crisis",
                                "🫀 Cardiac Patient (Rehabilitation)",
                                "🩸 Diabetic Patient (Management)",
                                "🏥 Post-Surgery (Recovery)",
                                "🧠 Mental Health (Anxiety/Depression)",
                                "👴 Elderly Patient (Chronic Care)",
                                "🏃 Athletic Patient (Injury/Training)"
                            ],
                            value="👤 Default Profile (Serhii)",
                            label="Select Profile",
                            interactive=True
                        )
                        
                        load_profile_btn = gr.Button("📥 Load Profile", variant="primary")
                        profile_status = gr.HTML(value="", visible=True)
                    
                    with gr.Column(scale=2):
                        gr.Markdown("### ⚙️ Current Patient Settings")
                        
                        with gr.Row():
                            patient_name = gr.Textbox(
                                label="Patient Name",
                                value="Serhii",
                                interactive=True
                            )
                            patient_phone = gr.Textbox(
                                label="Phone Number",
                                value="",
                                placeholder="(555) 123-4567",
                                interactive=True
                            )
                            patient_age = gr.Number(
                                label="Age",
                                value=52,
                                interactive=True
                            )
                        
                        with gr.Row():
                            conditions = gr.Textbox(
                                label="Medical Conditions (comma-separated)",
                                value="Atrial fibrillation, Deep vein thrombosis, Obesity, Hypertension",
                                lines=3,
                                interactive=True
                            )
                        
                        with gr.Row():
                            primary_goal = gr.Textbox(
                                label="Primary Goal",
                                value="Weight reduction and cardiovascular fitness improvement",
                                lines=2,
                                interactive=True
                            )
                        
                        with gr.Row():
                            exercise_prefs = gr.Textbox(
                                label="Exercise Preferences (comma-separated)",
                                value="Swimming, Walking, Light cardio",
                                lines=2,
                                interactive=True
                            )
                        
                        with gr.Row():
                            exercise_limits = gr.Textbox(
                                label="Exercise Limitations (comma-separated)",
                                value="Anticoagulation therapy, Post-thrombotic recovery",
                                lines=2,
                                interactive=True
                            )
                        
                        with gr.Row():
                            save_profile_btn = gr.Button("💾 Save Current Profile", variant="primary", scale=1)
                            reset_profile_btn = gr.Button("🔄 Reset to Default", scale=1)
                        
                        profile_save_status = gr.HTML(value="", visible=True)

            # Help tab
            with gr.TabItem("📖 Help", id="help"):
                gr.Markdown(SIMPLIFIED_HELP_CONTENT)

        
        # Event handlers
        
        # Initialize session
        demo.load(
            initialize_session,
            outputs=[session_data, session_info]
        )
        
        # Send message
        send_btn.click(
            simplified_chat_handlers.handle_message_simplified,
            inputs=[msg, chatbot, session_data],
            outputs=[chatbot, status_box, session_data, msg, conversation_stats]
        )
        
        msg.submit(
            simplified_chat_handlers.handle_message_simplified,
            inputs=[msg, chatbot, session_data],
            outputs=[chatbot, status_box, session_data, msg, conversation_stats]
        )
        
        # Clear chat
        clear_btn.click(
            simplified_chat_handlers.handle_clear_simplified,
            inputs=[session_data],
            outputs=[chatbot, status_box, session_data]
        )
        
        # Refresh status
        refresh_btn.click(
            stats_handlers.get_status,
            inputs=[session_data],
            outputs=[status_box, conversation_stats]
        )
        
        # Example buttons
        example_medical.click(
            lambda h, s: simplified_chat_handlers.send_example_simplified("I am fine", h, s),
            inputs=[chatbot, session_data],
            outputs=[chatbot, status_box, session_data, msg, conversation_stats]
        )
        
        example_wellness.click(
            lambda h, s: simplified_chat_handlers.send_example_simplified("I'm feeling stressed and overwhelmed lately", h, s),
            inputs=[chatbot, session_data],
            outputs=[chatbot, status_box, session_data, msg, conversation_stats]
        )
        
        example_help.click(
            lambda h, s: simplified_chat_handlers.send_example_simplified("I am currently experiencing an emotional crisis", h, s),
            inputs=[chatbot, session_data],
            outputs=[chatbot, status_box, session_data, msg, conversation_stats]
        )

        # Conversation Verification events
        generate_conv_verification_btn.click(
            verification_handlers._generate_conv_verification_with_summary,
            inputs=[session_data],
            outputs=[conv_verify_state, conv_verify_records, conv_verify_index, conv_verify_status, conv_verify_exchange, conv_position, conv_stats]
        )

        conv_verify_download_csv_btn.click(
            verification_handlers._download_reviewed_csv,
            inputs=[conv_verify_state, conv_verify_records],
            outputs=[conv_verify_download_csv_btn]
        )

        conv_correct_btn.click(
            verification_handlers._mark_conv_correct,
            inputs=[conv_verify_records, conv_verify_index],
            outputs=[
                conv_verify_records,
                conv_verify_index,
                conv_verify_status,
                conv_verify_exchange,
                conv_position,
                conv_stats,
                conv_incorrect_comment_row,
                conv_incorrect_comment,
                conv_correct_classification,
            ]
        )

        conv_incorrect_btn.click(
            verification_handlers._show_incorrect_comment_ui,
            inputs=[conv_verify_records, conv_verify_index],
            outputs=[
                conv_verify_records,
                conv_verify_index,
                conv_verify_status,
                conv_verify_exchange,
                conv_position,
                conv_stats,
                conv_incorrect_comment_row,
                conv_incorrect_comment,
                conv_correct_classification,
            ]
        )

        conv_save_comment_btn.click(
            verification_handlers._save_incorrect_comment,
            inputs=[conv_verify_records, conv_verify_index, conv_incorrect_comment, conv_correct_classification],
            outputs=[
                conv_verify_records,
                conv_verify_index,
                conv_verify_status,
                conv_verify_exchange,
                conv_position,
                conv_stats,
                conv_incorrect_comment_row,
                conv_incorrect_comment,
                conv_correct_classification,
            ]
        )

        conv_prev_btn.click(
            lambda records, idx: verification_handlers._nav_conv(records, idx, -1),
            inputs=[conv_verify_records, conv_verify_index],
            outputs=[conv_verify_index, conv_verify_exchange, conv_position, conv_stats, conv_incorrect_comment_row, conv_incorrect_comment, conv_correct_classification]
        )

        conv_next_btn.click(
            lambda records, idx: verification_handlers._nav_conv(records, idx, 1),
            inputs=[conv_verify_records, conv_verify_index],
            outputs=[conv_verify_index, conv_verify_exchange, conv_position, conv_stats, conv_incorrect_comment_row, conv_incorrect_comment, conv_correct_classification]
        )

        # Prompt editing events
        load_prompt_btn.click(
            prompt_handlers.load_prompt,
            inputs=[prompt_selector, session_data],
            outputs=[prompt_editor, prompt_info_display, prompt_status]
        )
        
        apply_prompt_btn.click(
            prompt_handlers.apply_prompt_changes,
            inputs=[prompt_selector, prompt_editor, session_data],
            outputs=[prompt_status, session_data]
        )
        
        reset_prompt_btn.click(
            prompt_handlers.reset_prompt,
            inputs=[prompt_selector, session_data],
            outputs=[prompt_editor, prompt_info_display, prompt_status, session_data]
        )
        
        promote_prompt_btn.click(
            prompt_handlers.promote_prompt_to_file,
            inputs=[prompt_selector, session_data],
            outputs=[prompt_status, session_data]
        )
        
        validate_prompt_btn.click(
            prompt_handlers.validate_prompt_syntax,
            inputs=[prompt_editor],
            outputs=[prompt_status]
        )
        
        # Auto-load prompt when selector changes
        prompt_selector.change(
            prompt_handlers.load_prompt,
            inputs=[prompt_selector, session_data],
            outputs=[prompt_editor, prompt_info_display, prompt_status]
        )
        
        # Bind model selection events
        apply_models_btn.click(
            model_handlers.apply_model_settings,
            inputs=[spiritual_model, soft_spiritual_triage_model, triage_evaluate_model, medical_model, soft_triage_model, spiritual_care_message_model, session_data],
            outputs=[model_status, session_data]
        )
        
        reset_models_btn.click(
            model_handlers.reset_model_settings,
            inputs=[session_data],
            outputs=[model_status, session_data]
        )
        
        # Bind profile events
        load_profile_btn.click(
            profile_handlers.load_profile,
            inputs=[profile_selector, session_data],
            outputs=[patient_name, patient_phone, patient_age, conditions, primary_goal, exercise_prefs, exercise_limits, profile_status]
        )
        
        save_profile_btn.click(
            profile_handlers.save_profile,
            inputs=[patient_name, patient_phone, patient_age, conditions, primary_goal, exercise_prefs, exercise_limits, session_data],
            outputs=[profile_save_status]
        )
        
        reset_profile_btn.click(
            profile_handlers.reset_profile,
            inputs=[session_data],
            outputs=[patient_name, patient_phone, patient_age, conditions, primary_goal, exercise_prefs, exercise_limits, profile_save_status]
        )
    
    # Add CSS for prompt status container
    demo.css = """
    .prompt-status-container {
        max-height: 300px !important;
        overflow-y: auto !important;
        margin: 0.5em 0 !important;
    }
    .prompt-status-container > div {
        max-height: 280px !important;
        overflow-y: auto !important;
    }
    """
    
    return demo


def main():
    """Launch the simplified Gradio interface."""
    demo = create_simplified_interface()
    
    # Get configuration
    server_name = os.getenv("GRADIO_SERVER_NAME", "0.0.0.0")
    server_port = int(os.getenv("GRADIO_SERVER_PORT", "7860"))  # Use standard port for HF Spaces
    share = os.getenv("GRADIO_SHARE", "false").lower() == "true"
    
    print(f"🚀 Starting Simplified Medical Assistant Interface...")
    print(f"📍 Server: http://{server_name}:{server_port}")
    print(f"📋 Based on customer feedback: Or_4.txt")
    
    demo.launch(
        server_name=server_name,
        server_port=server_port,
        share=share
    )


if __name__ == "__main__":
    main()