import streamlit as st
import pandas as pd
import numpy as np
import os
import hashlib
from datetime import datetime
from groq import Groq
import pdfplumber
import plotly.graph_objects as go
import folium
from streamlit_folium import st_folium
from streamlit_geolocation import streamlit_geolocation

# --- 1. CORE SYSTEM CONFIG ---
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
if not GROQ_API_KEY:
    st.error("⚠️ API Key Missing! Please check your environment variables.")
    st.stop()

st.set_page_config(page_title="IntelliCare Portal | Hassan Naseer", layout="wide", page_icon="🏥")

# --- 2. PREMIUM CSS ---
st.markdown("""
<style>
    [data-testid="stSidebar"] { background: #1e3a8a !important; }
    [data-testid="stSidebar"] * { color: #ffffff !important; }
    div.stButton > button {
        background-color: #1e3a8a !important; color: white !important;
        border: 1px solid #3b82f6 !important; transition: none !important;
    }
    .user-msg {
        background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%);
        color: white; padding: 15px; border-radius: 18px 18px 2px 18px;
        margin-bottom: 15px; margin-left: 20%; box-shadow: 0 4px 10px rgba(37,99,235,0.2);
    }
    .ai-msg {
        background: #f1f5f9; color: #1e293b; padding: 15px; border-radius: 18px 18px 18px 2px;
        margin-bottom: 15px; margin-right: 20%; border: 1px solid #e2e8f0;
    }
</style>
""", unsafe_allow_html=True)

# --- 3. DATA UTILS ---
USER_DB, CALL_LOG_DB, CALL_SIGNAL_DB = "users_secure.csv", "call_history.csv", "active_calls.csv"
def hash_pass(pwd): return hashlib.sha256(str.encode(pwd)).hexdigest()
def load_db(file, cols):
    if os.path.exists(file): return pd.read_csv(file)
    return pd.DataFrame(columns=cols)

# Session State Initialization
if "logged_in" not in st.session_state: st.session_state.logged_in = False
if "msgs" not in st.session_state: st.session_state.msgs = []
if "active_doc" not in st.session_state: st.session_state.active_doc = ""
if "last_voice_hash" not in st.session_state: st.session_state.last_voice_hash = None
if "audio_key" not in st.session_state: st.session_state.audio_key = 0 # To reset Mic

# --- 4. AUTHENTICATION ---
if not st.session_state.logged_in:
    st.markdown("<h1 style='text-align: center; color: #1e3a8a;'>🏥 IntelliCare Portal</h1>", unsafe_allow_html=True)
    c2 = st.columns([1, 2, 1])[1]
    with c2:
        tab1, tab2 = st.tabs(["🔐 Login", "📝 Create Account"])
        with tab1:
            u, p = st.text_input("Username"), st.text_input("Password", type="password")
            if st.button("Sign In"):
                db = load_db(USER_DB, ["username", "password", "role"])
                match = db[(db['username'] == u) & (db['password'] == hash_pass(p))]
                if not match.empty:
                    st.session_state.logged_in, st.session_state.username = True, u
                    st.session_state.role = match.iloc[0]['role']
                    st.rerun()
        with tab2:
            nu, np = st.text_input("New ID"), st.text_input("New Pass", type="password")
            nr = st.selectbox("Role", ["Patient", "Doctor"])
            if st.button("Register"):
                df = load_db(USER_DB, ["username", "password", "role"])
                pd.concat([df, pd.DataFrame([{"username": nu, "password": hash_pass(np), "role": nr}])]).to_csv(USER_DB, index=False)
                st.success("Registered!")
    st.stop()

# --- 5. SIDEBAR NAVIGATION ---
with st.sidebar:
    st.markdown(f"### 👤 {st.session_state.username}")
    lang = st.radio("🌐 Language", ["English", "Urdu"])
    
    # PERMANENT DOWNLOAD
    chat_txt = "\n".join([f"{m['role']}: {m['content']}" for m in st.session_state.msgs])
    st.download_button("📥 Download History", data=chat_txt, file_name="chat.txt")
    
    if st.button("🗑️ Clear Chat"): 
        st.session_state.msgs, st.session_state.active_doc = [], ""
        st.rerun()
    if st.button("🚪 Logout"): st.session_state.logged_in = False; st.rerun()
    st.divider()
    
    if st.session_state.role == "Patient":
        nav = st.radio("Menu", ["💬 AI Chat", "🧪 Health Lab", "📍 Nearby Clinics", "📞 Video Consult", "📜 History"])
    else:
        nav = st.radio("Menu", ["🖥️ Consultation Desk", "📜 Call Logs"])

# --- 6. AI CHAT MODULE (FIXED MIC & INPUT) ---
if nav == "💬 AI Chat":
    st.markdown("### 💬 Clinical Intelligence Assistant")
    with st.expander("📄 Upload Report/Prescription"):
        up = st.file_uploader("Drop PDF", type=['pdf'], label_visibility="collapsed")
        if up:
            with pdfplumber.open(up) as f:
                st.session_state.active_doc = " ".join([p.extract_text() for p in f.pages if p.extract_text()])
                st.success("✅ Analysis Ready!")

    for m in st.session_state.msgs:
        st.markdown(f'<div class="{"user-msg" if m["role"]=="user" else "ai-msg"}">{m["content"]}</div>', unsafe_allow_html=True)
    
    st.divider()
    col_v, col_q = st.columns([1, 9])
    with col_v: # FIXED MIC RESET
        v = st.audio_input("🎤", key=f"mic_{st.session_state.audio_key}", label_visibility="collapsed")
    with col_q: q = st.chat_input("Ask about your health...")

    final_q = q if q else None
    if v and not final_q:
        vh = hashlib.md5(v.getvalue()).hexdigest()
        if vh != st.session_state.last_voice_hash:
            final_q = Groq(api_key=GROQ_API_KEY).audio.transcriptions.create(file=("a.wav", v.getvalue()), model="whisper-large-v3", response_format="text")
            st.session_state.last_voice_hash = vh
            st.session_state.audio_key += 1 # Forces Mic refresh

    if final_q:
        st.session_state.msgs.append({"role": "user", "content": final_q})
        sys_p = f"Professional Medical Assistant. Language: {lang}. Analyze the PDF text provided. Refuse non-medical tasks."
        ans = Groq(api_key=GROQ_API_KEY).chat.completions.create(model="llama-3.3-70b-versatile", messages=[{"role": "system", "content": sys_p}, {"role": "system", "content": f"PDF: {st.session_state.active_doc}"}] + st.session_state.msgs)
        st.session_state.msgs.append({"role": "assistant", "content": ans.choices[0].message.content})
        st.rerun()

