app.py

import os
import streamlit as st
from PyPDF2 import PdfMerger
from google.cloud import vision
from google.oauth2 import service_account
import fitz
from PIL import Image
from io import BytesIO
import tempfile
import json
from streamlit_sortables import sort_items
import subprocess
from tqdm import tqdm
from datetime import datetime
import uuid

#os.environ["PATH"] += os.pathsep + "/usr/bin" + os.pathsep + "/usr/local/bin" #Removed redundant path setting
print(f"Current PATH: {os.environ['PATH']}")


# Check if /usr/bin/pdfinfo exists
if os.path.exists("/usr/bin/pdfinfo"):
    print("pdfinfo exists at /usr/bin/pdfinfo")
    # Check the file permissions
    permissions = os.stat("/usr/bin/pdfinfo").st_mode
    print(f"File permissions: {oct(permissions)}")
else:
    print("pdfinfo does not exist at /usr/bin/pdfinfo")


# Load Google Cloud Vision credentials from secret
credentials_json = os.getenv("GOOGLE_CREDENTIALS_JSON")
if credentials_json:
    credentials_dict = json.loads(credentials_json)
    credentials = service_account.Credentials.from_service_account_info(credentials_dict)
    client = vision.ImageAnnotatorClient(credentials=credentials)
else:
    client = None

# Function to extract text using Google Cloud Vision
def extract_text_with_google_vision(image_bytes):
    """Extracts text using Google Cloud Vision."""
    image = vision.Image(content=image_bytes)
    response = client.document_text_detection(image=image)
    if response.error.message:
        raise Exception(f"Google Cloud Vision API Error: {response.error.message}")
    return response.full_text_annotation.text if response.full_text_annotation else ""

# Function to process PDF for transcription
def process_pdf(file):
    """Converts PDF pages to images and extracts text."""
    text = ""
    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_pdf:
        temp_pdf.write(file.read())
        temp_pdf_path = temp_pdf.name

    try:
        doc = fitz.open(temp_pdf_path)
        for i in tqdm(range(len(doc)), desc="Processing pages"):
            page = doc.load_page(i)
            pix = page.get_pixmap()
            image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            image_bytes = BytesIO()
            image.save(image_bytes, format="PNG")
            image_bytes.seek(0)
            try:
                page_text = extract_text_with_google_vision(image_bytes.getvalue())
                text += f"--- Page {i + 1} ---\n{page_text}\n\n"
            except Exception as e:
                st.error(f"Error on page {i + 1}: {e}")
    finally:
        os.remove(temp_pdf_path)

    with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as temp_txt:
      temp_txt.write(text.encode("utf-8"))
      temp_txt_path = temp_txt.name
    return temp_txt_path


# Function to generate thumbnail from PDF
def get_pdf_thumbnail(uploaded_file, page_num, rotation=0):
    """Generates a thumbnail image of a PDF page."""
    uploaded_file.seek(0)
    doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
    page = doc.load_page(page_num)
    page.set_rotation(rotation)
    pix = page.get_pixmap()
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    img.thumbnail((400, 560))
    return img

# Function to merge PDFs
def merge_pdfs(reordered_pages, uploaded_file):
    """Merges multiple PDFs into one."""
    merger = PdfMerger()
    uploaded_file.seek(0)
    doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
    temp_files = []

    for page_data in reordered_pages:
        page = doc.load_page(page_data['page_index'])
        page.set_rotation(page_data['rotation'])
        temp_pdf_page = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False)
        page.get_pixmap().save(temp_pdf_page.name)
        temp_files.append(temp_pdf_page.name)
        merger.append(temp_pdf_page.name)

    output_filename = "combined_document.pdf"
    with open(output_filename, "wb") as output_file:
      merger.write(output_file)

    for file in temp_files:
      os.remove(file)
    return output_filename

def download_file(output_file, file_name, mime_type):
    with open(output_file, "rb") as f:
        st.download_button(
            label="Download File",
            data=f,
            file_name=file_name,
            mime=mime_type,
        )

def generate_user_filename(filename, file_type):
    """Generates a filename based on user input and file type."""
    return f"{filename}{file_type}"

def generate_unique_filename(original_filename, suffix, file_type):
    """Generates a unique filename based on date, time, and original filename."""
    now = datetime.now()
    timestamp = now.strftime("%Y%m%d_%H%M%S")
    name, ext = os.path.splitext(original_filename)
    return f"{name}_{timestamp}{suffix}{file_type}"

def rotate_pdf(pdf_file, rotation_angle):
    """Rotates all pages of a PDF by a specified angle."""
    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_pdf:
        temp_pdf.write(pdf_file.read())
        temp_pdf_path = temp_pdf.name

    doc = fitz.open(temp_pdf_path)
    for page in doc:
        page.set_rotation(rotation_angle)
    os.remove(temp_pdf_path)
    return doc

