import gradio as gr
from transformers import pipeline

# Load model (lightweight, works on HF Spaces CPU)
pipe = pipeline(
    "text-generation",
    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0"
)

def analyze(symptoms):

    prompt = f"""
You are MediGuide AI, a medical assistant for basic guidance only.

Strict format:

1. Possible Conditions
2. Risk Level (Low / Medium / High)
3. Basic Guidance
4. Emergency Advice

Rules:
- Do NOT give prescription drugs
- Do NOT give final diagnosis
- Keep it simple and clear

Symptoms:
{symptoms}
"""

    result = pipe(
        prompt,
        max_new_tokens=250,
        temperature=0.7,
        do_sample=True
    )

    return result[0]["generated_text"]


with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown("""
    # 🩺 MediGuide AI
    ### Offline Medical Reasoning Assistant
    For underserved communities 🌍
    """)

    inp = gr.Textbox(
        label="Describe your symptoms",
        lines=5,
        placeholder="e.g. fever, headache, body pain"
    )

    btn = gr.Button("Analyze")

    out = gr.Textbox(
        label="Medical Analysis",
        lines=15
    )

    btn.click(analyze, inp, out)

demo.launch()