import os
import streamlit as st
from dotenv import load_dotenv
from groq import Groq

# Load environment variables
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# Initialize Groq client
client = Groq(api_key=GROQ_API_KEY)

# Symptom checker for specific diseases
def symptom_checker(symptoms):
    symptom_mapping = {
        "fever": "You might have a viral or bacterial infection. Drink fluids, rest, and take paracetamol if needed.",
        "flu": "Symptoms of flu include fever, body aches, and fatigue. Consider antiviral medication if severe.",
        "cough": "Dry cough may indicate allergies, while a wet cough could be a respiratory infection. Drink warm fluids.",
        "malaria": "Malaria symptoms include high fever, chills, and sweating. Seek medical attention immediately.",
        "typhoid": "Symptoms include prolonged fever, weakness, and digestive issues. Antibiotics are required; consult a doctor."
    }
    return symptom_mapping.get(symptoms.lower(), "Sorry, I do not have knowledge of it. Please ask only health-related questions.")

# Function to query LLaMA model
def ask_health_assistant(question):
    try:
        response = client.chat.completions.create(
            messages=[{"role": "user", "content": question}],
            model="llama-3.3-70b-versatile"
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: {str(e)}"

# Streamlit UI
st.title("🩺 Health Assistant Chatbot")
st.write("Ask me anything related to health!")

# User input
user_input = st.text_input("Enter your health-related question:")

# Button to process user input
if st.button("Ask"):
    if user_input.lower() in ["fever", "flu", "cough", "malaria", "typhoid"]:
        response = symptom_checker(user_input)
    else:
        response = ask_health_assistant(user_input)
    
    st.write("🤖 **Bot:**", response)