import streamlit as st
from groq import Client

# API Key Configuration (Use environment variables in production)
GROQ_API_KEY = "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"

def main():
    st.set_page_config(page_title="🧠 AI ассистент врача-невролога", layout="wide")
    st.title("🧠 AI ассистент врача-невролога")

    # Load chatbot model
    @st.cache_resource
    def load_model():
        return Client(api_key=GROQ_API_KEY)

    chatbot = load_model()

    # Strict system prompt with appointment booking feature
    system_prompt = {
        "role": "system",
        "content": ("You are a professional neurology assistant. Your role is to provide accurate and up-to-date medical insights related to neurological disorders, brain health, symptoms, treatments, and neuroscience research."
                    " Always ensure that responses are factual, empathetic, and professional. If a query is **unrelated to neurology** or medical concerns, politely redirect the user by saying: 'I specialize in neurology-related assistance. Let me know how I can help with your neurological health concerns.'"
                    " If a user wants to book an appointment with a neurologist, ask for their preferred date, time, and location. Provide a confirmation message once details are received. Ответы прошу давать на русском языке")
    }

    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state["messages"] = [system_prompt]

    # Display chat history (excluding system message)
    for message in st.session_state["messages"]:
        if message["role"] != "system":
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

    # User input
    user_input = st.chat_input("Прошу здесь задавать вопросы по неврологии..")

    if user_input:
        # Add user message to session state
        st.session_state["messages"].append({"role": "user", "content": user_input})
        with st.chat_message("user"):
            st.markdown(user_input)

        # Generate response with strict system prompt
        response = chatbot.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[system_prompt] + st.session_state["messages"]  # Always include system prompt
        )

        bot_response = response.choices[0].message.content

        # Add assistant's response to session state
        st.session_state["messages"].append({"role": "assistant", "content": bot_response})
        with st.chat_message("assistant"):
            st.markdown(bot_response)

if __name__ == "__main__":
    main()