import gradio as gr
import os
from groq import Groq

# Initialize Groq client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

def chat_with_medbrain(message, history):
    """
    MedBRAIN AI - Medical Assistant Chatbot
    Uses Groq's Llama model for medical information assistance
    """
    try:
        # Build conversation history for context
        messages = [
            {
                "role": "system",
                "content": """You are MedBRAIN AI, a specialized medical information assistant chatbot.
                
                STRICT SCOPE: You ONLY answer questions related to:
                - Medical conditions, diseases, and disorders
                - Symptoms, diagnosis, and medical treatments
                - Medications, drugs, and pharmaceuticals
                - Human anatomy, physiology, and biology
                - Healthcare, wellness, and nutrition
                - Medical procedures and tests
                - Mental health and psychology
                - Public health and epidemiology
                
                IMPORTANT RESTRICTIONS:
                - If a question is NOT medical/health-related, politely decline and redirect
                - Do NOT answer questions about: technology, programming, mathematics, general knowledge, entertainment, sports, politics, or any non-medical topics
                - If the question IS medical/health-related, answer directly without the disclaimer                
                IMPORTANT DISCLAIMERS:
                - You provide educational information only, not medical advice
                - Always recommend consulting healthcare professionals for diagnosis and treatment
                - You cannot diagnose conditions or prescribe treatments
                - In emergencies, advise calling emergency services immediately
                
                Provide clear, accurate medical information while being empathetic and professional."""
            }
        ]
        
        # Add conversation history
        for human, assistant in history:
            messages.append({"role": "user", "content": human})
            messages.append({"role": "assistant", "content": assistant})
        
        # Add current message
        messages.append({"role": "user", "content": message})
        
        # Get response from Groq
        response = client.chat.completions.create(
            model="llama-3.1-8b-instant",            messages=messages,
            temperature=0.7,
            max_tokens=1024,
        )
        
        return response.choices[0].message.content
    
    except Exception as e:
        return f"I apologize, but I encountered an error: {str(e)}. Please try again or contact support if the issue persists."

# Create Gradio interface
demo = gr.ChatInterface(
    fn=chat_with_medbrain,
    title="🧠 MedBRAIN AI - Medical Assistant",
    description="""A powerful medical assistant chatbot powered by advanced AI. 
    Ask questions about medical conditions, symptoms, treatments, and general health information.""",
    examples=[
        "What are the symptoms of diabetes?",
        "How can I manage high blood pressure?",
        "What are the common causes of headaches?",
        "Explain the difference between bacteria and viruses",
        "What are the benefits of regular exercise?"
    ],
    theme=gr.themes.Soft(),
)

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