# import os
# import streamlit as st
# from groq import Groq
# import re
# from datetime import datetime

# # Initialize Groq client
# client = Groq(api_key="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")

# # Set page configuration
# st.set_page_config(
#     page_title="Virtual Medical Assistant",
#     page_icon="👨‍⚕️",
#     layout="wide",
#     initial_sidebar_state="expanded",
# )

# # Custom CSS with improved chat styling
# st.markdown(
#     """
#     <style>
#     .main {
#         background-color: #f5f7f9;
#     }
#     .stButton>button {
#         width: 100%;
#         border-radius: 5px;
#         height: 3em;
#         background-color: #0083B8;
#         color: white;
#         border: none;
#     }
#     .stButton>button:hover {
#         background-color: #00669e;
#     }
#     .chat-container {
#         background-color: #f0f2f6;
#         padding: 20px;
#         border-radius: 10px;
#         box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
#         margin-bottom: 20px;
#     }
#     .user-message {
#         background-color: #0083B8;
#         padding: 10px 15px;
#         border-radius: 15px 15px 0 15px;
#         margin: 10px 0;
#         margin-left: 20%;
#         color: white;
#         text-align: right;
#         box-shadow: 0 1px 2px rgba(0, 0, 0, 0.1);
#     }
#     .assistant-message {
#         background-color: #e9ecef;
#         padding: 10px 15px;
#         border-radius: 15px 15px 15px 0;
#         color: #333;
#         margin: 10px 0;
#         margin-right: 20%;
#         text-align: left;
#         box-shadow: 0 1px 2px rgba(0, 0, 0, 0.1);
#     }
#     .sidebar-content {
#         padding: 20px;
#     }
#     .main-title {
#         text-align: center;
#         color: #0083B8;
#         padding: 20px 0;
#     }
#     .delete-button {
#         background-color: #dc3545 !important;
#     }
#     .delete-button:hover {
#         background-color: #c82333 !important;
#     }
#     </style>
#     """,
#     unsafe_allow_html=True,
# )

# # System prompt remains the same
# SYSTEM_PROMPT = """You are a virtual medical assistant with great expertise and empathy. Your role is:

#     1. Scope of Assistance:
#     - Diagnose common symptoms
#     - Suggest possible conditions
#     - Provide recommendations for common over-the-counter medications
#     - Offer advice on managing health issues
#     - Answer medical questions only
#     - Inform the patient about symptoms of illnesses

#     2. Important Limitations:
#     - Do not answer non-medical questions
#     - When asked a non-medical question, simply say, "Sorry, this question is outside my medical expertise."
#     - Provide definitive diagnoses only for very common conditions
#     - Do not prescribe prescription medications
#     - Do not provide advice on serious medical conditions
#     - Do not greet the user in every response. Do so only if the user starts with a greeting.

#     3. When dealing with symptoms:

#     A. Gather information - Ask specific questions about:
#     - Duration of symptoms
#     - Severity of symptoms
#     - Age
#     - Medical history
#     - Current medications
#     - Associated symptoms

#     B. Assessment:
#     - Analyze the reported symptoms
#     - Link symptoms to common possible conditions
#     - Determine the severity of the condition

#     C. Recommendations:
#     - Suggest safe home remedies
#     - Recommend over-the-counter medications
#     - Provide prevention and self-care tips

#     4. When to refer the patient to a doctor:
#     - In case of severe symptoms
#     - If symptoms persist for a long time
#     - If the condition worsens
#     - When specific medical tests are needed
#     - In emergencies

#     5. Mandatory Reminder:
#     - Always emphasize that the advice provided is not a substitute for consulting a doctor
#     - Encourage visiting a healthcare provider in serious cases
#     - Explain that this advice is general and not an official medical diagnosis

#     6. Language and Communication:
#     - Use simple and clear language
#     - Maintain a professional and empathetic tone
#     - Avoid complex medical terminology

#     7. In emergencies:
#     - Direct the patient immediately to the nearest emergency department
#     - Provide simple first aid instructions if necessary
#     - Emphasize the importance of seeking immediate medical help

#     Always remember: Patient safety is the top priority, and in case of doubt, always recommend visiting a doctor."""

# # Initialize session state
# if "chat_history" not in st.session_state:
#     st.session_state.chat_history = {}
# if "current_chat_id" not in st.session_state:
#     st.session_state.current_chat_id = None
# if "temp_chat" not in st.session_state:
#     st.session_state.temp_chat = None