# --- 7. DIAGNOSTIC TOOLS (3 TOOLS RESTORED) ---
elif nav == "🧪 Health Lab":
    st.markdown("### 🧪 Interactive Clinical Lab")
    tool = st.selectbox("Select Diagnostic Tool", ["⚖️ BMI Analyzer", "🩸 Glucose Tracker", "🫀 Heart Rate Sim"])
    
    if tool == "⚖️ BMI Analyzer":
        
        w, h = st.number_input("Weight (kg)", 30, 200, 70), st.number_input("Height (cm)", 100, 250, 175)
        bmi = round(w / ((h/100)**2), 1)
        st.metric("Your BMI", bmi)
        st.plotly_chart(go.Figure(go.Indicator(mode="gauge+number", value=bmi, gauge={'axis': {'range': [10, 40]}, 'bar': {'color': "#3b82f6"}})), use_container_width=True)

    elif tool == "🩸 Glucose Tracker":
        
        st.write("Recent Blood Sugar Levels (mg/dL)")
        df = pd.DataFrame({'Time': range(1, 11), 'Level': np.random.randint(80, 150, 10)})
        st.plotly_chart(go.Figure(data=go.Scatter(x=df['Time'], y=df['Level'], mode='lines+markers', line=dict(color='#2563eb'))), use_container_width=True)

    elif tool == "🫀 Heart Rate Sim":
        
        bpm = st.slider("Select Current BPM", 40, 180, 72)
        st.metric("Heart Rate", f"{bpm} BPM")
        t = np.linspace(0, 2, 1000)
        heart_wave = np.sin(2 * np.pi * (bpm/60) * t) + 0.1 * np.random.normal(0, 1, 1000)
        st.plotly_chart(go.Figure(data=go.Scatter(x=t, y=heart_wave, line=dict(color='#ef4444'))), use_container_width=True)

# --- 8. MAPS & VIDEO CALLS (RESTORED) ---
elif nav == "📍 Nearby Clinics":
    st.markdown("### 📍 Specialist Hospital Locator")
    loc = streamlit_geolocation()
    if loc.get("latitude"):
        m = folium.Map(location=[loc["latitude"], loc["longitude"]], zoom_start=14)
        folium.Marker([loc["latitude"], loc["longitude"]], popup="You").add_to(m)
        st_folium(m, width=1000, height=500)

elif nav == "📞 Video Consult":
    db = load_db(USER_DB, ["username", "role"])
    docs = db[db['role'] == 'Doctor']['username'].tolist()
    sel_doc = st.selectbox("Select Specialist", docs)
    if st.button("Request Video Consultation"):
        room = f"IntelliCare-{st.session_state.username}-{sel_doc}"
        pd.DataFrame([{"Caller": st.session_state.username, "Receiver": sel_doc, "RoomID": room, "Status": "Active"}]).to_csv(CALL_SIGNAL_DB, index=False)
        pd.concat([load_db(CALL_LOG_DB, ["Time", "Caller", "Receiver", "RoomID"]), pd.DataFrame([{"Time": datetime.now().strftime("%Y-%m-%d %H:%M"), "Caller": st.session_state.username, "Receiver": sel_doc, "RoomID": room}])]).to_csv(CALL_LOG_DB, index=False)
        st.session_state.p_room = room
    if "p_room" in st.session_state:
        st.components.v1.html(f'<iframe src="https://meet.jit.si/{st.session_state.p_room}" width="100%" height="600px"></iframe>', height=650)

elif nav == "🖥️ Consultation Desk":
    signals = load_db(CALL_SIGNAL_DB, ["Caller", "Receiver", "RoomID"])
    my_calls = signals[signals['Receiver'] == st.session_state.username]
    if not my_calls.empty:
        st.info(f"🔔 Request from {my_calls.iloc[0]['Caller']}")
        c1, c2 = st.columns(2)
        with c1:
            if st.button("✅ Accept"): st.session_state.d_room = my_calls.iloc[0]['RoomID']
        with c2:
            if st.button("❌ Decline"): 
                load_db(CALL_SIGNAL_DB, ["Caller", "Receiver", "RoomID", "Status"])[load_db(CALL_SIGNAL_DB, ["Caller", "Receiver", "RoomID", "Status"])['Receiver'] != st.session_state.username].to_csv(CALL_SIGNAL_DB, index=False)
                st.rerun()
    if "d_room" in st.session_state:
        st.components.v1.html(f'<iframe src="https://meet.jit.si/{st.session_state.d_room}" width="100%" height="600px"></iframe>', height=650)

elif nav in ["📜 History", "📜 Call Logs"]:
    st.dataframe(load_db(CALL_LOG_DB, ["Time", "Caller", "Receiver", "RoomID"]), use_container_width=True)