#!/usr/bin/env python3
"""
Maternal Health RAG Chatbot - Gradio Interface
Complete chatbot interface for Sri Lankan maternal health guidelines
"""

import gradio as gr
import json
import time
from typing import List, Tuple, Dict, Any
from datetime import datetime
from pathlib import Path

from maternal_health_rag import MaternalHealthRAG, QueryResponse

class MaternalHealthChatbot:
    """Maternal Health Chatbot with Gradio interface"""
    
    def __init__(self):
        self.rag_system = None
        self.chat_history = []
        self.session_stats = {
            'queries_processed': 0,
            'total_response_time': 0.0,
            'session_start': datetime.now()
        }
        
        # Initialize RAG system
        self.initialize_chatbot()
    
    def initialize_chatbot(self):
        """Initialize the RAG system for the chatbot"""
        try:
            print("🚀 Initializing Maternal Health RAG Chatbot...")
            self.rag_system = MaternalHealthRAG(use_mock_llm=True)
            print("✅ Chatbot initialized successfully!")
        except Exception as e:
            print(f"❌ Failed to initialize chatbot: {e}")
            raise
    
    def process_query(self, message: str, history: List[List[str]]) -> Tuple[str, List[List[str]]]:
        """Process user query and return response with updated history"""
        
        if not message.strip():
            return "", history
        
        try:
            # Process query through RAG system
            response = self.rag_system.query(message)
            
            # Update session statistics
            self.session_stats['queries_processed'] += 1
            self.session_stats['total_response_time'] += response.response_time
            
            # Format response with metadata
            formatted_response = self.format_response(response)
            
            # Update chat history
            history.append([message, formatted_response])
            
            return "", history
            
        except Exception as e:
            error_response = f"I apologize, but I encountered an error: {str(e)}. Please try rephrasing your question."
            history.append([message, error_response])
            return "", history
    
    def format_response(self, response: QueryResponse) -> str:
        """Format the RAG response for display"""
        
        # Main answer
        formatted_answer = f"**🏥 Clinical Response:**\n{response.answer}\n\n"
        
        # Confidence and metadata
        confidence_emoji = "🟢" if response.confidence >= 0.7 else "🟡" if response.confidence >= 0.4 else "🔴"
        formatted_answer += f"**📊 Response Metadata:**\n"
        formatted_answer += f"{confidence_emoji} Confidence: {response.confidence:.1%}\n"
        formatted_answer += f"⏱️ Response Time: {response.response_time:.2f}s\n"
        formatted_answer += f"📚 Sources: {response.metadata['num_sources']} guidelines\n"
        
        if response.metadata['content_types']:
            content_types = ", ".join(response.metadata['content_types'])
            formatted_answer += f"📋 Content Types: {content_types}\n"
        
        # Source details (for high-confidence responses)
        if response.confidence >= 0.6 and response.sources:
            formatted_answer += f"\n**📖 Key Sources:**\n"
            for i, source in enumerate(response.sources[:3], 1):  # Show top 3 sources
                source_preview = source.content[:150] + "..." if len(source.content) > 150 else source.content
                formatted_answer += f"{i}. **{source.chunk_type.title()}** (Score: {source.score:.2f})\n"
                formatted_answer += f"   {source_preview}\n\n"
        
        # Safety disclaimer
        formatted_answer += "\n---\n"
        formatted_answer += "⚠️ **Medical Disclaimer:** This information is based on Sri Lankan maternal health guidelines and is for educational purposes only. Always consult with qualified healthcare professionals for medical decisions."
        
        return formatted_answer
    
    def get_example_queries(self) -> List[str]:
        """Get example queries for the interface"""
        return [
            "What is the recommended dosage of magnesium sulfate for preeclampsia?",
            "How should postpartum hemorrhage be managed in emergency situations?",
            "What are the signs and symptoms of puerperal sepsis?",
            "What is the normal fetal heart rate range during labor?",
            "When is cesarean section indicated during delivery?",
            "How to manage gestational diabetes during pregnancy?",
            "What are the contraindications for vaginal delivery?",
            "How to recognize and manage eclampsia?",
            "What is the proper management of prolonged labor?",
            "How to handle breech presentation during delivery?"
        ]
    
    def clear_chat(self) -> List[List[str]]:
        """Clear chat history"""
        self.chat_history = []
        return []
    
    def get_system_info(self) -> str:
        """Get system information and statistics"""
        if not self.rag_system:
            return "❌ RAG system not initialized"
        
        stats = self.rag_system.get_system_stats()
        session_time = (datetime.now() - self.session_stats['session_start']).total_seconds()
        
        avg_response_time = (
            self.session_stats['total_response_time'] / self.session_stats['queries_processed']
            if self.session_stats['queries_processed'] > 0 else 0
        )
        
        info = f"""
## 🏥 Maternal Health RAG Assistant - System Information

### 📊 Knowledge Base Statistics
- **Total Medical Chunks:** {stats['vector_store']['total_chunks']:,}
- **Embedding Model:** {stats['vector_store']['embedding_model']}
- **Vector Store Size:** {stats['vector_store']['vector_store_size_mb']:.1f} MB
- **Clinical Content Types:** {len(stats['vector_store']['chunk_type_distribution'])}

### 🧠 RAG Configuration
- **Default Results:** {stats['rag_config']['default_k']} sources per query
- **Context Length:** {stats['rag_config']['max_context_length']:,} characters max
- **LLM Type:** {stats['rag_config']['llm_type'].title()}

### 📈 Session Statistics
- **Queries Processed:** {self.session_stats['queries_processed']}
- **Average Response Time:** {avg_response_time:.2f}s
- **Session Duration:** {session_time:.0f}s
- **System Status:** {stats['status'].title()}

### 📚 Document Coverage
This assistant covers **15 Sri Lankan maternal health guidelines** including:
- National Guidelines for Maternal Care
- SLJOG Clinical Guidelines
- Emergency Management Protocols
- Dosage and Treatment Guidelines
- Postnatal Care Guidelines
"""
        return info

