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

# Configure the Gemini API
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

def upload_to_gemini(file_path, mime_type=None):
    """Uploads the given file to Gemini, returning the URI for use in chat."""
    file = genai.upload_file(file_path, mime_type=mime_type)
    print(f"Uploaded file '{file.display_name}' as: {file.uri}")
    return file.uri

def wait_for_files_active(file_uris):
    """Ensure all files are processed and ready for use."""
    print("Waiting for file processing...")
    for uri in file_uris:
        while True:
            file = genai.get_file(uri)
            if file.state.name == "ACTIVE":
                print("File ready.")
                break
            elif file.state.name == "FAILED":
                raise Exception(f"File {uri} failed to process")
            time.sleep(10)

def handle_files(files):
    """Handle file uploads, process them, and interact with Gemini API."""
    # Save uploaded files locally
    saved_files = []
    for file in files:
        local_path = f"/tmp/{file.name}"
        with open(local_path, "wb") as f:
            f.write(file.read())
        saved_files.append((local_path, mimetypes.guess_type(local_path)[0]))

    # Upload files to Gemini and wait for them to be ready
    file_uris = [upload_to_gemini(path, mime_type) for path, mime_type in saved_files]
    wait_for_files_active(file_uris)
    
    # Start a chat session with uploaded files
    model = genai.GenerativeModel(
        model_name="gemini-2.0-flash-exp",
        generation_config={
            "temperature": 0.4,
            "top_p": 0.95,
            "top_k": 40,
            "max_output_tokens": 8192,
            "response_mime_type": "text/plain",
        }
    )
    chat_session = model.start_chat(history=[{"role": "user", "parts": [{"file_uri": uri} for uri in file_uris]}])
    response = chat_session.send_message("Describe the contents.")
    return response.text

def gradio_interface():
    """Setup Gradio UI for file uploads and processing."""
    with gr.Blocks() as app:
        with gr.Row():
            file_input = gr.File(label="Upload files", file_types=[".pdf", ".png", ".jpg", ".mp3", ".mp4"], multiple=True)
            submit_button = gr.Button("Submit")
        output = gr.Textbox(label="Response from Gemini")

        submit_button.click(fn=handle_files, inputs=[file_input], outputs=[output])

    return app

app = gradio_interface()
app.launch()