import os
import base64
import streamlit as st
from docx import Document
import openpyxl
from pdfminer.high_level import extract_text
import csv
from pptx import Presentation
from io import StringIO
from bs4 import BeautifulSoup
import re
import google.generativeai as genai

# Setup
st.set_page_config(page_title="Gemini Q&A App", layout="centered")
st.title("🤖 Q&A - Document & Image")

# Initialize chat history
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Optional: Clear history button
if st.button("🧹 Clear Chat History"):
    st.session_state.chat_history = []

# Load API key
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY") or st.secrets.get("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
    st.error("GOOGLE_API_KEY not found in environment or Streamlit secrets.")
    st.stop()

genai.configure(api_key=GOOGLE_API_KEY)
model = genai.GenerativeModel(model_name="gemini-2.0-flash")

# Helper: Extract text from various documents
def load_and_extract_text(file_path):
    try:
        ext = file_path.split('.')[-1].lower()

        if ext == 'docx':
            doc = Document(file_path)
            return '\n'.join(p.text for p in doc.paragraphs)

        elif ext == 'xlsx':
            workbook = openpyxl.load_workbook(file_path)
            return '\n'.join(
                ' '.join(str(cell.value) for cell in row if cell.value is not None)
                for sheet in workbook for row in sheet.iter_rows()
            )

        elif ext == 'pdf':
            return extract_text(file_path)

        elif ext == 'csv':
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                return '\n'.join(' '.join(row) for row in reader)

        elif ext == 'txt':
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()

        elif ext == 'pptx':
            prs = Presentation(file_path)
            return '\n'.join(
                shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, 'text')
            )

        elif ext == 'tex':
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
                content = re.sub(r'\\[a-zA-Z]+\*?(?:\[[^\]]*\])?(?:\{[^}]*\})*', '', content)
                content = re.sub(r'%.*', '', content)
                return content.strip()

        elif ext in ['html', 'htm']:
            with open(file_path, 'r', encoding='utf-8') as f:
                soup = BeautifulSoup(f, 'html.parser')
                for script_or_style in soup(['script', 'style']):
                    script_or_style.decompose()
                return soup.get_text(separator='\n', strip=True)

        else:
            return "Unsupported file format."

    except Exception as e:
        return f"Error processing file: {e}"

# Multi-file uploader
uploaded_files = st.file_uploader(
    "Upload one or more documents or images",
    type=["docx", "xlsx", "pdf", "csv", "txt", "pptx", "tex", "html", "htm", "jpg", "jpeg", "png"],
    accept_multiple_files=True
)

# Question input
question = st.text_input("💬 Enter your question:")

# Processing
if uploaded_files and question:
    combined_context = ""
    image_contents = []

    for uploaded_file in uploaded_files:
        file_ext = uploaded_file.name.split('.')[-1].lower()

        if file_ext in ['jpg', 'jpeg', 'png']:
            image_bytes = uploaded_file.read()
            encoded_image = base64.b64encode(image_bytes).decode("utf-8")
            image_contents.append({
                "inline_data": {
                    "mime_type": uploaded_file.type,
                    "data": encoded_image
                }
            })
            st.image(image_bytes, caption=uploaded_file.name, use_container_width=True)
            continue

        temp_file_path = "temp." + file_ext
        with open(temp_file_path, "wb") as temp_file:
            temp_file.write(uploaded_file.read())

        extracted_text = load_and_extract_text(temp_file_path)
        os.remove(temp_file_path)

        if "Error" in extracted_text or "Unsupported" in extracted_text:
            st.error(f"{uploaded_file.name}: {extracted_text}")
        else:
            combined_context += f"\n\n---\nDocument: {uploaded_file.name}\n\n{extracted_text}"

    if combined_context.strip() or image_contents:
        with st.spinner("Generating answer..."):
            try:
                content_blocks = []

                # Add prior Q&A as part of the context
                if st.session_state.chat_history:
                    history_text = "\n\n".join(
                        f"Q: {entry['question']}\nA: {entry['answer']}"
                        for entry in st.session_state.chat_history
                    )
                    content_blocks.append({"text": f"Previous conversation:\n{history_text}"})

                # Add file text content
                if combined_context.strip():
                    content_blocks.append({"text": f"Context:\n{combined_context}"})

                # Add image blocks
                content_blocks.extend(image_contents)

                # Add current question
                content_blocks.append({"text": f"Question: {question}"})

                response = model.generate_content(contents=content_blocks)
                st.success("💡 Answer:")
                st.write(response.text)

                # Save to session history
                st.session_state.chat_history.append({
                    "question": question,
                    "answer": response.text
                })

            except Exception as e:
                st.error(f"Error generating answer: {e}")
    else:
        st.warning("No valid documents or images processed.")

# Optional: Show full conversation history
if st.session_state.chat_history:
    st.markdown("### 🗂️ Chat History")
    for i, entry in enumerate(st.session_state.chat_history, 1):
        st.markdown(f"**Q{i}:** {entry['question']}")
        st.markdown(f"**A{i}:** {entry['answer']}")