import gradio as gr
from rag_pipeline import HealthcareRAG
import os

# Initialize RAG system
rag = HealthcareRAG()

def process_query(query: str) -> tuple:
    """Process user query and return response with retrieved chunks."""
    result = rag.query(query)
    print(result['retrieved_chunks'])
    # Format retrieved chunks for display
    chunks_text = "\n\n---\n\n".join(result["retrieved_chunks"])
    
    return result["response"], chunks_text

# Create Gradio interface
with gr.Blocks(title="Healthcare Guidelines Assistant") as demo:
    gr.Markdown("""
    # Healthcare Guidelines Assistant
    
    This assistant provides information based on clinical guidelines. 
    Please note that this is for educational purposes only and not medical advice.
    
    **DISCLAIMER**: This information is for educational purposes only and not medical advice.
    Always consult with healthcare professionals for medical decisions.
    """)
    
    with gr.Row():
        with gr.Column():
            query_input = gr.Textbox(
                label="Your Question",
                placeholder="Enter your question about clinical guidelines...",
                lines=3
            )
            submit_btn = gr.Button("Submit")
        
    with gr.Row():
        with gr.Column():
            response_output = gr.Textbox(
                label="Response",
                lines=5
            )
            chunks_output = gr.Textbox(
                label="Retrieved Guidelines",
                lines=10
            )
    
    submit_btn.click(
        fn=process_query,
        inputs=query_input,
        outputs=[response_output, chunks_output]
    )

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