import streamlit as st
import pandas as pd
import plotly.express as px
import os
from openai import OpenAI
import bcrypt
from supabase import create_client, Client

# Set up Supabase client
supabase_url = st.secrets["SUPABASE_URL"]
supabase_key = st.secrets["SUPABASE_KEY"]
supabase: Client = create_client(supabase_url, supabase_key)

# Set up OpenAI client
client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])

# Initialize session state
if 'user' not in st.session_state:
    st.session_state.user = None
if 'user_type' not in st.session_state:
    st.session_state.user_type = None

def load_data(username=None):
    if username:
        response = supabase.table('entries').select('*').eq('username', username).execute()
    else:
        response = supabase.table('entries').select('*').execute()
    return pd.DataFrame(response.data)

def load_user_data():
    response = supabase.table('users').select('*').execute()
    return pd.DataFrame(response.data)

def save_data(entry):
    supabase.table('entries').insert(entry).execute()

def save_user_data(username, hashed_password, user_type):
    supabase.table('users').insert({
        'username': username,
        'password': hashed_password,
        'user_type': user_type
    }).execute()

def get_user(username):
    response = supabase.table('users').select('*').eq('username', username).execute()
    return pd.DataFrame(response.data)

def get_gpt_analysis(entry_text, system_prompt):
    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": entry_text}
            ]
        )
        return response.choices[0].message.content
    except Exception as e:
        st.error(f"Error in GPT analysis: {str(e)}")
        return "Analysis unavailable at this time."

def hash_password(password):
    return bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')

def verify_password(stored_password, provided_password):
    return bcrypt.checkpw(provided_password.encode('utf-8'), stored_password.encode('utf-8'))

def auth():
    if st.session_state.user is None:
        st.subheader("User Authentication")
        tabs = st.tabs(["Login", "Register"])
        
        with tabs[0]:
            st.subheader("Login")
            login_username = st.text_input("Username", key="login_username")
            login_password = st.text_input("Password", type="password", key="login_password")
            login_button = st.button("Login")
            
            if login_button:
                user_data = get_user(login_username)
                if not user_data.empty and verify_password(user_data.iloc[0]['password'], login_password):
                    st.session_state.user = login_username
                    st.session_state.user_type = user_data.iloc[0]['user_type']
                    st.success("Logged in successfully!")
                    st.rerun()
                else:
                    st.error("Invalid username or password.")
        
        with tabs[1]:
            st.subheader("Register")
            reg_username = st.text_input("Choose a Username", key="reg_username")
            reg_password = st.text_input("Choose a Password", type="password", key="reg_password")
            confirm_password = st.text_input("Confirm Password", type="password", key="confirm_password")
            user_type = st.selectbox("User Type", ["Patient", "Doctor"])
            register_button = st.button("Register")
            
            if register_button:
                existing_user = get_user(reg_username)
                if not existing_user.empty:
                    st.error("Username already exists. Please choose a different one.")
                elif reg_password != confirm_password:
                    st.error("Passwords do not match.")
                elif len(reg_password) < 8:
                    st.error("Password must be at least 8 characters long.")
                else:
                    hashed_password = hash_password(reg_password)
                    save_user_data(reg_username, hashed_password, user_type)
                    st.session_state.user = reg_username
                    st.session_state.user_type = user_type
                    st.success("Registered successfully!")
                    st.rerun()
    else:
        st.sidebar.write(f"Logged in as {st.session_state.user} ({st.session_state.user_type})")
        if st.sidebar.button("Logout"):
            st.session_state.user = None
            st.session_state.user_type = None
            st.rerun()

def main():
    st.set_page_config(page_title="Migraine Diary App", page_icon="🧠", layout="wide")
    st.title("Migraine Diary App")

    auth()

    if st.session_state.user:
        if st.session_state.user_type == "Patient":
            patient_interface()
        elif st.session_state.user_type == "Doctor":
            doctor_interface()

def patient_interface():
    menu = st.sidebar.selectbox("Menu", ["Add Entry", "View Entries", "Dashboard"])

    if menu == "Add Entry":
        add_entry()
    elif menu == "View Entries":
        view_entries(is_doctor=False)
    elif menu == "Dashboard":
        display_dashboard(is_doctor=False)

def doctor_interface():
    menu = st.sidebar.selectbox("Menu", ["View All Entries", "Patient Dashboard"])

    if menu == "View All Entries":
        view_entries(is_doctor=True)
    elif menu == "Patient Dashboard":
        display_dashboard(is_doctor=True)

