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

# Load environment variables from .env file
load_dotenv()
API_KEY = ("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")

# Check if API key is set
if not API_KEY:
    st.error("GROQ_API_KEY is missing. Set it as an environment variable or in a .env file.")
    st.stop()

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

# System prompt for the Gynecologist Chatbot
SYSTEM_PROMPT = """
You are an AI-powered Gynecologist, designed to provide accurate, science-backed information about women's health, 
pregnancy, menstrual cycles, fertility, and reproductive health.
You assist users by answering questions about gynecological health, prenatal care, contraception, and hormonal balance.
You do NOT provide medical diagnoses or prescriptions.
If a user asks about serious medical conditions, always recommend consulting a healthcare professional.
Maintain a professional, supportive, and non-judgmental tone.
"""

# Streamlit UI
st.set_page_config(page_title="👩‍⚕️ Gynecologist Chatbot", page_icon="💖")

st.title("👩‍⚕️ AI Gynecologist Chatbot")
st.write("Ask me anything about women's health, pregnancy, periods, or reproductive health!")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": SYSTEM_PROMPT}]

# Function to check if the question is gynecology-related
def is_gynecology_related(question):
    keywords = [
        "pregnancy", "menstrual", "fertility", "ovulation", "contraception", "hormones", "reproductive", 
        "periods", "gynecologist", "women's health", "vaginal", "uterus", "cervix", "PCOS", "fibroids",
        "ovaries", "menopause", "endometriosis", "prenatal", "postpartum", "STI", "sexual health"
    ]
    return any(re.search(rf"\b{kw}\b", question, re.IGNORECASE) for kw in keywords)

# Display chat history
for message in st.session_state.messages[1:]:  # Skip system prompt in UI
    with st.chat_message(message["role"]):
        st.write(message["content"])

# User input
user_input = st.chat_input("Type your question here...")

if user_input:
    if is_gynecology_related(user_input):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": user_input})
        
        # Send request to Groq API using LLaMA 3.3 70B Versatile model
        try:
            response = client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=st.session_state.messages
            )

            # Extract AI response
            ai_response = response.choices[0].message.content

            with st.chat_message("user"):
                st.markdown(f"**You:** {user_input}")

            # Display AI response
            with st.chat_message("assistant"):
                st.write(ai_response)

            # Add AI response to chat history
            st.session_state.messages.append({"role": "assistant", "content": ai_response})

        except Exception as e:
            st.error(f"Error: {e}")
    else:
        st.warning("Please ask a question related to gynecology, women's health, pregnancy, or reproductive health.")