app.py

import streamlit as st
import base64
import requests
from PIL import Image
from io import BytesIO
from pdf2image import convert_from_bytes
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
import schedule
import time
import threading

# Configuration - Replace with your API key
GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
GEMINI_MODEL = "gemini-2.0-flash"

DOCUMENT_TYPES = [
    "Insurance Policies", "Explanation of Benefits (EOBs)",
    "Claims (Approved, Denied, or Pending)", "Visit Summaries",
    "Test Results (Lab Reports, Imaging Reports)", 
    "Prescriptions (E-Prescriptions, Handwritten)",
    "Discharge Summaries", "Medical Bills", "Payment Statements",
    "Pharmacy Receipts", "Prior Authorization Requests",
    "Consent Forms", "Referral Letters", "Others"
]

def initialize_session_state():
    """Initialize all session state variables"""
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []
    if "processed_doc" not in st.session_state:
        st.session_state.processed_doc = None
    if "doc_preview" not in st.session_state:
        st.session_state.doc_preview = None
    if "reminder_set" not in st.session_state:
        st.session_state.reminder_set = False

def encode_file(uploaded_file):
    """Safely encode different file types to base64"""
    try:
        file_bytes = uploaded_file.getvalue()
        
        if uploaded_file.type == "application/pdf":
            images = convert_from_bytes(file_bytes, first_page=1, last_page=1)
            if not images:
                raise ValueError("Failed to convert PDF to image")
            img_byte_arr = BytesIO()
            images[0].save(img_byte_arr, format='JPEG')
            return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
        
        return base64.b64encode(file_bytes).decode('utf-8')
    except Exception as e:
        st.error(f"File processing error: {str(e)}")
        return None

def query_gemini(prompt, image_b64=None):
    """Handle Gemini API communication"""
    url = f"https://generativelanguage.googleapis.com/v1/models/{GEMINI_MODEL}:generateContent?key={GEMINI_API_KEY}"
    
    parts = [{"text": prompt}]
    if image_b64:
        parts.append({
            "inline_data": {
                "mime_type": "image/jpeg",
                "data": image_b64
            }
        })
    
    try:
        response = requests.post(
            url,
            json={"contents": [{"parts": parts}]},
            headers={"Content-Type": "application/json"},
            timeout=30
        )
        response.raise_for_status()
        return response.json()["candidates"][0]["content"]["parts"][0]["text"]
    except Exception as e:
        st.error(f"API Error: {str(e)}")
        return None

def process_document():
    """Handle document processing pipeline"""
    uploaded_file = st.session_state.uploaded_file
    if not uploaded_file:
        return
    
    try:
        with st.spinner("Analyzing document..."):
            # Convert to base64
            image_b64 = encode_file(uploaded_file)
            if not image_b64:
                return

            # Generate preview
            if uploaded_file.type == "application/pdf":
                images = convert_from_bytes(uploaded_file.getvalue(), first_page=1, last_page=1)
                st.session_state.doc_preview = images[0]
            else:
                st.session_state.doc_preview = Image.open(uploaded_file)

            # Classify document
            classify_prompt = f"Classify this healthcare document into one of: {DOCUMENT_TYPES}. Respond only with the category name."
            doc_type = query_gemini(classify_prompt, image_b64) or "Others"
            
            # Store results
            st.session_state.processed_doc = {
                "type": doc_type,
                "content": image_b64,
                "summary": query_gemini("Create a detailed structured summary of this healthcare document.", image_b64)
            }
            
    except Exception as e:
        st.error(f"Processing failed: {str(e)}")
        st.session_state.processed_doc = None

def handle_chat_query():
    """Process user chat input"""
    user_input = st.session_state.chat_input
    if not user_input or not st.session_state.processed_doc:
        return
    
    prompt = f"""
    Document Context:
    - Type: {st.session_state.processed_doc['type']}
    - Summary: {st.session_state.processed_doc['summary']}
    
    Question: {user_input}
    
    Answer concisely and factually. If unsure, state "Information not found".
    """
    
    with st.spinner("Generating response..."):
        response = query_gemini(prompt, st.session_state.processed_doc['content']) 
    
    st.session_state.chat_history.append(("user", user_input))
    st.session_state.chat_history.append(("assistant", response or "Could not generate response"))

def send_email_reminder(email, message):
    """Send an email reminder using ProtonMail Bridge"""
    try:
        sender_email = st.secrets["PROTONMAIL_USERNAME"]
        sender_password = st.secrets["PROTONMAIL_PASSWORD"]
        
        msg = MIMEMultipart()
        msg['From'] = sender_email
        msg['To'] = email
        msg['Subject'] = "Medicine Reminder"
        
        msg.attach(MIMEText(message, 'plain'))
        
        with smtplib.SMTP('127.0.0.1', 1025) as smtp:
            smtp.starttls()
            smtp.login(sender_email, sender_password)
            smtp.sendmail(sender_email, email, msg.as_string())
        
        st.success("Email reminder sent successfully!")
    except Exception as e:
        st.error(f"Failed to send email: {str(e)}")

def schedule_reminder(email, time_str, message):
    """Schedule a medicine reminder"""
    try:
        schedule.every().day.at(time_str).do(send_email_reminder, email, message)
        
        # Run the scheduler in a separate thread
        def run_scheduler():
            while True:
                schedule.run_pending()
                time.sleep(1)
        
        threading.Thread(target=run_scheduler, daemon=True).start()
        
        st.session_state.reminder_set = True
        st.success("Reminder set successfully!")
    except Exception as e:
        st.error(f"Failed to set reminder: {str(e)}")

# UI Layout
def main():
    st.set_page_config(page_title="Healthcare Document Assistant", layout="wide")
    initialize_session_state()

    # Sidebar Section
    with st.sidebar:
        st.header("Document Management")
        
        # Preview above upload button
        if st.session_state.doc_preview:
            st.subheader("Preview")
            st.image(st.session_state.doc_preview, use_container_width=True)
        
        # Upload button
        st.file_uploader(
            "Upload Document",
            type=["pdf", "png", "jpg", "jpeg"],
            key="uploaded_file",
            on_change=process_document
        )
        
        # Document type below upload button
        if st.session_state.processed_doc:
            st.divider()
            st.subheader("Document Type")
            st.markdown(f"**{st.session_state.processed_doc['type']}**")
        
        # Medicine Reminder
        st.divider()
        st.subheader("Medicine Reminder")
        
        if not st.session_state.reminder_set:
            email = st.text_input("Enter your email")
            time_str = st.text_input("Enter reminder time (24-hour format, e.g., 14:30)")
            message = st.text_area("Enter reminder message")
            
            if st.button("Set Reminder"):
                schedule_reminder(email, time_str, message)
        else:
            st.success("Reminder is already set. You will receive an email at the specified time.")

    # Main Content
    st.title("Healthcare Document Assistant")
    
    if st.session_state.processed_doc:
        # Document Summary
        st.subheader("Document Summary")
        st.markdown(st.session_state.processed_doc['summary'])
        
        # Chat Interface
        st.divider()
        st.subheader("Document Q&A")
        
        # Chat history
        for role, message in st.session_state.chat_history:
            with st.chat_message(role.capitalize()):
                st.markdown(message)
        
        # Chat input
        st.chat_input(
            "Ask about the document...",
            key="chat_input",
            on_submit=handle_chat_query
        )
    else:
        st.info("Please upload a document to begin analysis")

if __name__ == "__main__":
    main()