# def get_groq_response(messages):
#     try:
#         chat_completion = client.chat.completions.create(
#             messages=messages,
#             model="llama3-8b-8192",
#             temperature=0.7,
#             max_tokens=1000,
#         )
#         return chat_completion.choices[0].message.content
#     except Exception as e:
#         return f"Sorry, there was an error in the connection: {str(e)}"


# def create_new_chat():
#     chat_id = datetime.now().strftime("%Y%m%d_%H%M%S")
#     st.session_state.temp_chat = [{"role": "system", "content": SYSTEM_PROMPT}]
#     st.session_state.current_chat_id = chat_id
#     return chat_id


# def save_chat():
#     if st.session_state.temp_chat and len(st.session_state.temp_chat) > 1:
#         st.session_state.chat_history[st.session_state.current_chat_id] = (
#             st.session_state.temp_chat
#         )
#     st.session_state.temp_chat = None


# def delete_chat(chat_id):
#     if chat_id in st.session_state.chat_history:
#         del st.session_state.chat_history[chat_id]
#     if st.session_state.current_chat_id == chat_id:
#         st.session_state.current_chat_id = None


# def get_chat_preview(chat):
#     for message in chat:
#         if message["role"] == "user":
#             preview = message["content"][:30]
#             return f"{preview}..." if len(message["content"]) > 30 else preview
#     return "New Chat"


# # Main layout
# st.markdown(
#     '<h1 class="main-title">👨‍⚕️ Virtual Medical Assistant</h1>', unsafe_allow_html=True
# )

# # Sidebar
# with st.sidebar:
#     st.markdown('<div class="sidebar-content">', unsafe_allow_html=True)

#     if st.button("New Chat ➕", key="new_chat"):
#         create_new_chat()

#     st.markdown("### Previous Chats")

#     for chat_id in st.session_state.chat_history:
#         col1, col2 = st.columns([4, 1])
#         with col1:
#             if st.button(
#                 get_chat_preview(st.session_state.chat_history[chat_id]),
#                 key=f"select_{chat_id}",
#             ):
#                 st.session_state.current_chat_id = chat_id
#                 st.session_state.temp_chat = None
#         with col2:
#             if st.button("🗑️", key=f"delete_{chat_id}", help="Delete Chat"):
#                 delete_chat(chat_id)
#                 st.rerun()

#     st.markdown("</div>", unsafe_allow_html=True)

# # Main chat interface
# current_chat = st.session_state.temp_chat or (
#     st.session_state.chat_history.get(st.session_state.current_chat_id, None)
# )

# if current_chat:
#     st.markdown('<div class="chat-container">', unsafe_allow_html=True)

#     # Display chat history
#     for message in current_chat[1:]:  # Skip system prompt
#         if message["role"] == "user":
#             st.markdown(
#                 f'<div class="user-message">{message["content"]}</div>',
#                 unsafe_allow_html=True,
#             )
#         else:
#             st.markdown(
#                 f'<div class="assistant-message">{message["content"]}</div>',
#                 unsafe_allow_html=True,
#             )

#     st.markdown("</div>", unsafe_allow_html=True)
    
#     # User input
#     user_input = st.chat_input("Type your question here...")
#     if user_input:
#         # Check if current_chat_id exists, if not create a new one
#         if not st.session_state.current_chat_id:
#             create_new_chat()
#             current_chat = st.session_state.temp_chat
            
#         # Add user message
#         current_chat.append({"role": "user", "content": user_input})
        
#         # Get assistant response
#         with st.spinner("Thinking..."):
#             assistant_response = get_groq_response(current_chat)
        
#         # Add assistant response to chat
#         current_chat.append({"role": "assistant", "content": assistant_response})

#         # Save chat if temporary
#         if st.session_state.temp_chat:
#             save_chat()
#         st.rerun()
# else:
#     st.markdown(
#         """
#         <div style='text-align: center; padding: 20px;'>
#             <p style='color: black; font-size: 16px;'>
#                 Start a new chat by clicking the 'New Chat' ➕ button
#             </p>
#         </div>
#         """,
#         unsafe_allow_html=True,
#     )

# # Footer
# st.markdown(
#     """
#     <div style='text-align: center; color: #666; padding: 20px;'>
#         <p>Reminder: This medical assistant is for general consultations only. Please consult a doctor for serious medical conditions.</p>
#     </div>
#     """,
#     unsafe_allow_html=True,
# )




