app.py

"""Prescription Explainer - AI-powered prescription understanding for SE Asia."""

import logging

import streamlit as st
from PIL import Image

from src.constants import SUPPORTED_LANGUAGES
from src.fhir_generator import FhirGenerator, parse_medications_to_dict
from src.medgemma_service import MedGemmaService, load_medgemma_model
from src.translation_service import TranslationService, load_translation_model

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Page configuration
st.set_page_config(
    page_title="Prescription Explainer",
    page_icon="💊",
    layout="centered",
    initial_sidebar_state="collapsed",
)


@st.cache_resource
def get_medgemma():
    """Load and cache MedGemma model."""
    with st.spinner("Loading AI model... This may take a moment on first run."):
        model, processor = load_medgemma_model()
        return MedGemmaService(model, processor)


@st.cache_resource
def get_translation_service():
    """Load and cache translation model."""
    with st.spinner("Loading translation model..."):
        model, tokenizer = load_translation_model()
        return TranslationService(model, tokenizer)


@st.cache_resource
def get_fhir_generator():
    """Get FHIR generator instance."""
    return FhirGenerator()


def main():
    """Main application entry point."""
    # Header
    st.title("💊 Prescription Explainer")
    st.markdown(
        "Upload a prescription image to get a clear explanation in your language."
    )

    # Language selector
    selected_language = st.selectbox(
        "Choose your language / เลือกภาษา / Pilih bahasa",
        options=list(SUPPORTED_LANGUAGES.keys()),
        index=0,
    )

    # File uploader
    uploaded_file = st.file_uploader(
        "Upload prescription image",
        type=["jpg", "jpeg", "png", "webp"],
        help="Take a photo of your prescription or upload an existing image",
    )

    if uploaded_file is not None:
        # Display uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="Your prescription", use_container_width=True)

        # Process button
        if st.button("📋 Explain My Prescription", type="primary", use_container_width=True):
            try:
                # Load services
                medgemma = get_medgemma()
                translation_service = get_translation_service()
                fhir_generator = get_fhir_generator()

                # Step 1: Extract medications
                with st.spinner("Reading your prescription..."):
                    extraction = medgemma.extract_medications(image)

                # Step 2: Generate explanation using translation model
                with st.spinner("Creating explanation..."):
                    from src.prompts import EXPLANATION_PROMPT
                    explanation_prompt = EXPLANATION_PROMPT.format(medication_info=extraction)
                    explanation = translation_service.generate_text(explanation_prompt)

                # Step 3: Translate if needed
                if selected_language != "English":
                    with st.spinner(f"Translating to {selected_language}..."):
                        explanation = translation_service.translate_text(
                            explanation, selected_language
                        )

                # Display results
                st.success("Done!")
                st.markdown("---")
                st.subheader("📖 Your Prescription Explained")
                st.markdown(explanation)

                # FHIR Export section
                st.markdown("---")
                st.subheader("📤 Export Health Data (FHIR)")

                try:
                    medications = parse_medications_to_dict(extraction)

                    col1, col2 = st.columns(2)

                    with col1:
                        # Generate MedicationStatement for each medication
                        statements = []
                        for med in medications:
                            try:
                                statements.append(fhir_generator.generate_medication_statement(med))
                            except Exception as e:
                                logger.warning(f"Skipping FHIR generation for {med.get('drug_name', 'unknown')}: {e}")

                        if statements:
                            combined_statements = "[\n" + ",\n".join(statements) + "\n]"

                            st.download_button(
                                label="💾 MedicationStatement (JSON)",
                                data=combined_statements,
                                file_name="medication_statement.json",
                                mime="application/json",
                            )
                        else:
                            st.info("FHIR export not available for this prescription")

                    with col2:
                        # Generate MedicationRequest for each medication
                        requests = []
                        for med in medications:
                            try:
                                requests.append(fhir_generator.generate_medication_request(med))
                            except Exception as e:
                                logger.warning(f"Skipping FHIR generation for {med.get('drug_name', 'unknown')}: {e}")

                        if requests:
                            combined_requests = "[\n" + ",\n".join(requests) + "\n]"

                            st.download_button(
                                label="💾 MedicationRequest (JSON)",
                                data=combined_requests,
                                file_name="medication_request.json",
                                mime="application/json",
                            )
                        else:
                            st.info("FHIR export not available for this prescription")

                    st.caption(
                        "FHIR R4 compliant files for sharing with healthcare providers"
                    )
                except Exception as e:
                    logger.error(f"FHIR export failed: {e}")
                    st.warning("FHIR export not available for this prescription")

            except ValueError as e:
                st.error(str(e))
            except Exception as e:
                logger.error(f"Processing failed: {e}")
                st.error(
                    "Something went wrong. Please try again with a clearer image."
                )

    # Footer
    st.markdown("---")
    st.caption(
        "⚠️ This tool provides general information only. "
        "Always consult your healthcare provider for medical advice."
    )
    st.caption("Built with MedGemma for the Google AI Hackathon 2026")


if __name__ == "__main__":
    main()