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

import logging
import os
from dotenv import load_dotenv

import streamlit as st
from PIL import Image

from src.constants import SUPPORTED_LANGUAGES, MEDGEMMA_MODEL_ID
from src.fhir_generator import FhirGenerator, parse_medications_to_dict
from src.medgemma_service import MedGemmaService, load_medgemma_model
from src.sealion_service import SEALionService

# Load environment variables
load_dotenv()

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

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


@st.cache_resource
def get_medgemma():
    """Initialize local MedGemma service."""
    with st.spinner("Loading MedGemma model (this may take a moment)..."):
        model, processor = load_medgemma_model()
        return MedGemmaService(model, processor)


@st.cache_resource
def get_sealion_service():
    """Initialize SEA-LION API service."""
    with st.spinner("Connecting to SEA-LION API..."):
        return SEALionService()


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


def main():
    """Main application entry point."""
    # Header Section
    st.title("💊 Prescription Explainer")
    st.markdown(
        "Upload a prescription image to get a clear explanation in your language. "
        "Powered by MedGemma and SEA-LION AI."
    )
    st.markdown("---")

    # Create two columns layout
    left_col, right_col = st.columns([1, 1.5])

    # Left Column: Input Section
    with left_col:
        st.subheader("📝 Input")

        # 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
            process_button = st.button(
                "📋 Explain My Prescription",
                type="primary",
                use_container_width=True
            )
        else:
            process_button = False
            st.info("👆 Upload a prescription image to get started")

    # Right Column: Output Section
    with right_col:
        st.subheader("📖 Explanation")

        if uploaded_file is not None and process_button:
            try:
                # Load services
                medgemma = get_medgemma()
                sealion_service = get_sealion_service()
                fhir_generator = get_fhir_generator()

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

                # Step 2: Generate explanation using SEA-LION
                with st.spinner("Creating explanation..."):
                    from src.prompts import EXPLANATION_PROMPT
                    explanation_prompt = EXPLANATION_PROMPT.format(medication_info=extraction)
                    explanation = sealion_service.generate_text(explanation_prompt)

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

                # Display results
                st.success("✅ Analysis Complete!")
                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."
                )
        else:
            # Placeholder when no prescription uploaded
            st.info("Upload a prescription image to see the explanation here")
            st.markdown("""
            ### How it works:
            1. 🖼️ Upload your prescription image
            2. 🌐 Select your preferred language
            3. 🤖 AI analyzes the prescription
            4. 📖 Get easy-to-understand explanation
            5. 📤 Export to FHIR format (optional)
            """)

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


if __name__ == "__main__":
    main()