import streamlit as st
import os
import speech_recognition as sr
import pyttsx3
from openai import OpenAI
import time
import threading
import random

css = '''
<style>
body {
    font-family: Arial, sans-serif;
    background-color: #f0f2f6;
}
.chat-message {
    padding: 1.5rem;
    border-radius: 0.5rem;
    margin-bottom: 1rem;
    display: flex;
    box-shadow: 0 2px 5px rgba(0,0,0,0.1);
}
.chat-message.bot {
    background-color: #475063;
}
.chat-message.user {
    background-color: #2b313e;
}
.chat-message .message {
    width: 90%;
    padding: 0 1.5rem;
    color: #fff;
    line-height: 1.6;
}
.stButton>button {
    width: 100%;
    border-radius: 20px;
    font-weight: bold;
}
</style>
'''

bot_template = '''
<div class="chat-message bot">
    <div class="message">{{MSG}}</div>
</div>
'''

user_template = '''
<div class="chat-message user">
    <div class="message">{{MSG}}</div>
</div>
'''

def load_openai(api_key):
    client = OpenAI(api_key=api_key)
    return client

def validate_api_key(api_key):
    try:
        client = load_openai(api_key)
        client.models.list()
        return True
    except:
        return False

file_name = "output.txt"
def update_history(data, mode):
    with open(file_name, mode, encoding='utf-8') as file:
        file.write(data + "\n\n")
        file.close()

def download_history():
    try:
        if os.path.exists(file_name):
            with open(file_name) as file:
                file_data = file.read()
                file.close()

        st.download_button(
            label="Download data",
            data=file_data,
            file_name='history.txt',
            mime='text')
    except:
        st.write("No history to download")

def clear_history():
    st.session_state.history = []
    
    if os.path.exists(file_name):
        os.remove(file_name)
        st.write("Cleared history")
    else:
        st.write("No history to clear")    

def text_to_speech(text):
    def speak():
        engine = pyttsx3.init()
        engine.setProperty('rate', 150)
        engine.setProperty('volume', 1)
        engine.say(text)
        engine.runAndWait()

    st.session_state.tts_thread = threading.Thread(target=speak)
    st.session_state.tts_thread.start()

def stop_text_to_speech():
    if hasattr(st.session_state, 'tts_thread') and st.session_state.tts_thread.is_alive():
        engine = pyttsx3.init()
        engine.stop()
        st.session_state.tts_thread.join()

def get_random_question():
    questions = [
        "How about we set a fun fitness goal for this week? What would you like to achieve?",
        "Have you tried any new healthy recipes lately? I'd love to hear about them!",
        "What's your favorite way to stay active when you're short on time?",
        "How's your water intake been today? Remember, staying hydrated is key!",
        "If you could instantly master any sport, what would it be and why?",
        "What's the most fun you've ever had during a workout?",
        "Have you ever considered trying a new fitness class? Any that interest you?",
        "What's your go-to healthy snack when you need a quick energy boost?",
        "If you could work out anywhere in the world, where would you choose?",
        "What's one small healthy habit you'd like to incorporate into your daily routine?"
    ]
    return random.choice(questions)

def stream_response(prompt):
    placeholder = st.empty()
    full_response = ''
    
    messages = [{"role": "system", "content": system_prompt}]
    for i in range(0, len(st.session_state.history), 2):
        if i+1 < len(st.session_state.history):
            user_message = st.session_state.history[i]
            assistant_message = st.session_state.history[i+1]
            messages.append({"role": "user", "content": user_message})
            messages.append({"role": "assistant", "content": assistant_message})
    messages.append({"role": "user", "content": prompt})

    for chunk in st.session_state.client.chat.completions.create(
        model=st.session_state.model,
        messages=messages,
        stream=True
    ):
        if chunk.choices[0].delta.content is not None:
            full_response += chunk.choices[0].delta.content
            placeholder.markdown(bot_template.replace("{{MSG}}", full_response), unsafe_allow_html=True)
    
    update_history(prompt, 'a')
    update_history(full_response, 'a')
    
    st.session_state.history.append(prompt)
    st.session_state.history.append(full_response)
    placeholder.markdown('')
    if len(st.session_state.history) > 10:
        st.session_state.history = st.session_state.history[-10:]
    
    st.session_state.response = full_response
    text_to_speech(full_response)

    # Add a random question if the response is short
    if len(full_response.split()) < 20:
        random_question = get_random_question()
        st.session_state.history.append(random_question)
        st.markdown(bot_template.replace("{{MSG}}", random_question), unsafe_allow_html=True)
        text_to_speech(random_question)

