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: white; 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,
)