import gradio as gr
from transformers import pipeline

# Load model
generator = pipeline("text2text-generation", model="google/flan-t5-base")

# Knowledge base
disaster_knowledge = {
    "earthquake": """
🌍 Earthquake Safety:
1. Drop, cover, and hold on immediately.
2. Stay indoors until shaking stops completely.
3. Stay away from windows, shelves, and heavy objects.
4. If outside, move to an open area away from buildings, trees, and power lines.
5. After shaking stops, check for injuries and help others.
6. Avoid using elevators.
📞 Helpline: 1078 (National Disaster Helpline, India)
""",
    "flood": """
🌊 Flood Safety:
1. Move to higher ground immediately.
2. Avoid walking or driving through floodwaters – even 6 inches of water can knock you down.
3. Disconnect electrical appliances.
4. Keep emergency supplies (water, food, torch, first aid kit).
5. Follow evacuation orders from authorities.
📞 Helpline: 1070 (State Emergency Operations Center, India)
""",
    "fire": """
🔥 Fire Safety:
1. Stay low to avoid smoke inhalation.
2. Use the stairs, not elevators.
3. Call emergency services immediately.
4. Stop, drop, and roll if your clothes catch fire.
5. Exit quickly and leave belongings behind.
📞 Fire Helpline: 101 (India)
""",
    "cyclone": """
🌪️ Cyclone Safety:
1. Stay indoors and away from windows.
2. Keep emergency kits ready (water, food, medicines, torch).
3. Disconnect electrical appliances.
4. Follow official weather updates and evacuation orders.
📞 Helpline: 108 (Emergency Medical Services, India), 1098 (Childline)
"""
}

# Function to pick the right context
def get_context(user_input: str):
    text = user_input.lower()
    if "earthquake" in text:
        return disaster_knowledge["earthquake"]
    elif "flood" in text:
        return disaster_knowledge["flood"]
    elif "fire" in text:
        return disaster_knowledge["fire"]
    elif "cyclone" in text or "hurricane" in text or "typhoon" in text:
        return disaster_knowledge["cyclone"]
    else:
        return "\n".join(disaster_knowledge.values())  # default: all

# Chat function
def disaster_chat(user_input, history):
    context = get_context(user_input)
    prompt = f"""{context}

User question: {user_input}

Give a detailed, step-by-step safety guide with helpline numbers.
"""
    response = generator(prompt, max_length=400, temperature=0.7)[0]["generated_text"]
    return response

# Welcome message
welcome_msg = (
    "👋 Hello! I am your **Disaster Safety Assistant**.\n\n"
    "Ask me questions like:\n"
    "• What should I do during an earthquake?\n"
    "• How can I stay safe in a flood?\n"
    "• What should I do if there is a fire?\n"
    "• How to prepare for a cyclone?\n\n"
    "I will always give you **step-by-step safety precautions with helpline numbers**. 🌟"
)

# Gradio Chat UI
demo = gr.ChatInterface(
    fn=disaster_chat,
    title="🌍 Disaster Guidance AI Chatbot",
    description="Get detailed step-by-step safety guidance (earthquake, flood, fire, cyclone) including helpline numbers.",
    chatbot=gr.Chatbot(value=[("assistant", welcome_msg)])
)

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