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

# Load environment variables
load_dotenv()

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

# Custom CSS for styling
tabs_css = """
    <style>
        body, .stApp {
            background-color: #E3F2FD;
        }
        h2 {
            color: #003366;
        }
        .stTextInput, .stButton > button {
            border-radius: 10px;
            font-size: 16px;
        }
        .response-box {
            background-color: #ffffff;
            padding: 15px;
            border-radius: 10px;
            box-shadow: 2px 2px 10px rgba(0, 0, 0, 0.1);
        }
        .header-image {
            text-align: center;
        }
        .stTabs [data-baseweb="tab-list"] {
            background-color: white;
            border-radius: 20px;
            padding: 5px;
            box-shadow: 2px 2px 10px rgba(0, 0, 0, 0.1);
        }
        .stTabs [data-baseweb="tab"] {
            color: #003366;
            font-weight: bold;
            border-radius: 20px;
            padding: 10px 20px;
        }
        .stTabs [aria-selected="true"] {
            background-color: #E3F2FD !important;
            color: #003366 !important;
            font-weight: bold;
        }
    </style>
"""

st.markdown(tabs_css, 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")
except Exception as e:
    st.error(f"❌ Error initializing the Gemini-Pro model: {e}")
    st.stop()

# Ensure session state for authentication and chat session
if "logged_in" not in st.session_state:
    st.session_state.logged_in = False
if "chat_session" not in st.session_state or st.session_state.chat_session is None:
    try:
        st.session_state.chat_session = model.start_chat(history=[])
    except Exception as e:
        st.error(f"❌ Error initializing chat session: {e}")
        st.stop()

# Store user credentials (For simplicity, stored in session state - No database)
if "users" not in st.session_state:
    st.session_state.users = {}

# User Login/Signup Page
def login_page():
    st.title("🧑‍⚕️🤖 AI Healthcare Assistant Login")
    
    tab1, tab2 = st.tabs(["Login", "Signup"])
    
    with tab1:
        username = st.text_input("👤 Username", key="login_user")
        password = st.text_input("🔑 Password", type="password", key="login_pass")
        
        if st.button("Login"):
            if username in st.session_state.users and st.session_state.users[username] == password:
                st.session_state.logged_in = True
                st.experimental_rerun()
            else:
                st.error("❌ Invalid username or password")
    
    with tab2:
        new_username = st.text_input("👤 Choose a Username", key="signup_user")
        new_password = st.text_input("🔑 Choose a Password", type="password", key="signup_pass")
        
        if st.button("Sign Up"):
            if new_username and new_password:
                if new_username in st.session_state.users:
                    st.error("❌ Username already exists! Choose a different one.")
                else:
                    st.session_state.users[new_username] = new_password
                    st.success("✅ Signup successful! Please login now.")
            else:
                st.error("❌ Please fill in all fields.")

# Display login page if user is not authenticated
if not st.session_state.logged_in:
    login_page()
else:
    # Display Header Image
    st.image("https://cdn.pixabay.com/photo/2017/03/09/12/31/doctor-2133209_1280.jpg", use_column_width=True)
    st.title("🩺 AI Healthcare Assistant")
    
    # Create Tabs
    tab1, tab2 = st.tabs(["Enter Symptoms", "Suggestions"])

    with tab1:
        st.subheader("Enter your symptoms to get advice and treatment guidance")
        user_symptoms = st.text_input("Describe your symptoms (e.g., fever, cough, headache)...")

        if st.button("Get Advice") and user_symptoms:
            st.write(f"**Your Symptoms:** {user_symptoms}")

            # Send message to Gemini-Pro for response
            try:
                response = st.session_state.chat_session.send_message(
                    f"A user is experiencing the following symptoms: {user_symptoms}. Provide possible causes, advice, and treatment guidance in a structured format including causes, recommended actions, and when to see a doctor."
                )

                # Display AI response in a structured format
                st.subheader("🩺 AI Advice & Treatment Guidance:")
                st.markdown(f"<div class='response-box'>{response.text}</div>", unsafe_allow_html=True)
            except Exception as e:
                st.error(f"❌ Error processing your request: {e}")

    with tab2:
        st.subheader("💡 General Health Tips & Suggestions")
        st.markdown(
            """
            - 🥦 **Eat a Balanced Diet**: Include fruits, vegetables, and whole grains.
            - 🚶 **Stay Active**: Aim for at least 30 minutes of exercise daily.
            - 💧 **Drink Water**: Stay hydrated by drinking 8+ glasses of water a day.
            - 😴 **Get Enough Sleep**: Adults need 7-9 hours of sleep per night.
            - 🧘 **Manage Stress**: Practice mindfulness and relaxation techniques.
            - 💉 **Stay Vaccinated**: Keep up with routine vaccinations and boosters.
            - 🚭 **Avoid Smoking & Alcohol**: Reduces risk of various diseases.
            - 🔍 **Regular Checkups**: Visit your doctor for periodic health screenings.
            """
        )