def display_history():
    for i in range(0, len(st.session_state.history), 2):
        user_msg = st.session_state.history[i]
        st.markdown(user_template.replace("{{MSG}}", user_msg), unsafe_allow_html=True)
        if i+1 < len(st.session_state.history):
            bot_msg = st.session_state.history[i+1]
            st.markdown(bot_template.replace("{{MSG}}", bot_msg), unsafe_allow_html=True)

def recognize_speech():
    r = sr.Recognizer()
    with sr.Microphone() as source:
        st.write("Listening...")
        audio = r.listen(source)
        st.write("Processing...")

    try:
        text = r.recognize_google(audio)
        return text
    except sr.UnknownValueError:
        return "Sorry, I couldn't understand that."
    except sr.RequestError:
        return "Sorry, there was an error processing your speech."

st.set_page_config(page_title="Your Personal Fitness Buddy", page_icon=":muscle:")
st.markdown(css, unsafe_allow_html=True)

st.header("Meet ZenFit, Your Personal Fitness Buddy :muscle:")

if 'api_key_validated' not in st.session_state:
    st.session_state.api_key_validated = False

if not st.session_state.api_key_validated:
    api_key = st.text_input("Enter your OpenAI API key:", type="password")
    if st.button("Validate API Key"):
        if validate_api_key(api_key):
            st.session_state.api_key_validated = True
            st.session_state.client = load_openai(api_key)
            st.session_state.model = "gpt-3.5-turbo"
            st.session_state.history = []
            st.success("API key validated successfully!")
            st.experimental_rerun()
        else:
            st.error("Invalid API key. Please try again.")
else:
    system_prompt = """
    You are ZenFit, a friendly and enthusiastic personal fitness buddy. Your goal is to provide concise, relevant responses to the user's questions or comments about fitness and health. Keep your answers brief and focused on the specific topic at hand. When appropriate, offer encouragement and practical tips. If the conversation lulls, ask engaging questions about the user's fitness journey or healthy lifestyle choices.
    """

    if 'initial_message' not in st.session_state:
        initial_message = """
        Hey there! I'm ZenFit, your personal fitness buddy. 🏋️‍♂️💪 I'm here to help you on your journey to a healthier, more active lifestyle. Whether you're a fitness newbie or a seasoned pro, I've got your back!

        Feel free to ask me anything about workouts, nutrition, or general wellness. I'm all ears and ready to chat. So, what's on your mind today? How can I help you take a step towards your fitness goals?
        """
        st.session_state.history.append(initial_message)
        st.session_state.initial_message = True
        text_to_speech(initial_message)

    display_history()

    col1, col2 = st.columns(2)
    with col1:
        if st.button("Push to Talk"):
            try:
                user_input = recognize_speech()
                if user_input and user_input not in ["Sorry, I couldn't understand that.", "Sorry, there was an error processing your speech."]:
                    st.write(f"You said: {user_input}")
                    stream_response(user_input)
                else:
                    st.write("No valid input detected. Please try again.")
            except Exception as e:
                st.write(f"An error occurred: {str(e)}. Please try again.")
            st.experimental_rerun()
    with col2:
        if st.button("Stop TTS"):
            stop_text_to_speech()

    col3, col4 = st.columns(2)
    with col3:
        if st.button("Clear History", key="clear"):
            clear_history()
            st.session_state.initial_message = False
            st.experimental_rerun()
    with col4:
        if st.button("Download History", key="download"):
            download_history()