import streamlit as st
import pymongo
import pandas as pd
import pytesseract
from PIL import Image
import fitz  # pymupdf for PDFs
import os

# MongoDB Connection
MONGO_URI = "mongodb+srv://viranggupta:<db_password>@gmp-21-19-virang.srl59.mongodb.net/?retryWrites=true&w=majority&appName=GMP-21-19-Virang"  # Change if using a remote DB
client = pymongo.MongoClient(MONGO_URI)
db = client["financial_analysis"]
collection = db["documents"]

# CSV File Path
CSV_FILE = "financial_analysis.csv"

# Streamlit App
st.title("📊 AI-Powered Financial Document Analyzer (Without Sentiment Model)")
st.write("Upload financial documents (PDFs or images) for text extraction.")

# File Upload
uploaded_file = st.file_uploader("Upload a PDF or Image", type=["pdf", "png", "jpg", "jpeg"])

def extract_text_from_pdf(pdf_file):
    """Extract text from a PDF using pymupdf"""
    text = ""
    doc = fitz.open(pdf_file)
    for page in doc:
        text += page.get_text("text") + "\n"
    return text

def extract_text_from_image(image_file):
    """Extract text from an image using pytesseract"""
    image = Image.open(image_file)
    text = pytesseract.image_to_string(image)
    return text

def save_to_csv(data):
    """Save extracted text data to a CSV file"""
    df = pd.DataFrame([data])
    
    # Check if file exists, append or create new
    if os.path.exists(CSV_FILE):
        df.to_csv(CSV_FILE, mode='a', index=False, header=False)
    else:
        df.to_csv(CSV_FILE, index=False)

if uploaded_file:
    file_extension = uploaded_file.name.split(".")[-1].lower()
    
    if file_extension in ["pdf"]:
        extracted_text = extract_text_from_pdf(uploaded_file)
    elif file_extension in ["png", "jpg", "jpeg"]:
        extracted_text = extract_text_from_image(uploaded_file)
    else:
        st.error("Unsupported file format.")
        st.stop()
    
    st.subheader("Extracted Text")
    st.text_area("Extracted Content", extracted_text, height=200)

    # Save to MongoDB
    doc_data = {"filename": uploaded_file.name, "text": extracted_text}
    collection.insert_one(doc_data)

    # Save to CSV
    save_to_csv(doc_data)

    st.success("📁 Document analyzed and saved to both MongoDB & CSV!")

# Download CSV Option
if os.path.exists(CSV_FILE):
    st.subheader("📥 Download Extracted Data")
    with open(CSV_FILE, "rb") as file:
        st.download_button(label="Download CSV", data=file, file_name="financial_analysis.csv", mime="text/csv")

st.write("🚀 Powered by MongoDB & CSV Storage")