import os
import streamlit as st
from dotenv import load_dotenv
import google.generativeai as gen_ai

# Load environment variables
load_dotenv()

# Configure Streamlit page settings
st.set_page_config(
    page_title="Health Assistant Chatbot",
    page_icon="🩺",
    layout="wide",
)

# Custom CSS for styling
st.markdown(
    """
    <style>
        body, .stApp {
            background-color: #E3F2FD;
        }
        body {
            background-color: #ADD8E6;
        }
        .stTabs [data-baseweb="tab-list"] {
            justify-content: center;
        }
        .stTabs [data-baseweb="tab"] {
            background-color: #E3F2FD;
            border-radius: 20px;
            font-weight: bold;
            color: black;
            padding: 10px;
        }
        .stTabs [data-baseweb="tab"][aria-selected="true"] {
            background-color: #1E90FF;
            color: white;
            border-radius: 20px;
        }
        h2 {
            color: #003366;
        }
    </style>
    """,
    unsafe_allow_html=True,
)

# Retrieve the Google API key from the environment
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

if not GOOGLE_API_KEY:
    st.error("🚨 API key not found! Please set the GOOGLE_API_KEY in your .env file.")
    st.stop()

# Configure the Generative AI model
try:
    gen_ai.configure(api_key=GOOGLE_API_KEY)
    model = gen_ai.GenerativeModel("gemini-1.5-pro")  # Using a model for health-related queries
except Exception as e:
    st.error(f"❌ Error initializing the Gemini-Pro model: {e}")
    st.stop()

# Initialize the chat session
if "chat_session" not in st.session_state:
    try:
        st.session_state.chat_session = model.start_chat(history=[])
    except Exception as e:
        st.error(f"❌ Error initializing chat session: {e}")
        st.stop()

# Display chatbot title
st.title("🩺 Health Assistant Chatbot")

# Symptom categories for the user to select
symptom_tabs = ["🤧 Sneezing", "😷 Headache", "🤢 Stomach Ache", "🦠 Fever", "💪 Fatigue"]
selected_symptom_tab = st.tabs(symptom_tabs)

# Sample symptoms and associated treatment information
symptom_treatment = {
    "🤧 Sneezing": [
        "Possible cause: Allergies or common cold.",
        "Treatment recommendation: Try antihistamines for allergies or rest and fluids for a cold."
    ],
    "😷 Headache": [
        "Possible cause: Tension, dehydration, or migraines.",
        "Treatment recommendation: Drink water, take over-the-counter pain relievers, or rest in a dark room."
    ],
    "🤢 Stomach Ache": [
        "Possible cause: Indigestion, food poisoning, or gas.",
        "Treatment recommendation: Drink ginger tea, rest, and avoid heavy meals."
    ],
    "🦠 Fever": [
        "Possible cause: Infection or viral illness.",
        "Treatment recommendation: Take fever reducers like acetaminophen, stay hydrated, and rest."
    ],
    "💪 Fatigue": [
        "Possible cause: Stress, sleep deprivation, or illness.",
        "Treatment recommendation: Ensure proper sleep, hydrate, and reduce stress."
    ]
}

# Display symptom recommendations under the selected tab
for i, tab in enumerate(selected_symptom_tab):
    with tab:
        st.subheader(f"📌 {symptom_tabs[i]} Symptoms")
        for recommendation in symptom_treatment[symptom_tabs[i]]:
            st.write(f"- {recommendation}")

# Input field for user's symptoms
user_symptom = st.chat_input("💬 Enter your symptom...")

if user_symptom:
    # Display the user's symptom
    st.chat_message("user").markdown(user_symptom)

    # Send symptom message to Gemini-Pro for response
    try:
        gemini_response = st.session_state.chat_session.send_message(
            f"Provide possible causes and treatments for the following symptom: {user_symptom}"
        )
        
        # Display AI response
        with st.chat_message("assistant"):
            st.markdown(gemini_response.text)
    except Exception as e:
        st.error(f"❌ Error processing your message: {e}")