def create_chatbot_interface():
    """Create the Gradio chatbot interface"""
    
    # Initialize chatbot
    chatbot = MaternalHealthChatbot()
    
    # Create Gradio interface
    with gr.Blocks(
        title="Maternal Health Assistant", 
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            font-family: 'Arial', sans-serif;
        }
        .chat-message {
            font-size: 16px;
            line-height: 1.5;
        }
        """
    ) as demo:
        
        # Header
        gr.Markdown("""
        # 🏥 Sri Lankan Maternal Health RAG Assistant
        
        **Your AI assistant for Sri Lankan maternal health guidelines**
        
        Ask questions about:
        - 💊 Medication dosages and protocols
        - 🚨 Emergency management procedures  
        - 🤱 Maternal and fetal care guidelines
        - 📋 Clinical decision-making support
        - 🔬 Diagnostic criteria and procedures
        
        *Based on official Sri Lankan maternal health guidelines and SLJOG recommendations*
        """)
        
        with gr.Tab("💬 Chat Assistant"):
            # Chat interface
            chatbot_interface = gr.Chatbot(
                label="Maternal Health Assistant",
                height=500,
                elem_classes=["chat-message"]
            )
            
            msg = gr.Textbox(
                label="Your Question",
                placeholder="Ask me about maternal health guidelines, emergency protocols, dosages, or clinical procedures...",
                lines=2
            )
            
            with gr.Row():
                submit_btn = gr.Button("🔍 Ask Question", variant="primary")
                clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
            
            # Example queries
            gr.Markdown("### 💡 Example Questions:")
            with gr.Row():
                examples = chatbot.get_example_queries()
                for i in range(0, len(examples), 2):
                    with gr.Column():
                        if i < len(examples):
                            gr.Button(
                                examples[i], 
                                variant="outline",
                                size="sm"
                            ).click(
                                lambda x=examples[i]: x,
                                outputs=msg
                            )
                        if i+1 < len(examples):
                            gr.Button(
                                examples[i+1], 
                                variant="outline",
                                size="sm"
                            ).click(
                                lambda x=examples[i+1]: x,
                                outputs=msg
                            )
        
        with gr.Tab("📊 System Information"):
            system_info = gr.Markdown(
                chatbot.get_system_info(),
                label="System Information"
            )
            
            refresh_btn = gr.Button("🔄 Refresh Stats", variant="secondary")
            refresh_btn.click(
                chatbot.get_system_info,
                outputs=system_info
            )
        
        with gr.Tab("ℹ️ About"):
            gr.Markdown("""
            ## About This Assistant
            
            This **Maternal Health RAG Assistant** provides information based on official Sri Lankan maternal health guidelines using Retrieval-Augmented Generation (RAG) technology.
            
            ### 🔧 Technical Features
            - **Vector-based search** through 542 medical content chunks
            - **Semantic similarity** using all-MiniLM-L6-v2 embeddings
            - **Clinical importance scoring** for prioritizing critical information
            - **Medical context filtering** by content type (dosage, emergency, procedure, etc.)
            - **Sub-second response times** with confidence scoring
            
            ### 📚 Knowledge Base
            - **15 comprehensive documents** covering maternal health
            - **479 pages** of clinical guidelines processed
            - **48 clinical tables** with dosage and protocol information
            - **107,010 words** of medical content indexed
            
            ### ⚠️ Important Disclaimers
            1. **For Educational Use Only:** This tool provides information based on guidelines but should not replace professional medical judgment
            2. **Always Consult Healthcare Professionals:** Medical decisions should always involve qualified healthcare providers
            3. **Regular Updates:** Guidelines may change - always verify with the latest official sources
            4. **Emergency Situations:** In medical emergencies, contact emergency services immediately
            
            ### 🏗️ Built With
            - **LangChain** for RAG pipeline orchestration
            - **FAISS** for efficient vector similarity search
            - **Sentence Transformers** for medical text embeddings
            - **Gradio** for the user interface
            - **pdfplumber** for medical document processing
            
            ---
            *Developed for educational and clinical reference purposes*
            """)
        
        # Event handlers
        submit_btn.click(
            chatbot.process_query,
            inputs=[msg, chatbot_interface],
            outputs=[msg, chatbot_interface]
        )
        
        msg.submit(
            chatbot.process_query,
            inputs=[msg, chatbot_interface],
            outputs=[msg, chatbot_interface]
        )
        
        clear_btn.click(
            chatbot.clear_chat,
            outputs=chatbot_interface
        )
    
    return demo

def main():
    """Main function to launch the chatbot"""
    print("🚀 Launching Maternal Health RAG Chatbot...")
    
    try:
        # Create and launch interface
        demo = create_chatbot_interface()
        
        print("✅ Chatbot interface created successfully!")
        print("🌐 Launching on http://localhost:7860")
        print("📱 Access from other devices using the public link")
        
        # Launch with public sharing for easier access
        demo.launch(
            server_name="0.0.0.0",  # Allow external access
            server_port=7860,
            share=True,  # Create public link
            show_error=True,
            quiet=False
        )
        
    except Exception as e:
        print(f"❌ Failed to launch chatbot: {e}")
        raise

if __name__ == "__main__":
    main()