def display_pdf_preview(doc, page_num=0):
    """Displays a preview of a PDF page."""
    page = doc.load_page(page_num)
    pix = page.get_pixmap()
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    st.image(img, use_container_width=True)

# Sidebar Navigation
st.sidebar.title("Tool Selector")
selection = st.sidebar.radio("Choose a tool:", ["PDF Combiner", "PDF Transcriber", "PDF Rotator", "PDF Document Separator"])

# PDF Document Separator Tool
if selection == "PDF Document Separator":
    st.title("PDF Document Separator")
    st.write("Upload a multi-document PDF and separate the documents out.")

    uploaded_file = st.file_uploader("Upload a PDF", type="pdf")

    if uploaded_file:
        # Initialize session state variables if not already defined
        if "uploaded_pdf" not in st.session_state:
            st.session_state.uploaded_pdf = None
        if "target_pages" not in st.session_state:
            st.session_state.target_pages = [] # List of dictionaries of "page_index", "rotation", "uuid", "image"
        if "available_pages" not in st.session_state:
            st.session_state.available_pages = []
        if "pdf_filename" not in st.session_state:
            st.session_state.pdf_filename = "combined_document"
        if "transcribe_filename" not in st.session_state:
            st.session_state.transcribe_filename = "transcribed_text"
        st.session_state.uploaded_pdf = uploaded_file

        def reset_state():
            """Resets the session state for the PDF Document Separator tool."""
            st.session_state.uploaded_pdf = None
            st.session_state.target_pages = []
            st.session_state.available_pages = []


        if not st.session_state.available_pages:
          uploaded_file.seek(0)
          doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
          for page_num in range(len(doc)):
            image = get_pdf_thumbnail(uploaded_file, page_num)
            st.session_state.available_pages.append({"page_index":page_num, "rotation":0, "uuid":str(uuid.uuid4()), "image":image})

        def move_page_to_target(page_uuid):
          """Moves a page from available_pages to target_pages."""
          page_to_move = next((page for page in st.session_state.available_pages if page["uuid"] == page_uuid), None)
          if page_to_move:
            st.session_state.target_pages.append(page_to_move)
            st.session_state.available_pages = [page for page in st.session_state.available_pages if page["uuid"] != page_uuid]

        def remove_page_from_available(page_uuid):
            """Removes a page from available_pages."""
            st.session_state.available_pages = [
                page for page in st.session_state.available_pages if page["uuid"] != page_uuid
            ]


        left_col, right_col = st.columns(2)

        with left_col:
            st.header("Document Overview")
            st.write("Drag pages to the right to start creating the new document.")
            for page in st.session_state.available_pages:
              col1, col2 = st.columns([1,2])
              with col1:
                  st.image(page["image"], caption=f"Page {page['page_index'] + 1}", use_container_width=True)
              with col2:
                if st.button(f"Add Page", key=f"add_{page['uuid']}"):
                  move_page_to_target(page["uuid"])
                if st.button("Remove Page", key=f"remove_{page['uuid']}"):
                  remove_page_from_available(page["uuid"])


        with right_col:
            st.header("Target Document Builder")
            st.write("Reorder the pages, change rotation, and create the final document.")

            if st.session_state.target_pages:
                # Use sort_items for the right-hand column to manage reordering
                target_page_uuids = [page["uuid"] for page in st.session_state.target_pages]
                reordered_page_uuids = sort_items(target_page_uuids)
                reordered_target_pages = [next(page for page in st.session_state.target_pages if page['uuid'] == uuid) for uuid in reordered_page_uuids]
                st.session_state.target_pages = reordered_target_pages

                for page in st.session_state.target_pages:
                    with st.expander(f"Page {page['page_index'] + 1} Options", expanded=True):
                        col1, col2 = st.columns([1,1])
                        with col1:
                          st.image(page["image"], use_container_width=True)
                        with col2:
                          rotation = st.selectbox(f"Rotation", [0, 90, 180, 270], key=page["uuid"], index = [0,90,180,270].index(page['rotation']))
                          if rotation != page['rotation']:
                            page['rotation'] = rotation
                            page['image'] = get_pdf_thumbnail(st.session_state.uploaded_pdf, page['page_index'], rotation)


            st.text_input("Enter output PDF filename", key='pdf_filename_input', value = st.session_state.pdf_filename)
            if st.button("Create PDF"):
              st.session_state.pdf_filename = st.session_state.pdf_filename_input
              with st.spinner("Creating PDF..."):
                  output_file = merge_pdfs(st.session_state.target_pages, st.session_state.uploaded_pdf)
                  st.success("PDF created successfully!")
                  download_file(output_file, generate_user_filename(st.session_state.pdf_filename, ".pdf"), "application/pdf")
                  os.remove(output_file)
            st.text_input("Enter output transcription filename", key='transcribe_filename_input', value = st.session_state.transcribe_filename)
            if st.button("Transcribe PDF"):
                st.session_state.transcribe_filename = st.session_state.transcribe_filename_input
                with st.spinner("Processing..."):
                    output_file = process_pdf(st.session_state.uploaded_pdf)
                    if output_file:
                        st.success("Text extraction complete!")
                        st.text_area("Extracted Text", open(output_file, encoding="utf-8").read(), height=400)
                        output_file_name = generate_unique_filename(st.session_state.transcribe_filename, "(T)", ".txt")
                        with open(output_file, "rb") as f:
                            st.download_button("Download Extracted Text", f, file_name=output_file_name, mime="text/plain")
                        os.remove(output_file)
                    else:
                        st.error("Could not process the PDF, please try again")

        if st.button("Reset State"):
          reset_state()


