app.py

import os
import gradio as gr
import fitz  # PyMuPDF
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# --- 1. SETUP & SECRETS ---
# Ensure "GROQ_API_KEY" is added to Hugging Face Secrets
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

# Use Llama 3 for medical reasoning
llm = ChatGroq(
    temperature=0.1, 
    model_name="llama-3.3-70b-versatile", 
    api_key=GROQ_API_KEY
)

# Best open-source embeddings for medical text
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# --- 2. THE RAG LOGIC ---

def process_medical_report(pdf_file, user_question):
    if not pdf_file:
        return "⚠️ Error: Please upload a medical PDF report first."
    
    try:
        # Step A: PDF Text Extraction
        doc = fitz.open(pdf_file.name)
        text = ""
        for page in doc:
            text += page.get_text()
        
        if not text.strip():
            return "⚠️ Error: The PDF seems empty or is an image. Please provide a text-based PDF."

        # Step B: Recursive Chunking
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=700,
            chunk_overlap=150
        )
        chunks = text_splitter.split_text(text)
        
        # Step C: Temporary Vector Store
        vector_db = FAISS.from_texts(chunks, embeddings)
        retriever = vector_db.as_retriever(search_kwargs={"k": 3})

        # Step D: Prompt Engineering (Medical Cross-Border Specialist)
        medical_template = """
        You are a Medical Cross-Border AI Agent. Your goal is to help patients understand reports 
        from different countries.

        CONTEXT FROM REPORT:
        {context}

        INSTRUCTIONS:
        1. Summarize findings in simple terms.
        2. Detect URGENCY: Label as LOW, MEDIUM, or HIGH.
        3. Convert Units: If you see international units, explain them clearly.
        4. If a language other than English is requested, provide a high-quality translation.
        5. CRITICAL: If values are dangerous, tell the user to seek immediate care.

        USER QUESTION: {question}
        """

        prompt = ChatPromptTemplate.from_template(medical_template)

        # Step E: Modern LCEL Chain (Replaces RetrievalQA)
        def format_docs(docs):
            return "\n\n".join(doc.page_content for doc in docs)

        rag_chain = (
            {"context": retriever | format_docs, "question": RunnablePassthrough()}
            | prompt
            | llm
            | StrOutputParser()
        )

        # Execute
        return rag_chain.invoke(user_question)

    except Exception as e:
        return f"🚨 System Error: {str(e)}"

# --- 3. UI DESIGN (Gradio) ---

# Cyber-Luxe Theme setup
theme = gr.themes.Soft(
    primary_hue="cyan",
    secondary_hue="blue",
    neutral_hue="slate",
).set(
    button_primary_background_fill="*primary_500",
    button_primary_text_color="white",
)

with gr.Blocks(theme=theme) as demo:
    gr.HTML("<h1 style='text-align: center;'>🏥 AI Cross-Border Health Navigator</h1>")
    gr.Markdown("---")
    
    with gr.Row():
        with gr.Column(scale=1):
            file_upload = gr.File(label="1. Upload Medical PDF", file_types=[".pdf"])
            chat_query = gr.Textbox(
                label="2. Ask AI Agent", 
                placeholder="e.g. 'Summarize and check for urgency' or 'Translate to Urdu'",
                lines=2
            )
            analyze_btn = gr.Button("🚀 Start Analysis", variant="primary")
            
        with gr.Column(scale=2):
            output_display = gr.Markdown("### 🔍 AI Insights will appear here...")

    # Action
    analyze_btn.click(
        fn=process_medical_report, 
        inputs=[file_upload, chat_query], 
        outputs=output_display
    )

    gr.Markdown("---")
    gr.HTML("<p style='text-align: center; color: gray;'>Note: This is an AI prototype for hackathons. Not a substitute for professional medical advice.</p>")

if __name__ == "__main__":
    demo.launch()