import streamlit as st
import os
import google.generativeai as genai
import tempfile
import time

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

# Create the model
generation_config = {
    "temperature": 0.9,
    "top_p": 1.0,
    "top_k": 32,
    "max_output_tokens": 8192,
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-pro",
    generation_config=generation_config,
)

def upload_to_gemini(file_path, mime_type=None):
    """Uploads the given file to Gemini."""
    try:
        file = genai.upload_file(file_path, mime_type=mime_type)
        time.sleep(2)  # Short delay after uploading
        return file
    except Exception as e:
        st.error(f"Error uploading file: {str(e)}")
        return None

def process_file(file, prompt, mime_type):
    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{mime_type.split('/')[-1]}") as tmp_file:
        tmp_file.write(file.getvalue())
        tmp_file.flush()
        tmp_file_path = tmp_file.name

    try:
        uploaded_file = upload_to_gemini(tmp_file_path, mime_type=mime_type)
        if uploaded_file is None:
            return "File upload failed."

        response = model.generate_content([uploaded_file, prompt])
        return response.text
    except Exception as e:
        return f"Error processing file: {str(e)}"
    finally:
        os.unlink(tmp_file_path)

# Streamlit UI
st.title("File Analysis with Gemini")

# Sidebar for file type selection
file_type = st.sidebar.selectbox(
    "Choose file type",
    ["Image", "Video", "Audio", "PDF"]
)

# Main content area
st.subheader(f"Upload {file_type}")

uploaded_file = st.file_uploader(f"Choose a {file_type.lower()} file", type={"Image": ["png", "jpg", "jpeg"],
                                                                             "Video": ["mp4"],
                                                                             "Audio": ["mp3"],
                                                                             "PDF": ["pdf"]}[file_type])

user_prompt = st.text_area("Enter your prompt for analysis:", 
                           {"Image": "Describe this image in detail.",
                            "Video": "Provide a description of the video.",
                            "Audio": "Summarize the audio content and provide key points.",
                            "PDF": "Summarize the main points of this document."}[file_type])

if st.button("Analyze"):
    if uploaded_file is not None:
        with st.spinner(f"Processing {file_type.lower()}..."):
            mime_type = {"Image": "image/jpeg",
                         "Video": "video/mp4",
                         "Audio": "audio/mpeg",
                         "PDF": "application/pdf"}[file_type]
            
            result = process_file(uploaded_file, user_prompt, mime_type)
            
            if "Error" not in result:
                st.success(f"{file_type} processed successfully!")
                st.subheader("Analysis Result:")
                st.write(result)
            else:
                st.error(result)
    else:
        st.error(f"Please upload a {file_type.lower()} file.")

# Display the uploaded file
if uploaded_file is not None:
    if file_type == "Image":
        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
    elif file_type == "Video":
        st.video(uploaded_file)
    elif file_type == "Audio":
        st.audio(uploaded_file)
    elif file_type == "PDF":
        st.write("PDF uploaded successfully. Content cannot be displayed directly in Streamlit.")