"""
app_fixed.py
-----------

Simplified version of the ID Agents app with working Gradio authentication.
This version removes the complex custom authentication that was causing initialization failures.
"""

# --- Imports ---
import gradio as gr
import json
import re
import os
import asyncio
import logging
from typing import Dict, Optional, Any, cast

# Try to import OpenAI
try:
    import openai
    from openai import RateLimitError, APIError, APIConnectionError, OpenAI
    OPENAI_AVAILABLE = True
    print("✅ OpenAI client loaded")
except ImportError as e:
    print(f"⚠️ OpenAI not available: {e}")
    OPENAI_AVAILABLE = False

# Try to import core modules
try:
    from orchestrator import AIOrchestrator
    from llm_connector import LLMConnector
    from rag import RAGRetriever
    CORE_MODULES_AVAILABLE = True
    print("✅ Core modules loaded")
except ImportError as e:
    print(f"⚠️ Core modules not available: {e}")
    CORE_MODULES_AVAILABLE = False

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

# --- Global Variables ---
current_llm_connector = None
current_orchestrator = None
current_rag_retriever = None

# --- Helper Functions ---

def initialize_components():
    """Initialize core components if available"""
    global current_llm_connector, current_orchestrator, current_rag_retriever
    
    if not CORE_MODULES_AVAILABLE:
        return False
    
    try:
        # Initialize LLM Connector
        current_llm_connector = LLMConnector()
        
        # Initialize RAG Retriever
        current_rag_retriever = RAGRetriever()
        
        # Initialize Orchestrator
        current_orchestrator = AIOrchestrator(current_llm_connector, current_rag_retriever)
        
        print("✅ All components initialized successfully")
        return True
    except Exception as e:
        print(f"❌ Component initialization failed: {e}")
        return False

def safe_chat_response(message: str, history: list) -> tuple:
    """Safe chat response with fallback"""
    if not CORE_MODULES_AVAILABLE or not current_orchestrator:
        response = """
        🔧 **System Status**: Core modules are not fully available.
        
        This is a demo version of ID Agents. In the full version, you would have access to:
        
        • **AI-Powered Clinical Assistance**: Advanced infectious disease consultation
        • **Multi-Agent Orchestration**: Coordinated responses from specialized agents
        • **Knowledge Base Integration**: Access to medical literature and guidelines
        • **Patient Scenario Analysis**: Context-aware clinical decision support
        
        **Available Test Accounts:**
        • dr_smith / idweek2025 (ID Physician)
        • id_fellow / hello (ID Fellow)
        • pharmacist / stewardship (Clinical Pharmacist)
        • ipc_nurse / infection (IPC Coordinator)
        • researcher / research (Clinical Researcher)
        
        Please contact the administrator for full access.
        """
        return response, history + [[message, response]]
    
    try:
        # Get response from orchestrator
        response = current_orchestrator.process_query(message)
        
        # Update history
        updated_history = history + [[message, response]]
        
        return response, updated_history
    
    except Exception as e:
        error_response = f"❌ Error processing your request: {str(e)}"
        return error_response, history + [[message, error_response]]

def convert_messages_for_gradio(messages):
    """Convert messages to Gradio 4.20.0 format"""
    if not messages:
        return []
    
    converted = []
    for msg in messages:
        if isinstance(msg, dict) and 'role' in msg and 'content' in msg:
            if msg['role'] == 'user':
                converted.append([msg['content'], None])
            elif msg['role'] == 'assistant' and converted:
                converted[-1][1] = msg['content']
            elif msg['role'] == 'assistant':
                converted.append([None, msg['content']])
        elif isinstance(msg, list) and len(msg) == 2:
            converted.append(msg)
    
    return converted

# Patient loading functions (simplified)
def load_patient_1():
    patient_data = {
        "name": "Maria Santos",
        "age": 67,
        "gender": "Female",
        "chief_complaint": "Fever and shortness of breath",
        "history": "3 days of progressive fever (up to 101.8°F), productive cough with yellow sputum, and increasing shortness of breath...",
        "vitals": "Temp: 101.8°F, HR: 110, BP: 145/85, RR: 22, O2 Sat: 89% on room air",
        "presentation": "Elderly female with fever, productive cough, and hypoxemia concerning for pneumonia"
    }
    
    context_message = f"""
🏥 **Patient Case Loaded: {patient_data['name']}**

**Demographics:** {patient_data['age']}-year-old {patient_data['gender']}
**Chief Complaint:** {patient_data['chief_complaint']}
**Current Vitals:** {patient_data['vitals']}

**Clinical Presentation:** {patient_data['presentation']}

How can I assist with this case?
"""
    
    history = [[f"Load patient case: {patient_data['name']}", context_message]]
    return convert_messages_for_gradio([]), history

