app.py

import streamlit as st
import os
import pandas as pd
import fitz  # PyMuPDF
from google import genai
from google.genai import types
from docx import Document
from pptx import Presentation
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet
from bs4 import BeautifulSoup
import re
import base64

# Set Gemini API Key
client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))

# --------- File Extractors ---------
def extract_text_from_docx(docx_file):
    document = Document(docx_file)
    return "\n".join([para.text for para in document.paragraphs])

def extract_text_from_csv(csv_file):
    df = pd.read_csv(csv_file)
    return df.to_string(index=False)

def extract_text_from_xlsx(xlsx_file):
    df = pd.read_excel(xlsx_file)
    return df.to_string(index=False)

def extract_text_from_pptx(pptx_file):
    prs = Presentation(pptx_file)
    text_runs = []
    for slide in prs.slides:
        for shape in slide.shapes:
            if hasattr(shape, "text"):
                text_runs.append(shape.text)
    return "\n".join(text_runs)

def extract_text_from_pdf(pdf_file):
    doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
    text = ""
    for page in doc:
        text += page.get_text()
    return text

def extract_text_from_html(html_file):
    soup = BeautifulSoup(html_file.read(), "html.parser")
    return soup.get_text()

def extract_text_from_tex(tex_file):
    content = tex_file.read().decode("utf-8")
    content = re.sub(r'\\[a-zA-Z]+\{[^}]*\}', '', content)
    content = re.sub(r'\\[a-zA-Z]+', '', content)
    return content

def process_image(image_file):
    image_bytes = image_file.read()
    encoded_image = base64.b64encode(image_bytes).decode("utf-8")
    return {
        "inline_data": {
            "mime_type": image_file.type,
            "data": encoded_image
        }
    }

def process_video(video_file):
    video_bytes = video_file.read()
    filetype = video_file.name.split(".")[-1].lower()
    mime_map = {
        "avi": "video/x-msvideo",
        "mkv": "video/x-matroska",
        "flv": "video/x-flv",
        "wmv": "video/x-ms-wmv",
        "mpeg": "video/mpeg"
    }
    mime_type = mime_map.get(filetype, video_file.type)
    encoded_video = base64.b64encode(video_bytes).decode("utf-8")
    return {
        "inline_data": {
            "mime_type": mime_type,
            "data": encoded_video
        }
    }

def export_conversation_to_pdf(conversation_history):
    pdf_path = "Conversation.pdf"
    doc = SimpleDocTemplate(pdf_path, pagesize=A4)
    styles = getSampleStyleSheet()
    elements = []

    for i, (user_q, gemini_a) in enumerate(conversation_history):
        elements.append(Paragraph(f"<b>Q{i+1}:</b> {user_q}", styles["Normal"]))
        elements.append(Spacer(1, 8))
        elements.append(Paragraph(f"<b>A{i+1}:</b> {gemini_a.replace(chr(10), '<br/>')}", styles["Normal"]))
        elements.append(Spacer(1, 16))

    doc.build(elements)
    return pdf_path

