import os
import google.generativeai as genai
import gradio as gr
from PIL import Image

# Load API key securely from environment variables
API_KEY = os.getenv("AIzaSyDeQ5VIcPbgzSG6QfFvn4tjP4G-V-frpwM")  # Set this in your environment
if not API_KEY:
    raise ValueError("API Key is missing! Set GEMINI_API_KEY as an environment variable.")

# Configure Gemini API
try:
    genai.configure(api_key=API_KEY)
    model = genai.GenerativeModel("gemini-1.5-flash")  # Corrected model name
    chat = model.start_chat(history=[])
except Exception as e:
    print(f"Error initializing Gemini model: {e}")
    exit()

# Function to process chatbot interaction
def chatbot_interface(text_input, image_input=None, audio_input=None):
    try:
        if not text_input:
            return "Please enter a message."

        content = [{"text": text_input}]

        # Handle image input
        if image_input:
            try:
                image_input.save("temp_image.png")  # Save and load as a file
                with open("temp_image.png", "rb") as img_file:
                    image_data = img_file.read()
                content.append({"image": image_data})
            except Exception as e:
                return f"Error processing image: {e}"

        # Handle audio input
        if audio_input and os.path.exists(audio_input):
            try:
                with open(audio_input, "rb") as audio_file:
                    audio_data = audio_file.read()
                content.append({"audio": audio_data})
            except Exception as e:
                return f"Error processing audio: {e}"
        elif audio_input:
            return "Audio file not found."

        # Send message to Gemini model
        response = chat.send_message(content)
        return response.text if response else "No response from the model."

    except Exception as e:
        return f"An error occurred: {e}"

# Create Gradio UI
interface = gr.Interface(
    fn=chatbot_interface,
    inputs=[
        gr.Textbox(label="Your Message", placeholder="Type your question here..."),
        gr.Image(label="Upload an Image", type="pil"),
        gr.Audio(label="Upload or Record Audio", type="filepath")
    ],
    outputs=gr.Textbox(label="Bot Response"),
    title="FINANCE Chatbot",
    description="Chat using text, images, and audio! Ensure your API key is valid.",
    allow_flagging="never"
)

# Launch the Gradio interface
if __name__ == "__main__":
    try:
        interface.launch(debug=True, share=True)
    except Exception as e:
        print(f"Error launching interface: {e}")