import os
import streamlit as st
from groq import Groq
import re
from datetime import datetime
from dotenv import load_dotenv

load_dotenv()
# Initialize Groq client with environment variable
Groq_API_KEY = os.getenv("GROQ_API_KEY")
# Initialize Groq client
client = Groq(api_key= Groq_API_KEY)

# Set page configuration
st.set_page_config(
    page_title="Virtual Medical Assistant",
    page_icon="👨‍⚕️",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Custom CSS with improved chat styling
st.markdown(
    """
    <style>
    .main {
        background-color: #f5f7f9;
    }
    .stButton>button {
        width: 100%;
        border-radius: 5px;
        height: 3em;
        background-color: #0083B8;
        color: white;
        border: none;
    }
    .stButton>button:hover {
        background-color: #00669e;
    }
    .chat-container {
        background-color: #f0f2f6;
        padding: 20px;
        border-radius: 10px;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
        margin-bottom: 20px;
    }
    .user-message {
        background-color: #0083B8;
        padding: 10px 15px;
        border-radius: 15px 15px 0 15px;
        margin: 10px 0;
        margin-left: 20%;
        color: white;
        text-align: right;
        box-shadow: 0 1px 2px rgba(0, 0, 0, 0.1);
    }
    .assistant-message {
        background-color: #e9ecef;
        padding: 10px 15px;
        border-radius: 15px 15px 15px 0;
        color: #333;
        margin: 10px 0;
        margin-right: 20%;
        text-align: left;
        box-shadow: 0 1px 2px rgba(0, 0, 0, 0.1);
    }
    .sidebar-content {
        padding: 20px;
    }
    .main-title {
        text-align: center;
        color: #0083B8;
        padding: 20px 0;
    }
    .delete-button {
        background-color: #dc3545 !important;
    }
    .delete-button:hover {
        background-color: #c82333 !important;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# Updated System prompt for multilingual support
SYSTEM_PROMPT = """You are a multilingual virtual medical assistant with great expertise and empathy. Your role is:

    1. Scope of Assistance:
    - Diagnose common symptoms
    - Suggest possible conditions
    - Provide recommendations for common over-the-counter medications
    - Offer advice on managing health issues
    - Answer medical questions only
    - Inform the patient about symptoms of illnesses

    2. Important Limitations:
    - Do not answer non-medical questions
    - When asked a non-medical question, simply say in the user's language: "Sorry, this question is outside my medical expertise."
    - Provide definitive diagnoses only for very common conditions
    - Do not prescribe prescription medications
    - Do not provide advice on serious medical conditions
    - Do not greet the user in every response. Do so only if the user starts with a greeting.

    3. When dealing with symptoms:

    A. Gather information - Ask specific questions about:
    - Duration of symptoms
    - Severity of symptoms
    - Age
    - Medical history
    - Current medications
    - Associated symptoms

    B. Assessment:
    - Analyze the reported symptoms
    - Link symptoms to common possible conditions
    - Determine the severity of the condition

    C. Recommendations:
    - Suggest safe home remedies
    - Recommend over-the-counter medications
    - Provide prevention and self-care tips

    4. When to refer the patient to a doctor:
    - In case of severe symptoms
    - If symptoms persist for a long time
    - If the condition worsens
    - When specific medical tests are needed
    - In emergencies

    5. Mandatory Reminder:
    - Always emphasize that the advice provided is not a substitute for consulting a doctor
    - Encourage visiting a healthcare provider in serious cases
    - Explain that this advice is general and not an official medical diagnosis

    6. Language and Communication:
    - ALWAYS respond in the same language the user is using
    - If the user writes in Spanish, respond in Spanish
    - If the user writes in Arabic, respond in Arabic
    - If the user writes in French, respond in French
    - For any language the user chooses, respond in that same language
    - Use simple and clear language
    - Maintain a professional and empathetic tone
    - Avoid complex medical terminology

    7. In emergencies:
    - Direct the patient immediately to the nearest emergency department
    - Provide simple first aid instructions if necessary
    - Emphasize the importance of seeking immediate medical help

    Always remember: Patient safety is the top priority, and in case of doubt, always recommend visiting a doctor."""

# Initialize session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = {}
if "current_chat_id" not in st.session_state:
    st.session_state.current_chat_id = None
if "temp_chat" not in st.session_state:
    st.session_state.temp_chat = None


def get_groq_response(messages):
    try:
        chat_completion = client.chat.completions.create(
            messages=messages,
            model="llama3-8b-8192",
            temperature=0.7,
            max_tokens=1000,
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        return f"Sorry, there was an error in the connection: {str(e)}"


def create_new_chat():
    chat_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    st.session_state.temp_chat = [{"role": "system", "content": SYSTEM_PROMPT}]
    st.session_state.current_chat_id = chat_id
    return chat_id


def save_chat():
    if st.session_state.temp_chat and len(st.session_state.temp_chat) > 1:
        st.session_state.chat_history[st.session_state.current_chat_id] = (
            st.session_state.temp_chat
        )
    st.session_state.temp_chat = None


def delete_chat(chat_id):
    if chat_id in st.session_state.chat_history:
        del st.session_state.chat_history[chat_id]
    if st.session_state.current_chat_id == chat_id:
        st.session_state.current_chat_id = None


def get_chat_preview(chat):
    for message in chat:
        if message["role"] == "user":
            preview = message["content"][:30]
            return f"{preview}..." if len(message["content"]) > 30 else preview
    return "New Chat"


# Main layout
st.markdown(
    '<h1 class="main-title">👨‍⚕️ Call on Doc Virtual Medical Assistant</h1>', unsafe_allow_html=True
)

# Add language indicator in the sidebar
with st.sidebar:
    st.markdown('<div class="sidebar-content">', unsafe_allow_html=True)
    
    st.markdown("### Multilingual Support")
    st.markdown("This assistant automatically responds in the language you use to ask your question.")
    
    if st.button("New Chat ➕", key="new_chat"):
        create_new_chat()

    st.markdown("### Previous Chats")

    for chat_id in st.session_state.chat_history:
        col1, col2 = st.columns([4, 1])
        with col1:
            if st.button(
                get_chat_preview(st.session_state.chat_history[chat_id]),
                key=f"select_{chat_id}",
            ):
                st.session_state.current_chat_id = chat_id
                st.session_state.temp_chat = None
        with col2:
            if st.button("🗑️", key=f"delete_{chat_id}", help="Delete Chat"):
                delete_chat(chat_id)
                st.rerun()

    st.markdown("</div>", unsafe_allow_html=True)

# Main chat interface
current_chat = st.session_state.temp_chat or (
    st.session_state.chat_history.get(st.session_state.current_chat_id, None)
)

if current_chat:
    st.markdown('<div class="chat-container">', unsafe_allow_html=True)

    # Display chat history
    for message in current_chat[1:]:  # Skip system prompt
        if message["role"] == "user":
            st.markdown(
                f'<div class="user-message">{message["content"]}</div>',
                unsafe_allow_html=True,
            )
        else:
            st.markdown(
                f'<div class="assistant-message">{message["content"]}</div>',
                unsafe_allow_html=True,
            )

    st.markdown("</div>", unsafe_allow_html=True)
    
    # User input
    user_input = st.chat_input("Type your question in any language...")
    if user_input:
        # Check if current_chat_id exists, if not create a new one
        if not st.session_state.current_chat_id:
            create_new_chat()
            current_chat = st.session_state.temp_chat
            
        # Add user message
        current_chat.append({"role": "user", "content": user_input})
        
        # Get assistant response
        with st.spinner("Thinking..."):
            assistant_response = get_groq_response(current_chat)
        
        # Add assistant response to chat
        current_chat.append({"role": "assistant", "content": assistant_response})

        # Save chat if temporary
        if st.session_state.temp_chat:
            save_chat()
        st.rerun()
else:
    st.markdown(
        """
        <div style='text-align: center; padding: 20px;'>
            <p style='color: black; font-size: 16px;'>
                Start a new chat by clicking the 'New Chat ➕' button
            </p>
            <p style='color: white; font-size: 14px;'>
                You can ask questions in any language, and the assistant will respond in the same language.
            </p>
        </div>
        """,
        unsafe_allow_html=True,
    )

# Footer
st.markdown(
    """
    <div style='text-align: center; color: #666; padding: 20px;'>
        <p>Reminder: This medical assistant is for general consultations only. Please consult a doctor for serious medical conditions.</p>
    </div>
    """,
    unsafe_allow_html=True,
)