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

# Initialize Gemini API
try:
    GOOGLE_API_KEY = os.environ["GOOGLE_API_KEY"]
    genai.configure(api_key=GOOGLE_API_KEY)
    model = genai.GenerativeModel("gemini-pro-vision")
except KeyError:
    print("Error: GEMINI_API_KEY not found in environment variables. Please set it in Hugging Face Space secrets.")
    exit()
except Exception as e:
    print(f"Error initializing Gemini API: {e}")
    exit()

def generate_response(text_input, file_input):
    """Generates a response using Gemini's multiprompt feature with various file types."""
    try:
        contents = []
        if text_input:
            contents.append(text_input)

        if file_input:
            file_path = file_input
            mime_type, _ = mimetypes.guess_type(file_path)
            if mime_type is None:
                return "Unsupported file type. Please upload a valid image, video, audio, or PDF."
            
            try:
                with open(file_path, "rb") as f:
                    file_bytes = f.read()
            except Exception as e:
                return f"Error reading file: {e}"
            
            if mime_type.startswith("image"):
                try:
                    # Check if it's a valid image
                    Image.open(io.BytesIO(file_bytes))
                    contents.append({
                        "mime_type": mime_type,
                        "data": file_bytes
                    })
                except Exception as e:
                     return f"Error processing image: {e}"
            elif mime_type.startswith("video") or mime_type.startswith("audio") or mime_type == "application/pdf":
                contents.append({
                    "mime_type": mime_type,
                    "data": file_bytes
                })
            else:
                return "Unsupported file type. Please upload a valid image, video, audio, or PDF."

        if not contents:
            return "Please provide either text or a file."

        response = model.generate_content(contents)
        response.resolve()
        return response.text
    except Exception as e:
        return f"Error generating response: {e}"

# Gradio Interface
iface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(lines=3, placeholder="Enter text prompt here..."),
        gr.File(label="Upload a File (Image, Video, Audio, PDF)")
    ],
    outputs=gr.Textbox(label="Gemini Response"),
    title="Multimodal Gemini Chatbot",
    description="Upload an image, video, audio, or PDF, or enter text to interact with Gemini. You can provide both text and a file for multiprompt."
)

if __name__ == "__main__":
    iface.launch(share=False)