"""
Streamlit application for the medical chatbot
"""
import streamlit as st
import time
from embedding_service import EmbeddingService
from medical_chatbot import MedicalChatbot
import os

# Page configuration
st.set_page_config(
    page_title="Medical Chatbot",
    page_icon="🏥",
    layout="wide"
)

# Initialize session state
if "chatbot" not in st.session_state:
    st.session_state.chatbot = None
    st.session_state.embeddings_initialized = False
    st.session_state.messages = []

def initialize_chatbot():
    """Initialize the chatbot and embedding service"""
    try:
        if not st.session_state.embeddings_initialized:
            with st.spinner("Initializing medical chatbot..."):
                embedding_service = EmbeddingService()
                st.session_state.chatbot = MedicalChatbot(embedding_service)
                st.session_state.embeddings_initialized = True
    except Exception as e:
        st.error(f"Error initializing chatbot: {str(e)}")
        st.error("Please check your API keys in the .env file")
        st.stop()

# Initialize chatbot
initialize_chatbot()

# Header
st.title("🏥 Medical Chatbot")
st.markdown("Ask medical questions and get evidence-based answers with citations and confidence scores.")

# Sidebar
with st.sidebar:
    st.header("⚙️ Settings")
    
    st.markdown("### About")
    st.info("""
    This medical chatbot:
    - Uses Sentence Transformers for embeddings
    - Retrieves relevant information from medical databases
    - Uses Gemini 1.5 Flash for generating responses
    - Provides citations and confidence scores
    
    **Important:** This is not a substitute for professional medical advice. Always consult with healthcare professionals for medical decisions.
    """)
    
    st.markdown("### Status")
    if st.session_state.embeddings_initialized:
        st.success("✓ Chatbot Initialized")
    else:
        st.error("✗ Not Initialized")
    
    if st.button("🔄 Reload Database"):
        st.warning("This will reload the medical database. This may take a few minutes.")
        st.session_state.embeddings_initialized = False
        initialize_chatbot()
        st.rerun()

# Display chat history
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        
        # Display confidence and citations if available
        if "confidence" in message and "citations" in message:
            st.markdown(f"**Confidence:** {message['confidence']} ({message['confidence_score']:.2f})")
            
            if message["citations"]:
                with st.expander("📚 View Sources"):
                    for cit_id, cit_info in message["citations"].items():
                        st.markdown(f"""
                        **{cit_id}**
                        - Source: {cit_info['metadata']['source']}
                        - Similarity: {cit_info['similarity_score']}
                        - Text: {cit_info['text']}
                        """)

# Chat input
if prompt := st.chat_input("Ask a medical question..."):
    # Add user message to history
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Display user message
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # Generate response
    if st.session_state.chatbot:
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                response = st.session_state.chatbot.generate_response(prompt)
                
                # Display response
                st.markdown(response['response'])
                
                # Display confidence score with color
                if response['confidence'] == "High":
                    confidence_color = "🟢"
                elif response['confidence'] == "Medium":
                    confidence_color = "🟡"
                else:
                    confidence_color = "🔴"
                
                st.markdown(f"{confidence_color} **Confidence:** {response['confidence']} ({response['confidence_score']:.2%})")
                
                # Display citations
                if response['citations']:
                    with st.expander("📚 Sources & Citations"):
                        for cit_id, cit_info in response['citations'].items():
                            col1, col2 = st.columns([3, 1])
                            with col1:
                                st.markdown(f"""
                                **{cit_id}** - {cit_info['metadata']['source']}
                                - Similarity: {cit_info['similarity_score']}
                                - Preview: {cit_info['text']}
                                """)
                
                # Add disclaimer
                st.warning("⚠️ This is not medical advice. Please consult with healthcare professionals for medical decisions.")
        
        # Add assistant message to history
        st.session_state.messages.append({
            "role": "assistant",
            "content": response['response'],
            "confidence": response['confidence'],
            "confidence_score": response['confidence_score'],
            "citations": response['citations']
        })
    else:
        st.error("Chatbot not initialized. Please check the sidebar for errors.")

# Footer
st.markdown("---")
st.markdown(
    """
    <div style='text-align: center; color: gray;'>
    <p>Powered by Gemini 1.5 Flash | Sentence Transformers | Pinecone DB</p>
    <p>Not a substitute for professional medical advice</p>
    </div>
    """,
    unsafe_allow_html=True
)