def load_patient_2():
    patient_data = {
        "name": "James Wilson",
        "age": 34,
        "gender": "Male", 
        "chief_complaint": "Painful urination and fever",
        "presentation": "Young male with UTI symptoms and systemic signs of infection"
    }
    
    context_message = f"""
🏥 **Patient Case Loaded: {patient_data['name']}**

**Demographics:** {patient_data['age']}-year-old {patient_data['gender']}
**Chief Complaint:** {patient_data['chief_complaint']}

**Clinical Presentation:** {patient_data['presentation']}

How can I assist with this case?
"""
    
    history = [[f"Load patient case: {patient_data['name']}", context_message]]
    return convert_messages_for_gradio([]), history

def load_patient_3():
    patient_data = {
        "name": "Sarah Chen",
        "age": 28,
        "gender": "Female",
        "chief_complaint": "Skin rash and joint pain", 
        "presentation": "Young female with systemic symptoms and dermatologic findings"
    }
    
    context_message = f"""
🏥 **Patient Case Loaded: {patient_data['name']}**

**Demographics:** {patient_data['age']}-year-old {patient_data['gender']}
**Chief Complaint:** {patient_data['chief_complaint']}

**Clinical Presentation:** {patient_data['presentation']}

How can I assist with this case?
"""
    
    history = [[f"Load patient case: {patient_data['name']}", context_message]]
    return convert_messages_for_gradio([]), history

def load_patient_4():
    patient_data = {
        "name": "Robert Martinez",
        "age": 45,
        "gender": "Male",
        "chief_complaint": "Persistent cough and weight loss",
        "presentation": "Middle-aged male with chronic respiratory symptoms and constitutional signs"
    }
    
    context_message = f"""
🏥 **Patient Case Loaded: {patient_data['name']}**

**Demographics:** {patient_data['age']}-year-old {patient_data['gender']}
**Chief Complaint:** {patient_data['chief_complaint']}

**Clinical Presentation:** {patient_data['presentation']}

How can I assist with this case?
"""
    
    history = [[f"Load patient case: {patient_data['name']}", context_message]]
    return convert_messages_for_gradio([]), history

def load_patient_5():
    patient_data = {
        "name": "Emma Thompson", 
        "age": 19,
        "gender": "Female",
        "chief_complaint": "Headache and neck stiffness",
        "presentation": "Young female with meningeal signs requiring urgent evaluation"
    }
    
    context_message = f"""
🏥 **Patient Case Loaded: {patient_data['name']}**

**Demographics:** {patient_data['age']}-year-old {patient_data['gender']}
**Chief Complaint:** {patient_data['chief_complaint']}

**Clinical Presentation:** {patient_data['presentation']}

How can I assist with this case?
"""
    
    history = [[f"Load patient case: {patient_data['name']}", context_message]]
    return convert_messages_for_gradio([]), history

def load_patient_6():
    patient_data = {
        "name": "Michael Brown",
        "age": 72,
        "gender": "Male", 
        "chief_complaint": "Abdominal pain and diarrhea",
        "presentation": "Elderly male with GI symptoms and potential infectious etiology"
    }
    
    context_message = f"""
🏥 **Patient Case Loaded: {patient_data['name']}**

**Demographics:** {patient_data['age']}-year-old {patient_data['gender']}
**Chief Complaint:** {patient_data['chief_complaint']}

**Clinical Presentation:** {patient_data['presentation']}

How can I assist with this case?
"""
    
    history = [[f"Load patient case: {patient_data['name']}", context_message]]
    return convert_messages_for_gradio([]), history

# --- Main UI Builder ---

