import os
import logging
import streamlit as st
import google.generativeai
from streamlit_chat import message
import whisper
from gtts import gTTS
import tempfile
from pydub import AudioSegment
from groq import Groq, GroqError

# Securely configure API keys
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# Configure Google Generative AI API
genai.configure(api_key=GOOGLE_API_KEY)

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize Groq Client
if not GROQ_API_KEY:
    raise ValueError("GROQ_API_KEY is not set.")
try:
    groq_client = Groq(api_key=GROQ_API_KEY)
    logger.info("Groq API key is set and client is initialized.")
except GroqError as e:
    logger.error(f"Failed to initialize Groq client: {e}")
    raise

# Load Whisper model for audio transcription
try:
    whisper_model = whisper.load_model("base")
    logger.info("Whisper model loaded successfully.")
except Exception as e:
    logger.error(f"Failed to load Whisper model: {e}")
    raise

# Initialize Google Generative Model for chatbot
model = genai.GenerativeModel(
    'gemini-1.5-flash',
    system_instruction=(
        "Persona: You are Dr. Assad Siddiqui, a heart specialist. Only provide information related to heart health, symptoms, and advice. "
        "Ask users about their heart-related symptoms and provide consultation and guidance based on their input. "
        "Always provide brief answers. If the inquiry is not related to heart health, politely say that you can only provide heart-related information. "
        "Responses should be in Urdu written in English and in English."
    )
)

# Function to get chatbot response
def get_chatbot_response(user_input):
    response = model.generate_content(user_input)
    return response.text.strip()

# Function to process audio using Whisper and Groq API
def process_audio(audio_file):
    try:
        result = whisper_model.transcribe(audio_file)
        user_text = result['text']
        logger.info(f"Transcription successful: {user_text}")
    except Exception as e:
        logger.error(f"Error in transcribing audio: {e}")
        return "Error in transcribing audio.", None

    try:
        chat_completion = groq_client.chat.completions.create(
            messages=[{"role": "user", "content": user_text}],
            model="llama3-8b-8192",
        )
        response_text = chat_completion.choices[0].message.content
        logger.info(f"Received response from Groq API: {response_text}")
    except GroqError as e:
        logger.error(f"Error in generating response with Groq API: {e}")
        return "Error in generating response with Groq API.", None

    try:
        tts = gTTS(text=response_text, lang='en')
        audio_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp3')
        tts.save(audio_file.name)
        logger.info("Text-to-speech conversion successful.")
    except Exception as e:
        logger.error(f"Error in text-to-speech conversion: {e}")
        return "Error in text-to-speech conversion.", None

    return response_text, audio_file.name

# Streamlit page configuration
st.set_page_config(page_title="Heart Health Chatbot", page_icon="👨‍⚕️", layout="centered")

# Background and header
st.markdown("""
    <style>
        .stApp { background-image: url('https://cdn.wallpapersafari.com/29/34/8Ak1Sf.png'); background-size: cover; }
        .chat-bubble { background-color: #128c7E; color: white; padding: 10px; border-radius: 10px; max-width: 70%; }
        .user-bubble { background-color: #075e54; color: white; padding: 10px; border-radius: 10px; max-width: 70%; }
        img.avatar { width: 50px; height: 50px; border-radius: 50%; }
    </style>
    <div style="padding:10px;text-align:center;color:white;">
        <h1>Heart Health Chatbot 🫀</h1>
        <p>Ask me anything about heart diseases!</p>
    </div>
""", unsafe_allow_html=True)

# Initialize session state for chat history
if "history" not in st.session_state:
    st.session_state.history = []

user_avatar_url = "https://img.freepik.com/free-photo/sad-cartoon-anatomical-heart_23-2149767987.jpg"
bot_avatar_url = "https://img.freepik.com/premium-photo/3d-render-man-doctor-avatar-round-sticker-with-cartoon-character-face-user-id-thumbnail.jpg"

# Function to display chat history
def display_chat_history():
    for chat in st.session_state.history:
        if chat["role"] == "user":
            st.markdown(f"""
                <div style="display: flex; justify-content: flex-end; margin-bottom: 10px;">
                    <div class="user-bubble"><p><b>You:</b> {chat['content']}</p></div>
                    <img src="{user_avatar_url}" class="avatar"/>
                </div>
            """, unsafe_allow_html=True)
        else:
            st.markdown(f"""
                <div style="display: flex; margin-bottom: 10px;">
                    <img src="{bot_avatar_url}" class="avatar"/>
                    <div class="chat-bubble"><p><b>Bot:</b> {chat['content']}</p></div>
                </div>
            """, unsafe_allow_html=True)

# Main application layout
def main():
    display_chat_history()
    with st.container():
        with st.form(key="user_input_form", clear_on_submit=True):
            user_input = st.text_input("Type your message...", placeholder="Ask about heart health...", max_chars=500)
            submit_button = st.form_submit_button("Send")

            if submit_button and user_input.strip():
                with st.spinner("Thinking..."):
                    bot_response = get_chatbot_response(user_input)

                # Update chat history
                st.session_state.history.append({"role": "user", "content": user_input})
                st.session_state.history.append({"role": "bot", "content": bot_response})

                display_chat_history()

# Footer
st.markdown("""
    <p style="text-align:center; color:white; margin-top:50px;">
        Check out the <a href="https://live-appointment-chatbot20.zapier.app/" target="_blank" style="color:#34c759;">Live Appointment</a>.
    </p>
""", unsafe_allow_html=True)

# Run the app
if __name__ == "__main__":
    main()