import gradio as gr
import os
import tempfile
from PyPDF2 import PdfReader, PdfWriter
from reportlab.pdfgen import canvas
from PIL import Image, ImageDraw
import shutil

class MultiDocMerger:
    def __init__(self):
        self.temp_dir = tempfile.mkdtemp()
        self.documents = []
        self.pages = []

    def add_document(self, file):
        file_path = os.path.join(self.temp_dir, file.name)
        with open(file_path, 'wb') as f:
            f.write(file.file.read())  # Use file.file.read() instead of file.read()
        self.documents.append(file_path)
        self._extract_pages(file_path)
        return self._generate_preview()
        self.documents.append(file_path)
        self._extract_pages(file_path)
        return self._generate_preview()

    def _extract_pages(self, file_path):
        reader = PdfReader(file_path)
        for page_num, page in enumerate(reader.pages):
            temp_pdf = os.path.join(self.temp_dir, f"page_{len(self.pages)}.pdf")
            writer = PdfWriter()
            writer.add_page(page)
            with open(temp_pdf, 'wb') as f:
                writer.write(f)
            self.pages.append(temp_pdf)

    def rearrange_documents(self, order):
        self.documents = [self.documents[i] for i in order]
        return self._generate_preview()

    def rearrange_pages(self, order):
        self.pages = [self.pages[i] for i in order]
        return self._generate_preview()

    def delete_page(self, index):
        if 0 <= index < len(self.pages):
            del self.pages[index]
        return self._generate_preview()

    def add_blank_page(self):
        blank_path = os.path.join(self.temp_dir, f"blank_{len(self.pages)}.pdf")
        can = canvas.Canvas(blank_path)
        can.showPage()
        can.save()
        self.pages.append(blank_path)
        return self._generate_preview()

    def generate_index(self):
        index_content = "Index Page\n\n"
        for i, doc in enumerate(self.documents):
            index_content += f"{i + 1}. {os.path.basename(doc)}\n"
        index_pdf = os.path.join(self.temp_dir, "index.pdf")
        can = canvas.Canvas(index_pdf)
        can.drawString(100, 750, index_content)
        can.save()
        self.pages.insert(0, index_pdf)

    def export_merged_pdf(self):
        self.generate_index()
        output = PdfWriter()
        for page in self.pages:
            reader = PdfReader(page)
            output.add_page(reader.pages[0])
        output_path = os.path.join(self.temp_dir, 'merged_output.pdf')
        with open(output_path, 'wb') as f:
            output.write(f)
        return output_path

    def _generate_preview(self):
        return [os.path.basename(page) for page in self.pages]

    def cleanup(self):
        shutil.rmtree(self.temp_dir)

merger = MultiDocMerger()

with gr.Blocks() as demo:
    gr.Markdown("# Multi-Document Merger and Organizer")
    with gr.Column():
        file_upload = gr.File(label="Upload Files", file_count="multiple")
        preview_gallery = gr.Gallery(label="Document Preview and Reordering", interactive=True)
        delete_index = gr.Number(label="Delete Page Index")
        blank_button = gr.Button("Add Blank Page")
        merge_button = gr.Button("Merge and Export PDF")
        output_file = gr.File()

    def upload_files(files):
        paths = []
        for file in files:
            paths = merger.add_document(file)
        return paths

    def merge_and_export():
        return merger.export_merged_pdf()

    def delete_page(index):
        return merger.delete_page(index)

    def add_blank_page():
        return merger.add_blank_page()

    def reorder_pages(order):
        return merger.rearrange_pages(order)

    file_upload.upload(upload_files, inputs=file_upload, outputs=preview_gallery)
    delete_index.submit(delete_page, inputs=delete_index, outputs=preview_gallery)
    blank_button.click(add_blank_page, outputs=preview_gallery)
    preview_gallery.change(reorder_pages, inputs=preview_gallery, outputs=preview_gallery)
    merge_button.click(merge_and_export, outputs=output_file)

demo.launch()