def build_ui():
    """Build the main Gradio interface"""
    
    # Initialize components on startup
    initialize_components()
    
    # Custom CSS
    css = """
    :root {
      --id-primary: #1e40af;
      --id-secondary: #3b82f6;
      --id-accent: #06d6a0;
      --id-success: #059669;
      --id-warning: #d97706;
      --id-error: #dc2626;
      --id-bg: #f8fafc;
      --id-surface: #ffffff;
      --id-text: #1e293b;
      --id-text-muted: #64748b;
    }

    .gradio-container {
      background: linear-gradient(135deg, #f8fafc 0%, #e2e8f0 100%);
      min-height: 100vh;
    }

    .id-header {
      background: linear-gradient(90deg, var(--id-primary), var(--id-secondary));
      color: white;
      padding: 1.5rem;
      border-radius: 12px;
      margin-bottom: 1.5rem;
      box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
    }

    .id-card {
      background: var(--id-surface);
      border-radius: 12px;
      padding: 1.5rem;
      box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.1);
      border: 1px solid #e2e8f0;
    }

    .patient-card {
      background: linear-gradient(135deg, #eff6ff 0%, #dbeafe 100%);
      border: 2px solid var(--id-secondary);
      border-radius: 8px;
      padding: 1rem;
      margin: 0.5rem 0;
      cursor: pointer;
      transition: all 0.2s ease;
    }

    .patient-card:hover {
      transform: translateY(-2px);
      box-shadow: 0 4px 12px rgba(59, 130, 246, 0.15);
    }

    .chat-container {
      background: var(--id-surface);
      border-radius: 12px;
      border: 1px solid #e2e8f0;
    }
    """
    
    with gr.Blocks(title="🦠 ID Agents - Infectious Disease AI", css=css, theme=gr.themes.Soft()) as app:
        
        # Header
        gr.HTML("""
        <div class="id-header">
            <h1 style="margin: 0; display: flex; align-items: center; gap: 12px;">
                🦠 <span>ID Agents</span>
                <span style="font-size: 0.7em; background: rgba(255,255,255,0.2); padding: 4px 8px; border-radius: 6px;">BETA</span>
            </h1>
            <p style="margin: 8px 0 0 0; opacity: 0.9;">AI-Powered Infectious Disease Clinical Decision Support</p>
        </div>
        """)
        
        with gr.Row():
            # Left column - Chat interface
            with gr.Column(scale=7):
                with gr.Group(elem_classes="id-card"):
                    gr.Markdown("### 💬 Clinical Consultation Chat")
                    
                    chatbot = gr.Chatbot(
                        value=[],
                        height=400,
                        elem_classes="chat-container",
                        show_label=False,
                        avatar_images=(None, "🦠")
                    )
                    
                    with gr.Row():
                        msg_input = gr.Textbox(
                            placeholder="Ask about infectious diseases, patient cases, or clinical guidelines...",
                            show_label=False,
                            scale=4
                        )
                        submit_btn = gr.Button("Send", variant="primary", scale=1)
                    
                    gr.Examples(
                        examples=[
                            "What are the current guidelines for treating MRSA pneumonia?",
                            "Help me interpret these blood culture results",
                            "What empirical antibiotics should I start for severe sepsis?",
                            "Explain the mechanism of carbapenem resistance",
                            "What's the latest on COVID-19 treatment protocols?"
                        ],
                        inputs=msg_input,
                        label="💡 Example Queries"
                    )
            
            # Right column - Tools and patient scenarios
            with gr.Column(scale=3):
                with gr.Group(elem_classes="id-card"):
                    gr.Markdown("### 🏥 Patient Scenarios")
                    gr.Markdown("*Click to load a clinical case*")
                    
                    # Patient scenario buttons
                    with gr.Column():
                        patient1_btn = gr.Button("👵 Maria S. - Pneumonia (67F)", elem_classes="patient-card")
                        patient2_btn = gr.Button("👨 James W. - UTI/Sepsis (34M)", elem_classes="patient-card") 
                        patient3_btn = gr.Button("👩 Sarah C. - Rash/Arthritis (28F)", elem_classes="patient-card")
                        patient4_btn = gr.Button("👨 Robert M. - Chronic Cough (45M)", elem_classes="patient-card")
                        patient5_btn = gr.Button("👩 Emma T. - Meningitis (19F)", elem_classes="patient-card")
                        patient6_btn = gr.Button("👴 Michael B. - GI Infection (72M)", elem_classes="patient-card")
                
                with gr.Group(elem_classes="id-card"):
                    gr.Markdown("### 🔧 System Status")
                    
                    if CORE_MODULES_AVAILABLE:
                        status_msg = "✅ **Core modules loaded**\n✅ **AI orchestrator active**\n✅ **Knowledge base ready**"
                    else:
                        status_msg = "⚠️ **Limited functionality**\n🔧 **Core modules unavailable**\n💡 **Demo mode active**"
                    
                    gr.Markdown(status_msg)
        
        # Event handlers
        def handle_submit(message, history):
            if not message.strip():
                return "", history
            
            response, updated_history = safe_chat_response(message, history)
            return "", updated_history
        
        # Wire up the chat interface
        submit_btn.click(
            fn=handle_submit,
            inputs=[msg_input, chatbot],
            outputs=[msg_input, chatbot]
        )
        
        msg_input.submit(
            fn=handle_submit,
            inputs=[msg_input, chatbot],
            outputs=[msg_input, chatbot]
        )
        
        # Wire up patient loading buttons
        patient1_btn.click(fn=load_patient_1, outputs=[msg_input, chatbot])
        patient2_btn.click(fn=load_patient_2, outputs=[msg_input, chatbot])
        patient3_btn.click(fn=load_patient_3, outputs=[msg_input, chatbot])
        patient4_btn.click(fn=load_patient_4, outputs=[msg_input, chatbot])
        patient5_btn.click(fn=load_patient_5, outputs=[msg_input, chatbot])
        patient6_btn.click(fn=load_patient_6, outputs=[msg_input, chatbot])
    
    return app

