import streamlit as st
from openai import OpenAI
import time
import os
import logging
from groq import Groq
import os
import streamlit as st
import logging
from groq import Groq

# Logging setup
logging.basicConfig(level=logging.INFO)

# Streamlit page configuration
st.set_page_config(
    page_title="Gazal.ai-o1-preview",
    page_icon="🦌",
    layout="centered"
)

# Improved CSS for better contrast and design
st.markdown("""
<style>
    /* Background and main content styling */
    .stApp {
        background-color: #121212;
        color: #F5F5F5;
    }
    h1, h2, h3, h4, h5, h6 {
        color: #F5F5F5;
    }

    /* Chat bubbles styling */
    .user-message {
        text-align: right;
        background-color: #007AFF;
        color: white;
        padding: 0.5rem 1rem;
        border-radius: 15px;
        margin: 0.5rem 0;
        display: inline-block;
        max-width: 70%;
        box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.2);
    }
    .bot-message {
        text-align: left;
        background-color: #333333;
        color: #F5F5F5;
        padding: 0.5rem 1rem;
        border-radius: 15px;
        margin: 0.5rem 0;
        display: inline-block;
        max-width: 70%;
        box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.2);
    }

    /* Input field styling */
    input {
        background-color: #1E1E1E;
        color: #F5F5F5;
        border: 1px solid #333333;
        padding: 0.5rem;
        border-radius: 5px;
    }

    /* Button styling */
    button {
        background-color: #007AFF;
        color: white;
        border: none;
        border-radius: 5px;
        padding: 0.5rem 1rem;
        font-size: 1rem;
        cursor: pointer;
        box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.2);
    }
    button:hover {
        background-color: #005BB5;
    }

    /* Footer styling */
    .footer {
        text-align: center;
        margin-top: 2rem;
        color: #888888;
    }

    /* Disclaimer styling */
    .disclaimer {
        background-color: #333333;
        color: #F5F5F5;
        padding: 1rem;
        border-radius: 10px;
        margin-bottom: 20px;
        box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.2);
        font-size: 0.9rem;
    }
    .disclaimer h4 {
        color: #FF5733;
        margin: 0;
    }
</style>
<div class="disclaimer">
    <h4>Disclaimer</h4>
    <p>
        This app is for demonstration purposes only.
        The purpose of this demo is to showcase the power of reasoning large language models (LLMs) in guiding clinical decision support systems.
        It is not intended for clinical use. Please consult medical professionals for accurate medical advice.
    </p>
</div>
""", unsafe_allow_html=True)

# Groq API client initialization
@st.cache_resource
def init_groq_client():
    return Groq(api_key=os.getenv("GROQ_API_KEY"))

# Chat with Groq model
def chat_with_groq(client, message, history):
    try:
        # Build the conversation context
        messages = [
            {"role": "system", "content": "You are a helpful medical and clinical decision support system. Think step by step before answering."},
            *[{"role": "user" if i % 2 == 0 else "assistant", "content": m} for h in history for i, m in enumerate(h)],
            {"role": "user", "content": message}
        ]

        # Call the Groq model
        completion = client.chat.completions.create(
            model="deepseek-r1-distill-llama-70b",
            messages=messages,
            temperature=0.6,
            max_tokens=4000,
            top_p=0.95,
            stream=True,  # Stream the response
        )

        # Stream the response chunk by chunk
        response = ""
        for chunk in completion:
            content = chunk.choices[0].delta.content or ""
            response += content
            yield response

    except Exception as e:
        logging.error(f"Error during Groq inference: {str(e)}")
        yield f"An error occurred: {str(e)}. Please check your API key and network connection."

# Initialize app state
if "history" not in st.session_state:
    st.session_state["history"] = []  # [(user_message, bot_response), ...]

# Display the app title and description
st.title("Gazal.ai-o1-preview 🦌")
st.write("Ask gazal.ai any healthcare question and it will provide step-by-step reasoning.")

# Input form
with st.form("chat_form", clear_on_submit=True):
    user_message = st.text_input("Your Message:", key="user_input")
    submitted = st.form_submit_button("Send")

# Process user input and display chat
if submitted and user_message:
    # Add user message to history
    st.session_state["history"].append((user_message, None))

    # Display chat history
    for user_text, bot_text in st.session_state["history"]:
        st.markdown(f'<div class="user-message">{user_text}</div>', unsafe_allow_html=True)
        if bot_text:
            st.markdown(f'<div class="bot-message">{bot_text}</div>', unsafe_allow_html=True)

    # Initialize Groq client
    groq_client = init_groq_client()

    # Generate bot response
    response_placeholder = st.empty()  # Placeholder for streaming updates
    bot_response = ""
    for partial_response in chat_with_groq(groq_client, user_message, st.session_state["history"][:-1]):
        bot_response = partial_response  # Update bot response incrementally
        response_placeholder.markdown(f'<div class="bot-message">{bot_response}</div>', unsafe_allow_html=True)

    # Update history with full bot response
    st.session_state["history"][-1] = (user_message, bot_response)

# Clear chat history button
# Clear chat history button
if st.button("Clear Chat"):
    st.session_state["history"] = []
    st.rerun()

# Footer
st.markdown('<div class="footer">Gazal.ai © 2025 Made with ❤️ by TachyHealth</div>', unsafe_allow_html=True)