import streamlit as st
from PyPDF2 import PdfMerger
import fitz  # PyMuPDF
from PIL import Image
from io import BytesIO
import os
from streamlit_sortables import sort_items

# Function to generate thumbnail from PDF
def get_pdf_thumbnail(uploaded_file):
    uploaded_file.seek(0)  # Reset the file pointer to the beginning
    doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
    page = doc.load_page(0)
    pix = page.get_pixmap()
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    img.thumbnail((100, 140))
    return img

# App layout
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 for each PDF
    thumbnails = []
    filenames = []

    for file in uploaded_files:
        thumbnails.append(get_pdf_thumbnail(file))
        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), caption=file.name, use_container_width=True)

    # Merge PDFs in the specified order
    if st.button("Merge PDFs"):
        merger = PdfMerger()
        for file in reordered_files:
            file.seek(0)  # Reset the file pointer before merging
            merger.append(file)
    
        output_filename = "combined_document.pdf"
        with open(output_filename, "wb") as output_file:
            merger.write(output_file)
    
        st.success("PDF pages combined successfully!")
    
        # Provide download link
        with open(output_filename, "rb") as f:
            st.download_button("Download Combined PDF", f, file_name=output_filename, mime="application/pdf")
    
        # Clean up
        os.remove(output_filename)