# The rest of your original app.py code should go here
# PDF Combiner with Preview and Reordering
elif selection == "PDF Combiner":
    st.title("PDF Combiner with Preview & Reordering")
    st.write("Upload individual PDF pages, visualize them, reorder, and merge into a single PDF.")

    uploaded_files = st.file_uploader("Upload PDF pages", type="pdf", accept_multiple_files=True)

    if uploaded_files:
        # Generate thumbnails and filenames for each uploaded PDF
        thumbnails = []
        filenames = []

        for file in uploaded_files:
            thumbnails.append(get_pdf_thumbnail(file, 0))
            filenames.append(file.name)

        # Display thumbnails with filenames for reordering
        st.write("**Drag and drop to reorder the PDFs:**")
        reordered_filenames = sort_items(filenames)

        # Map the filenames back to the corresponding files
        reordered_files = [uploaded_files[filenames.index(name)] for name in reordered_filenames]

        # Display the thumbnails in the new order
        st.write("**Preview of selected order:**")
        cols = st.columns(len(reordered_files))
        for idx, file in enumerate(reordered_files):
            with cols[idx]:
                st.image(get_pdf_thumbnail(file, 0), caption=file.name, use_container_width=True)

        # Merge PDFs in the specified order
        if st.button("Merge PDFs"):
            output_file = merge_pdfs(reordered_files)
            st.success("PDF pages combined successfully!")
            download_file(output_file, generate_unique_filename("combined_document", "", ".pdf"), "application/pdf")
            os.remove(output_file)

# PDF Transcriber Tool
elif selection == "PDF Transcriber":
    st.title("PDF Transcriber Tool")
    st.write("Upload a scanned PDF to transcribe the text.")
    if not client:
        st.error("Google Cloud credentials are not set. Please configure the secret.")
    uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"])
    if uploaded_file and st.button("Transcribe PDF"):
        with st.spinner("Processing..."):
            output_file = process_pdf(uploaded_file)
            if output_file:
                st.success("Text extraction complete!")
                st.text_area("Extracted Text", open(output_file, encoding="utf-8").read(), height=400)
                output_file_name = generate_unique_filename(uploaded_file.name, "(T)", ".txt")
                with open(output_file, "rb") as f:
                  st.download_button("Download Extracted Text", f, file_name=output_file_name, mime="text/plain")
                os.remove(output_file)
            else:
                st.error("Could not process the PDF, please try again")

# PDF Rotator Tool
elif selection == "PDF Rotator":
    st.title("PDF Rotator")
    st.write("Upload a PDF and rotate all pages by a specified angle.")

    pdf_file = st.file_uploader("Upload PDF", type=["pdf"])
    rotation_angle = st.selectbox("Select rotation angle:", [90, 180, 270])

    if pdf_file:
      if 'rotated_doc' not in st.session_state:
          st.session_state.rotated_doc = None
      if st.button("Rotate PDF"):
        with st.spinner("Rotating PDF..."):
          st.session_state.rotated_doc = rotate_pdf(pdf_file, rotation_angle)
          st.success("PDF rotated successfully!")
      if st.session_state.rotated_doc:
          display_pdf_preview(st.session_state.rotated_doc)
          if st.button("Download Rotated PDF"):
              with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_pdf:
                  st.session_state.rotated_doc.save(temp_pdf.name)
                  download_file(temp_pdf.name, generate_unique_filename(pdf_file.name, "(R)", ".pdf"), "application/pdf")
                  os.remove(temp_pdf.name)