import os
import streamlit as st
from config import STATIC_DIR, HF_TOKEN, GOOGLE_API_KEY, DEVICE

# App Configuration
st.set_page_config(page_title="RxGuard Prescription Validator", page_icon="⚕️", layout="wide")

# Initialize directories and session state
UPLOADS_DIR = os.path.join(STATIC_DIR, "uploads")
os.makedirs(UPLOADS_DIR, exist_ok=True)

if "analysis_result" not in st.session_state:
    st.session_state.analysis_result = None
if "uploaded_filename" not in st.session_state:
    st.session_state.uploaded_filename = None

def show_service_status():
    """Displays service connectivity status."""
    st.caption("Service Status")
    cols = st.columns(3)
    cols[0].metric("HuggingFace Models", "✅" if HF_TOKEN else "❌")
    cols[1].metric("Google AI Services", "✅" if GOOGLE_API_KEY else "❌")
    cols[2].metric("Hardware Accelerator", DEVICE.upper())
    st.divider()

def main():
    st.title("⚕️ RxGuard Prescription Validator")
    st.caption("Advanced, multi-source AI verification system")
    show_service_status()

    # Only enable upload if required services are available
    if all([HF_TOKEN, GOOGLE_API_KEY]):
        uploaded_file = st.file_uploader(
            "Upload a prescription image (PNG/JPG/JPEG):",
            type=["png", "jpg", "jpeg"],
            help="Upload a clear image of the prescription for analysis."
        )

        if uploaded_file and uploaded_file.name != st.session_state.uploaded_filename:
            with st.status("Analyzing prescription...", expanded=True) as status:
                try:
                    st.session_state.uploaded_filename = uploaded_file.name
                    file_path = os.path.join(UPLOADS_DIR, uploaded_file.name)
                    with open(file_path, "wb") as f:
                        f.write(uploaded_file.getvalue())

                    # Lazily import the processing function
                    from validate_prescription import extract_prescription_info
                    st.session_state.analysis_result = extract_prescription_info(file_path)
                    status.update(label="Analysis complete!", state="complete", expanded=False)
                except Exception as e:
                    st.error(f"A critical error occurred during processing: {str(e)}")
                    st.session_state.analysis_result = {"error": str(e)}
                    status.update(label="Analysis failed", state="error")

    else:
        st.error("Missing API Keys. Please configure HF_TOKEN and GOOGLE_API_KEY in your Space secrets.")

    # Display results if available in the session state
    if result := st.session_state.get("analysis_result"):
        if error := result.get("error"):
            st.error(f"❌ Analysis Error: {error}")
        else:
            info = result.get("info", {})
            tab1, tab2 = st.tabs(["**👤 Patient & Prescription Info**", "**⚙️ Technical Details**"])

            with tab1:
                col1, col2 = st.columns([1, 2])
                with col1:
                    if uploaded_file:
                        st.image(uploaded_file, use_column_width=True, caption="Uploaded Prescription")
                with col2:
                    st.subheader("Patient Details")
                    st.info(f"**Name:** {info.get('Name', 'Not detected')}")
                    st.info(f"**Age:** {info.get('Age', 'N/A')}")
                    st.subheader("Prescription Details")
                    st.info(f"**Date:** {info.get('Date', 'N/A')}")
                    st.info(f"**Physician:** {info.get('PhysicianName', 'N/A')}")

                st.divider()
                st.subheader("💊 Medications")
                for med in info.get("Medications", []):
                    st.success(f"**Drug:** {med.get('drug_raw')} | **Dosage:** {med.get('dosage', 'N/A')} | **Frequency:** {med.get('frequency', 'N/A')}")

            with tab2:
                st.subheader("Debug Information from AI Pipeline")
                st.json(result.get("debug_info", {}))

if __name__ == "__main__":
    main()