import os, json, streamlit as st
from backend.rag_engine import get_embedder, get_chroma, retrieve, seed_index
from backend.soap_generator import compose_soap
from backend.pdf_utils import generate_pdf
from backend.chat_textgen import chat
from utils.constants import DOCS_DIR, RETRIEVAL_K_DEFAULT

st.set_page_config(
    page_title="MediAssist — Clinical Decision Support",
    page_icon="🩺",
    layout="wide",
    initial_sidebar_state="expanded"
)

@st.cache_resource(show_spinner=False)
def _embedder(): 
    return get_embedder()

@st.cache_resource(show_spinner=False)
def _col(): 
    return get_chroma()[1]

# Sidebar Configuration
with st.sidebar:
    st.markdown("## ⚙️ Configuration")
    
    with st.expander("RAG Index Management", expanded=False):
        if st.button("🔄 Seed / Refresh RAG Index"):
            with st.spinner("Indexing medical guidelines..."):
                try:
                    n = seed_index(_col(), _embedder(), DOCS_DIR)
                    st.success(f"✅ Indexed {n} chunks from {DOCS_DIR}")
                except Exception as e:
                    st.error(f"❌ Indexing failed: {str(e)}")
        st.caption("Upload .txt/.md files to `data/guidelines/<specialty>/` then reseed.")
    
    st.divider()
    st.markdown("### About")
    st.info(
        "**MediAssist v15** combines clinical guidelines with AI to provide "
        "evidence-based decision support for healthcare professionals."
    )
    
    with st.expander("Features", expanded=False):
        st.markdown("""
        - 📚 RAG-based guideline retrieval
        - 🤖 SOAP note generation
        - 💬 Context-aware AI chat
        - 📄 PDF report generation
        """)

# Main Content
st.title("🩺 MediAssist — Clinical Decision Support System")
st.markdown(
    "AI-powered clinical guidelines with **RAG** • **SOAP** • **Chat** • **PDF Reports**"
)

# Main Input Section
st.markdown("### 📝 Patient Information")
narrative = st.text_area(
    "Patient Narrative",
    height=120,
    placeholder="e.g., 32-year-old female with 10 days of period delay, nausea, mild cramps. No fever. Medical history: PCOS.",
    help="Describe the patient's presenting complaint and relevant history"
)

col1, col2 = st.columns([3, 1])
with col1:
    k = st.slider(
        "Number of guidelines to retrieve",
        min_value=1,
        max_value=10,
        value=RETRIEVAL_K_DEFAULT,
        help="More results = more comprehensive but potentially noisy"
    )
with col2:
    st.metric("Retrieval K", k)

st.divider()

# Create Tabs for Different Functions
tab1, tab2, tab3, tab4 = st.tabs(["📊 SOAP Note", "💬 AI Chat", "📄 PDF Report", "📚 Guidelines"])

# Tab 1: SOAP Note Generation
with tab1:
    st.markdown("### Generate SOAP Note with Clinical Evidence")
    
    col_soap1, col_soap2 = st.columns(2)
    
    if st.button("🧾 Generate SOAP Note", key="soap_button"):
        if not narrative.strip():
            st.warning("⚠️ Please enter patient narrative first.")
        else:
            with st.spinner("Retrieving relevant guidelines..."):
                try:
                    items = retrieve(_col(), _embedder(), narrative, k=k)
                    soap = compose_soap(narrative, items)
                    
                    with col_soap1:
                        st.subheader("SOAP Note (JSON)")
                        st.code(json.dumps(soap, indent=2), language="json")
                        
                        # Copy button helper
                        st.caption("💡 Tip: Use the copy button in the code block to copy the JSON")
                    
                    with col_soap2:
                        st.subheader(f"📚 Citations ({len(items)} sources)")
                        if not items:
                            st.info("No relevant guidelines found.")
                        else:
                            for i, it in enumerate(items, 1):
                                with st.container(border=True):
                                    st.markdown(f"**Source {i}: {it.get('title', 'Unknown')}**")
                                    st.caption(f"📖 {it.get('source', 'N/A')}")
                                    st.markdown(f"> {it.get('text', '')[:350]}...")
                
                except Exception as e:
                    st.error(f"❌ Error generating SOAP note: {str(e)}")