# --------- Main App ---------
def main():
    st.set_page_config(page_title="Gemini 2.5 Flash Chat Q&A", layout="wide")
    st.title("\U0001F4C4 Product Prototype with Document Support and Internet Search")

    if "conversation" not in st.session_state:
        st.session_state.conversation = []
    if "documents_text" not in st.session_state:
        st.session_state.documents_text = []
    if "chat_active" not in st.session_state:
        st.session_state.chat_active = True
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

    uploaded_files = st.file_uploader(
        "Upload files (.txt, .docx, .csv, .xlsx, .pptx, .pdf, .html, .htm, .tex, .jpg, .jpeg, .png, video formats):",
        type=[
            "txt", "docx", "csv", "xlsx", "pptx", "pdf", "html", "htm", "tex",
            "jpg", "jpeg", "png",
            "mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"
        ],
        accept_multiple_files=True
    )

    if uploaded_files:
        all_text_content = ""
        for uploaded_file in uploaded_files:
            filetype = uploaded_file.name.split(".")[-1].lower()
            try:
                if filetype == "txt":
                    all_text_content += uploaded_file.read().decode("utf-8") + "\n"
                elif filetype == "docx":
                    all_text_content += extract_text_from_docx(uploaded_file) + "\n"
                elif filetype == "csv":
                    all_text_content += extract_text_from_csv(uploaded_file) + "\n"
                elif filetype == "xlsx":
                    all_text_content += extract_text_from_xlsx(uploaded_file) + "\n"
                elif filetype == "pptx":
                    all_text_content += extract_text_from_pptx(uploaded_file) + "\n"
                elif filetype == "pdf":
                    all_text_content += extract_text_from_pdf(uploaded_file) + "\n"
                elif filetype in ["html", "htm"]:
                    all_text_content += extract_text_from_html(uploaded_file) + "\n"
                elif filetype == "tex":
                    all_text_content += extract_text_from_tex(uploaded_file) + "\n"
                elif filetype in ["jpg", "jpeg", "png"]:
                    st.session_state.image_data = process_image(uploaded_file)
                    all_text_content += "Image uploaded and processed.\n"
                    st.image(uploaded_file, caption=uploaded_file.name, use_container_width=True)
                elif filetype in ["mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"]:
                    st.session_state.video_data = process_video(uploaded_file)
                    all_text_content += f"Video file '{uploaded_file.name}' uploaded and processed.\n"
                    st.video(uploaded_file)
                else:
                    st.error(f"Unsupported file format: {uploaded_file.name}")
            except Exception as e:
                st.error(f"Failed to extract/process {uploaded_file.name}: {e}")

        st.session_state.documents_text = all_text_content

    if st.session_state.documents_text:
        st.markdown("### \U0001F4AC Conversation")
        for user_q, gemini_a in st.session_state.conversation:
            st.markdown(f"**You:** {user_q}")
            st.markdown(f"**Gemini:**\n\n{gemini_a}")

        if st.session_state.chat_active:
            with st.form(key="chat_form", clear_on_submit=True):
                user_input = st.text_input("Ask a question about the documents (type 'exit' to stop):")
                submit = st.form_submit_button("Send")

                if submit and user_input:
                    if user_input.strip().lower() == "exit":
                        st.session_state.chat_active = False
                        st.success("Chat ended. Reload to start again.")
                    else:
                        with st.spinner("Let me think..."):
                            content_blocks = []

                            # System prompt inserted first here:
                            content_blocks.append({
                                "text": (
                                    "Your role is to analyze documents, images, and videos to help identify and classify driver states, "
                                    "including: drowsy, alert, yawning, microsleep, eyes/lips open or closed, and other distractions. "
                                    "The user may ask for timestamps of events in video, technical explanations from code/docs, "
                                    "or structured summaries. Prioritize your answers for this use case and act as a helpful assistant "
                                    "within this domain. "
                                    "Additionally, your classification may rely on the State Farm Distracted Driver Detection dataset, "
                                    "which defines 10 distraction categories such as texting, phone use, reaching behind, adjusting the radio, "
                                    "and safe driving. Use this labeling scheme when interpreting relevant images."
                                )
                            })

                            if st.session_state.conversation:
                                history_text = "\n\n".join(
                                    f"Q: {entry[0]}\nA: {entry[1]}"
                                    for entry in st.session_state.conversation
                                )
                                content_blocks.append({"text": f"Previous conversation:\n{history_text}"})

                            if st.session_state.documents_text.strip():
                                content_blocks.append({"text": f"Context:\n{st.session_state.documents_text}"})

                            if "image_data" in st.session_state:
                                content_blocks.append(st.session_state.image_data)

                            if "video_data" in st.session_state:
                                content_blocks.append(st.session_state.video_data)

                            content_blocks.append({"text": f"Question: {user_input}"})

                            response = client.models.generate_content(
                                model="gemini-2.5-flash",
                                contents=content_blocks,
                                config={"tools": [{"google_search": {}}]}
                            )

                            st.session_state.conversation.append((user_input, response.text))
                            st.success("\U0001F4A1 Answer:")
                            st.write(response.text)

                            st.session_state.chat_history.append({
                                "question": user_input,
                                "answer": response.text
                            })

        if st.session_state.conversation:
            pdf_path = export_conversation_to_pdf(st.session_state.conversation)
            with open(pdf_path, "rb") as f:
                st.download_button(
                    label="\U0001F4E5 Export Conversation as PDF",
                    data=f,
                    file_name="Conversation.pdf",
                    mime="application/pdf"
                )

if __name__ == "__main__":
    main()