import gradio as gr
from rag_dspy import MedicalRAG

rag_chain = MedicalRAG()

sample_questions = [
    "What are the most common symptoms of lupus?",
    "How is type 2 diabetes usually managed in adults?",
    "What are the first-line medications for treating hypertension?",
]

def qa_bot(user_question, history, year, specialty):
    history = history or []
    if not user_question.strip():
        return history, history
    history.append({"role": "user", "content": user_question})
    result = rag_chain.forward(user_question, year, specialty)
    answer = result.final_answer
    history.append({"role": "assistant", "content": answer})
    return history, history

css = """
body {
    background: #f5f7fa !important;
    font-family: 'Segoe UI', 'Roboto', 'Helvetica Neue', Arial, sans-serif !important;
    color: #212529 !important;
}
#main-card {
    background: #fff !important;
    border-radius: 18px;
    box-shadow: 0 4px 24px rgba(60, 80, 120, 0.07), 0 1.5px 4px rgba(60, 80, 120, 0.05);
    padding: 32px 32px 24px 32px;
    max-width: 750px;
    margin: 40px auto 0 auto;
    border: 1px solid #e3e7ef;
}
#chatbot {
    background: #f9fbfd !important;
    border-radius: 14px !important;
    box-shadow: 0 2px 8px rgba(60, 80, 120, 0.06);
    margin-bottom: 18px !important;
    border: 1px solid #e3e7ef !important;
    padding: 16px !important;
    color: #212529 !important;
}
.suggestion-btn {
    margin: 0 12px 18px 0 !important;
    background: #eaf4ff !important;
    color: #2563eb !important;
    border-radius: 10px !important;
    font-weight: 500 !important;
    font-size: 1rem !important;
    border: 1.5px solid #bcdfff !important;
    transition: background 0.2s, color 0.2s, border-color 0.2s;
}
.suggestion-btn:hover {
    background: #d2eaff !important;
    color: #174ea6 !important;
    border-color: #2563eb !important;
}
#medical-title {
    text-align: center;
    color: #2563eb;
    font-size: 2.3rem;
    font-weight: 700;
    margin-bottom: 18px;
    letter-spacing: 1px;
}
#user-input {
    border-radius: 10px !important;
    border: 1.5px solid #bcdfff !important;
    padding: 10px 16px !important;
    font-size: 1.1rem !important;
    background: #fff !important;
    color: #212529 !important;
}
#user-input::placeholder {
    color: #8ca0b3 !important;
    opacity: 1 !important;
}
#submit-btn {
    background: linear-gradient(90deg, #2563eb 0%, #38bdf8 100%) !important;
    color: #fff !important;
    border-radius: 10px !important;
    font-size: 1.3rem !important;
    min-width: 56px;
    min-height: 44px;
    border: none !important;
    box-shadow: 0 2px 8px rgba(60, 80, 120, 0.07);
    transition: background 0.2s;
}
#submit-btn:hover {
    background: linear-gradient(90deg, #174ea6 0%, #0ea5e9 100%) !important;
    color: #fff !important;
}
"""

with gr.Blocks(theme=gr.themes.Monochrome(), css=css) as demo:
    with gr.Column(elem_id="main-card"):
        gr.Markdown("""
        <div id='medical-title'>Medical QA Bot</div>
        """)
        chatbot = gr.Chatbot(label="", elem_id="chatbot", type="messages", height=300)
        state = gr.State([])


        with gr.Row():
            specialty_options = [
                "Rheumatology", "Psychiatry", "Pulmonology & Respiratory Medicine", "Nephrology", "Public Health & Epidemiology",
                "Medical Research & Methodology", "Pharmacy & Pharmacology", "Hematology", "Oncology", "Medical Ethics & Law",
                "Medical Technology & Informatics", "Infectious Disease", "Basic Medical Sciences", "Allergology", "Geriatrics",
                "Cardiology", "Gastroenterology & Hepatology", "General Surgery", "General Pediatrics", "Endocrinology & Metabolism",
                "Vascular Surgery", "Radiology & Imaging", "Obstetrics & Gynecology", "Orthopedic Surgery", "Neurology",
                "Family Medicine & Primary Care", "Psychology & Behavioral Health", "Otorhinolaryngology (ENT)", "General Internal Medicine",
                "Anesthesiology", "Physical & Rehabilitation Medicine", "Medical Education", "Healthcare Administration & Management",
                "Non-Medical Sciences & Disciplines", "Dermatology", "Critical Care & Intensive Care", "Urology", "Complementary & Alternative Medicine",
                "Cardiothoracic Surgery", "Neurosurgery", "Pediatric Subspecialties", "Occupational & Environmental Health", "Ophthalmology",
                "Emergency Medicine", "Dental & Oral Medicine", "Biomedical Engineering", "Pathology & Laboratory Medicine", "Transplant Surgery",
                "Preventive Medicine", "Genetics", "Nursing", "Allied Health Professions", "Plastic & Reconstructive Surgery", "Others",
                "Toxicology", "General Medicine"
            ]
            year_options = [str(y) for y in range(2021, 1792, -1)]
            specialty_dropdown = gr.Dropdown(choices=specialty_options, value="General Medicine", label="Specialty", scale=4, elem_id="specialty-dropdown")
            year_dropdown = gr.Dropdown(choices=year_options, value="2021", label="Year", scale=2, elem_id="year-dropdown")
        with gr.Row():
            user_input = gr.Textbox(placeholder="Type a medical question...", show_label=False, lines=1, scale=8, elem_id="user-input")
            submit_btn = gr.Button(value="➤", scale=1, elem_id="submit-btn")

        with gr.Row():
            suggestion_buttons = []
            for i, q in enumerate(sample_questions):
                btn = gr.Button(q, elem_id=f"suggestion-{i}", elem_classes=["suggestion-btn"])
                suggestion_buttons.append(btn)

    def suggestion_click(q, history, year, specialty):
        return qa_bot(q, history, year, specialty)

    submit_btn.click(qa_bot, inputs=[user_input, state, year_dropdown, specialty_dropdown], outputs=[chatbot, state])
    user_input.submit(qa_bot, inputs=[user_input, state, year_dropdown, specialty_dropdown], outputs=[chatbot, state])
    for btn, q in zip(suggestion_buttons, sample_questions):
        btn.click(suggestion_click, inputs=[gr.State(q), state, year_dropdown, specialty_dropdown], outputs=[chatbot, state])

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