# --- Main Application Entry Point ---

if __name__ == "__main__":
    try:
        print("🚀 Launching ID Agents with Beta Authentication...")
        
        # Create main app
        app = build_ui()
        
        # Simple Gradio authentication credentials
        auth_credentials = [
            ("dr_smith", "idweek2025"),
            ("id_fellow", "hello"), 
            ("pharmacist", "stewardship"),
            ("ipc_nurse", "infection"),
            ("researcher", "research"),
            ("educator", "education"),
            ("student", "learning"),
            ("admin", "idagents2025"),
            ("guest1", "guest123"),
            ("guest2", "guest456")
        ]
        
        auth_message = """
🦠 **ID Agents Beta Testing Access**

Welcome to the ID Agents beta testing environment!

**Available Test Accounts:**
• **dr_smith** / idweek2025 (ID Physician)
• **id_fellow** / hello (ID Fellow) 
• **pharmacist** / stewardship (Clinical Pharmacist)
• **ipc_nurse** / infection (IPC Coordinator)
• **researcher** / research (Clinical Researcher)
• **educator** / education (Medical Educator)
• **student** / learning (Medical Student - Limited Access)
• **admin** / idagents2025 (Administrator)
• **guest1** / guest123 (Guest Access - Limited)
• **guest2** / guest456 (Guest Access - Limited)

Please use your assigned credentials to access the application.
Built with ❤️ for ID Week 2025 — Empowering Infectious Diseases Innovation
        """
        
        # Check if running on Hugging Face Spaces
        try:
            import os
            if os.getenv("SPACE_ID"):
                # Running on HF Spaces
                launch_config = {
                    "auth": auth_credentials,
                    "auth_message": auth_message,
                    "show_error": True
                }
                print("🔐 Authentication enabled for HF Spaces deployment")
            else:
                # Local development
                launch_config = {
                    "auth": auth_credentials,
                    "auth_message": auth_message,
                    "share": False,
                    "server_name": "127.0.0.1",
                    "server_port": 7860,
                    "show_error": True
                }
                print("🔐 Authentication enabled for local testing")
        except Exception:
            # Fallback configuration with authentication
            launch_config = {
                "auth": auth_credentials,
                "auth_message": auth_message,
                "share": False,
                "server_name": "127.0.0.1", 
                "server_port": 7860,
                "show_error": True
            }
            print("🔐 Authentication enabled with fallback configuration")
        
        print("📋 Available beta test accounts:")
        for username, password in auth_credentials:
            print(f"   • {username} / {password}")
        
        app.launch(**launch_config)
        
    except Exception as e:
        print(f"❌ Failed to launch ID Agents: {e}")
        print("💡 Check your API keys and environment configuration")
        import traceback
        traceback.print_exc()