import streamlit as st
import google.generativeai as genai
from langchain.embeddings import GooglePalmEmbeddings
from langchain.vectorstores import FAISS
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains import RetrievalQA
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import os
import time

api_key = os.getenv("GOOGLE_API_KEY")

# Configure Gemini API
genai.configure()
model = genai.GenerativeModel("models/gemini-pro") # Use a more stable model
embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=api_key)

def calculate_bmi(weight, height):
    """Calculates BMI."""
    try:
        height_m = height / 100  # Convert cm to meters
        bmi = weight / (height_m ** 2)
        return bmi
    except ZeroDivisionError:
        return None

def get_bmi_category(bmi):
    """Categorizes BMI."""
    if bmi is None:
        return "Invalid input"
    elif bmi < 18.5:
        return "Underweight"
    elif 18.5 <= bmi < 25:
        return "Normal weight"
    elif 25 <= bmi < 30:
        return "Overweight"
    else:
        return "Obese"

def generate_workout_plan(weight, height, gender):
    """Generates workout plan using Gemini."""
    bmi = calculate_bmi(weight, height)
    if bmi is None:
        return None
    bmi_category = get_bmi_category(bmi)
    with st.spinner("Generating workout plan..."):  # Spinner during generation
        try:
            prompt = f"""
            I am a {gender} with a BMI indicating I am {bmi_category}. 
            Generate a full week workout plan suitable for me. 
            Include specific exercises, sets, reps, and rest times.
            Also give 3 fitness tips.
            """
            response = model.generate_content(prompt)
            time.sleep(1) # artificial delay to show spinner
            return response.text
        except Exception as e:
            st.error(f"Error generating workout plan: {e}")
            return None  # Return None in case of error

def create_or_load_vectorstore(notes):
    """Creates or loads a vectorstore from user notes."""
    if notes:
        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
        docs = text_splitter.create_documents([notes])
        vectorstore = FAISS.from_documents(docs, embeddings)
        return vectorstore
    return None

def query_vectorstore(vectorstore, query):
    """Queries the vectorstore using Langchain's RetrievalQA chain."""
    with st.spinner("Searching your notes..."):  # Spinner during search
        if vectorstore:
            qa = RetrievalQA.from_chain_type(llm=model, chain_type="stuff", retriever=vectorstore.as_retriever())
            try:
                result = qa.run(query)
                time.sleep(1) # artificial delay to show spinner
                return result
            except Exception as e:
                st.error(f"Error querying notes: {e}")
                return None  # Return None in case of error
        return "No notes provided."

st.title("💪GYM Fitness Chatbot🤖")

weight = st.number_input("Enter your weight (in kg)", min_value=0)
height = st.number_input("Enter your height (in cm)", min_value=0)
gender = st.selectbox("Select your gender", ["Male", "Female", "Other"])

notes = st.text_area("Enter any additional notes or details (e.g., injuries, preferences):", height=150)

if st.button("Calculate BMI and Get Workout Plan"):
    if weight and height and gender:
        with st.spinner("Calculating BMI..."):
            time.sleep(1) # artificial delay to show spinner
        workout_plan = generate_workout_plan(weight, height, gender)  # Pass weight, height, and gender
        if workout_plan:
            st.write("## Your Personalized Workout Plan:")
            st.write(workout_plan)

            # Notes and RAG Section
            vectorstore = create_or_load_vectorstore(notes)
            user_query = st.text_input("Ask a question about your notes:")
            if st.button("Get answer from Notes"):
                if user_query:
                    answer = query_vectorstore(vectorstore, user_query)
                    if answer:
                        st.write("## Answer from your notes:")
                        st.write(answer)
                    else:
                        st.warning("Could not retrieve answer from notes")
                else:
                    st.warning("Please enter a query to search your notes.")
    else:
        st.warning("Please enter your weight, height, and gender.")