import os
import time
import mimetypes
import google.generativeai as genai
import gradio as gr

# Configure the Gemini API
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
    raise ValueError("GEMINI_API_KEY not found. Please set it in your environment variables.")
genai.configure(api_key=api_key)

def upload_to_gemini(file_data, mime_type):
    """Uploads the given file data to Gemini and returns the file object if successful."""
    try:
        file = genai.upload_file(file_data, mime_type=mime_type)
        print(f"Uploaded file '{file.display_name}' as: {file.uri}")
        return file
    except Exception as e:
        print(f"Error uploading file: {e}")
        return None

def wait_for_files_active(files):
    """Waits for the uploaded files to be processed and become active."""
    for file in files:
        while file.state.name == "PROCESSING":
            print(".", end="", flush=True)
            time.sleep(10)
            file = genai.get_file(file.name)
        if file.state.name != "ACTIVE":
            raise Exception(f"File {file.display_name} failed to process, current state: {file.state.name}")
    print("All files ready.")

def gemini_interaction(input_text, *file_args):
    """Handles the interaction with Gemini API for both text and file inputs."""
    files = []
    for i in range(0, len(file_args), 2):
        if file_args[i] is not None:
            file_data = file_args[i].read()
            mime_type = file_args[i + 1]
            uploaded_file = upload_to_gemini(file_data, mime_type)
            if uploaded_file:
                files.append(uploaded_file)

    # Wait for all files to be ready before proceeding
    if files:
        wait_for_files_active(files)

    # Prepare the content for the Gemini API
    contents = []
    if input_text:
        contents.append({"text": input_text})
    contents.extend([{"file": file.uri} for file in files if file])

    # Initialize model and generate content
    model = genai.GenerativeModel("gemini-2.0-flash-exp")
    try:
        response_stream = model.generate_content(contents, stream=True)
        return ''.join([chunk.text for chunk in response_stream])
    except Exception as e:
        return f"Gemini API error: {e}"

# Define the Gradio interface
iface = gr.Interface(
    fn=gemini_interaction,
    inputs=[
        gr.Textbox(label="Enter your Prompt"),
        gr.File(label="Upload File 1"), gr.Dropdown(label="MIME Type for File 1", choices=["image/png", "audio/ogg", "application/pdf", "video/mp4"]),
        gr.File(label="Upload File 2"), gr.Dropdown(label="MIME Type for File 2", choices=["image/png", "audio/ogg", "application/pdf", "video/mp4"]),
        gr.File(label="Upload File 3"), gr.Dropdown(label="MIME Type for File 3", choices=["image/png", "audio/ogg", "application/pdf", "video/mp4"]),
        gr.File(label="Upload File 4"), gr.Dropdown(label="MIME Type for File 4", choices=["image/png", "audio/ogg", "application/pdf", "video/mp4"])
    ],
    outputs=gr.Textarea(),
    title="Interact with Gemini 2.0",
    description="Upload files and enter text to interact with the Gemini 2.0 model."
)

if __name__ == "__main__":
    iface.launch(debug=True)