# Tab 2: AI Chat
with tab2:
    st.markdown("### Conversational AI Assistant")
    
    col_chat1, col_chat2 = st.columns([1, 1])
    
    with col_chat1:
        mode = st.radio(
            "Chat Mode",
            ["Patient-facing explanation", "Doctor-facing analysis"],
            help="Choose audience for AI response"
        )
    
    if st.button("💬 Start Chat Session", key="chat_button"):
        if not narrative.strip():
            st.warning("⚠️ Please enter patient narrative first.")
        else:
            with st.spinner("AI is thinking..."):
                try:
                    reply = chat(
                        narrative,
                        mode="patient" if "Patient" in mode else "doctor"
                    )
                    st.markdown("### AI Response")
                    st.markdown(reply)
                except Exception as e:
                    st.error(f"❌ Chat error: {str(e)}")

# Tab 3: PDF Report
with tab3:
    st.markdown("### Generate Clinical PDF Report")
    
    col_pdf1, col_pdf2 = st.columns([2, 1])
    
    with col_pdf1:
        ai_summary = st.text_area(
            "Doctor-Reviewed Summary",
            height=100,
            placeholder="Enter clinical summary, assessment, and plan...",
            help="This will be included in the PDF report"
        )
    
    with col_pdf2:
        report_name = st.text_input(
            "Report Filename",
            value="MediAssist_Report",
            help="PDF will be saved as [filename].pdf"
        )
    
    if st.button("📄 Generate PDF Report", key="pdf_button"):
        if not narrative.strip():
            st.warning("⚠️ Please enter patient narrative first.")
        else:
            with st.spinner("Generating PDF..."):
                try:
                    items = retrieve(_col(), _embedder(), narrative, k=3)
                    soap = compose_soap(narrative, items)
                    pdf_path = f"{report_name}.pdf"
                    
                    generate_pdf(
                        pdf_path,
                        "MediAssist — Clinical Report",
                        soap,
                        ai_summary
                    )
                    
                    with open(pdf_path, "rb") as pdf_file:
                        st.download_button(
                            label="⬇️ Download PDF Report",
                            data=pdf_file,
                            file_name=pdf_path,
                            mime="application/pdf",
                            key="pdf_download"
                        )
                    st.success("✅ PDF generated successfully!")
                    
                except Exception as e:
                    st.error(f"❌ PDF generation error: {str(e)}")

# Tab 4: Guidelines Browse
with tab4:
    st.markdown("### Browse Medical Guidelines")
    
    if st.button("📚 Load Available Guidelines", key="guidelines_button"):
        with st.spinner("Loading guidelines..."):
            try:
                if os.path.exists(DOCS_DIR):
                    guidelines = []
                    for root, dirs, files in os.walk(DOCS_DIR):
                        for file in files:
                            if file.endswith(('.txt', '.md')):
                                guidelines.append(os.path.join(root, file))
                    
                    if guidelines:
                        st.info(f"Found {len(guidelines)} guidelines")
                        for guideline in sorted(guidelines)[:10]:  # Show first 10
                            st.caption(f"📄 {os.path.basename(guideline)}")
                    else:
                        st.warning("No guidelines found. Upload files to `data/guidelines/`")
                else:
                    st.error(f"Guidelines directory not found: {DOCS_DIR}")
            except Exception as e:
                st.error(f"❌ Error loading guidelines: {str(e)}")

st.divider()

# Footer
st.markdown("""
---
**Disclaimer:** MediAssist is a decision support tool and does not replace professional clinical judgment. 
Always consult with qualified healthcare professionals for medical decisions.

Made with ❤️ by the MediAssist Team | Deployed on 🤗 Hugging Face Spaces
""")