import streamlit as st
from google.cloud import aiplatform
#import google.generativeai as genai
import os
import logging
import whisper
from gtts import gTTS
import tempfile
from pydub import AudioSegment
from groq import Groq, GroqError


from google.cloud import aiplatform

# Initialize the client
def init_palm(api_key):
    aiplatform.init(api_key=api_key)

# Function to generate responses
def generate_palm_response(prompt):
    response = aiplatform.Model.predict(
        model_name="gemini-1.5-flash",
        instances=[{"prompt": prompt}]
    )
    return response.predictions[0]['content']

# Call the model
api_key = "your-google-api-key"
init_palm(api_key)
response = generate_palm_response("Heart health query...")
print(response)



# Configure API keys securely
#GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
#genai.configure(api_key=GOOGLE_API_KEY)

groq_api_key = os.getenv('GROQ_API_KEY')
if not groq_api_key:
    raise ValueError("GROQ_API_KEY is not set.")

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

# Initialize Groq API client
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
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

# Streamlit page configuration
st.set_page_config(
    page_title="Heart Health & Audio Processing App",
    page_icon="🫀",
    layout="centered",
    initial_sidebar_state="collapsed",
)

# Initialize the Generative Model for heart health 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 input
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

# Main application layout
def main():
    st.title("Heart Health & Audio Processing App 🫀🎙️")

    # Two tabs: one for the chatbot and one for audio processing
    tab1, tab2 = st.tabs(["Heart Health Chatbot", "Audio Processing"])

    # Tab 1: Heart Health Chatbot
    with tab1:
        st.header("Chat with Heart Health Specialist Dr. Assad Siddiqui")
        
        if "history" not in st.session_state:
            st.session_state.history = []

        user_input = st.text_input("Ask about heart health:", placeholder="Type here...")

        if st.button("Send") and user_input:
            bot_response = get_chatbot_response(user_input)
            st.session_state.history.append({"role": "user", "content": user_input})
            st.session_state.history.append({"role": "bot", "content": bot_response})
        
        for chat in st.session_state.history:
            if chat["role"] == "user":
                st.write(f"**You:** {chat['content']}")
            else:
                st.write(f"**Bot:** {chat['content']}")

    # Tab 2: Audio Processing
    with tab2:
        st.header("Audio Processing with Whisper and Groq")

        uploaded_audio = st.file_uploader("Upload an audio file for transcription and response", type=["mp3", "wav", "ogg"])
        
        if uploaded_audio:
            with st.spinner("Processing audio..."):
                response_text, audio_file_path = process_audio(uploaded_audio)
            
            if response_text:
                st.write(f"**Response:** {response_text}")
                st.audio(audio_file_path)

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