# app.py
import streamlit as st
from transformers import pipeline

# -------------------------
# Load model
# -------------------------
@st.cache_resource
def load_model():
    # Using FLAN-T5 (free, light, works well on Spaces CPU)
    return pipeline(
        "text2text-generation",
        model="google/flan-t5-base"
    )

model = load_model()

# -------------------------
# Streamlit UI
# -------------------------
st.set_page_config(page_title="Prescription Chatbot", page_icon="💊")

st.title("💊 Prescription Chatbot (Demo)")
st.caption("Academic demo only — not medical advice")

symptoms = st.text_area("Enter your symptoms:", "I have flu, body pain and runny nose")

if st.button("Get Prescription"):
    if symptoms.strip():
        prompt = f"""
You are a helpful medical-assistant for an academic demo only. 
The patient reports: {symptoms}.

Task: Provide a short structured response with:
1. Most likely condition(s)
2. Suggested non-prescription remedies
3. Precautions and red-flags
4. Disclaimer: This is a demo only — not medical advice.
"""

        output = model(prompt, max_length=256, do_sample=False)
        reply = output[0]["generated_text"].strip()

        st.subheader("Suggested Result (Demo)")
        st.write(reply)
    else:
        st.warning("Please enter symptoms first.")