def add_entry():
    st.header("Add New Migraine Entry")
    with st.form("migraine_entry"):
        date = st.date_input("Date")
        pain_level = st.slider("Pain Level", 1, 10)
        duration = st.selectbox("Duration", ["Less than 1 hour", "1-4 hours", "4-8 hours", "8-24 hours", "More than 24 hours"])
        
        triggers = st.multiselect("Triggers", [
            "Stress", "Lack of Sleep", "Dehydration", "Skipped Meals",
            "Alcohol", "Caffeine", "Chocolate", "Aged Cheeses",
            "Processed Meats", "Artificial Sweeteners", "MSG",
            "Weather Changes", "Barometric Pressure Changes",
            "Bright Lights", "Loud Noises", "Strong Smells",
            "Screen Time", "Reading", "Physical Exertion",
            "Hormonal Changes", "Medication Overuse",
            "Travel", "Altitude Changes", "Other"
        ])
        
        symptoms = st.multiselect("Symptoms", [
            "Throbbing Pain", "Pulsating Pain", "One-sided Pain",
            "Nausea", "Vomiting", "Sensitivity to Light",
            "Sensitivity to Sound", "Sensitivity to Smells",
            "Blurred Vision", "Visual Aura", "Blind Spots",
            "Zigzag Lines in Vision", "Tingling or Numbness",
            "Difficulty Speaking", "Weakness", "Dizziness",
            "Vertigo", "Neck Stiffness", "Confusion",
            "Mood Changes", "Food Cravings", "Frequent Urination",
            "Fatigue", "Yawning", "Other"
        ])
        
        medications = st.text_input("Medications taken")
        notes = st.text_area("Additional Notes")
        
        submitted = st.form_submit_button("Submit Entry")
        if submitted:
            entry_text = f"Date: {date}\nPain Level: {pain_level}\nDuration: {duration}\nTriggers: {', '.join(triggers)}\nSymptoms: {', '.join(symptoms)}\nMedications: {medications}\nNotes: {notes}"
            
            with st.spinner("Analyzing your entry..."):
                doctor_analysis = get_gpt_analysis(entry_text, "You are a neurologist specializing in migraine management. Provide a technical analysis of the patient's migraine diary entry, including potential correlations, patterns, and suggestions for the treating physician. Keep it short and to the point the doctor is busy.")
                patient_advice = get_gpt_analysis(entry_text, "You are a supportive health coach specializing in migraine management. Provide friendly, easy-to-understand advice for the patient based on their migraine diary entry. Include actionable tips for managing their condition and potential lifestyle adjustments.")

            new_entry = {
                'username': st.session_state.user,
                'entry_date': date.isoformat(),
                'pain_level': pain_level,
                'duration': duration,
                'triggers': ', '.join(triggers),
                'symptoms': ', '.join(symptoms),
                'medications': medications,
                'notes': notes,
                'doctor_analysis': doctor_analysis,
                'patient_advice': patient_advice
            }
            save_data(new_entry)
            st.success("Entry added successfully!")
            st.subheader("Advice for You:")
            st.write(patient_advice)

def view_entries(is_doctor):
    st.header("Migraine Entries")
    if is_doctor:
        user_entries = load_data()
        st.subheader("All Patient Entries")
    else:
        user_entries = load_data(st.session_state.user)
        st.subheader("Your Entries")
    
    user_entries = user_entries.sort_values(by='entry_date', ascending=False)
    
    if not user_entries.empty:
        for _, entry in user_entries.iterrows():
            with st.expander(f"Entry for {entry['username']} on {entry['entry_date']} - Pain Level: {entry['pain_level']}"):
                st.write(f"Duration: {entry['duration']}")
                st.write(f"Triggers: {entry['triggers']}")
                st.write(f"Symptoms: {entry['symptoms']}")
                st.write(f"Medications: {entry['medications']}")
                st.write(f"Notes: {entry['notes']}")
                if is_doctor:
                    st.write("Doctor's Analysis:", entry['doctor_analysis'])
                else:
                    st.write("Advice for Patient:", entry['patient_advice'])
    else:
        st.info("No entries found.")

def display_dashboard(is_doctor):
    st.header("Migraine Dashboard")
    
    if is_doctor:
        st.subheader("Select Patient")
        all_users = load_data()['username'].unique()
        selected_user = st.selectbox("Choose a patient", all_users)
        user_entries = load_data(selected_user)
    else:
        user_entries = load_data(st.session_state.user)

    if not user_entries.empty:
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Pain Level Over Time")
            fig = px.line(user_entries, x='entry_date', y='pain_level', title='Pain Level Over Time')
            st.plotly_chart(fig, use_container_width=True)

        with col2:
            st.subheader("Common Triggers")
            all_triggers = ', '.join(user_entries['triggers'].dropna()).split(', ')
            trigger_counts = pd.Series(all_triggers).value_counts().head(5)
            fig = px.bar(x=trigger_counts.index, y=trigger_counts.values, labels={'x': 'Trigger', 'y': 'Count'})
            st.plotly_chart(fig, use_container_width=True)

        col1, col2 = st.columns(2)

        with col1:
            st.subheader("Common Symptoms")
            all_symptoms = ', '.join(user_entries['symptoms'].dropna()).split(', ')
            symptom_counts = pd.Series(all_symptoms).value_counts().head(5)
            fig = px.bar(x=symptom_counts.index, y=symptom_counts.values, labels={'x': 'Symptom', 'y': 'Count'})
            st.plotly_chart(fig, use_container_width=True)

        st.subheader("Migraine Statistics")
        col1, col2, col3, col4 = st.columns(4)
        col1.metric("Total Entries", len(user_entries))
        col2.metric("Average Pain Level", f"{user_entries['pain_level'].mean():.2f}")
        col3.metric("Most Common Trigger", trigger_counts.index[0] if not trigger_counts.empty else "N/A")
        col4.metric("Most Common Symptom", symptom_counts.index[0] if not symptom_counts.empty else "N/A")

        st.subheader("Recent Entries")
        st.dataframe(user_entries[['entry_date', 'pain_level', 'duration', 'triggers', 'symptoms']].sort_values(by='entry_date', ascending=False).head())
    else:
        st.info("No entries found.")

if __name__ == "__main__":
    main()