import os
import gradio as gr
from PyPDF2 import PdfWriter, PdfReader

def split_pdf(input_file, pages_per_split):
    # Open the input PDF file
    with open(input_file.name, 'rb') as file:
        # Create a PDF reader object
        pdf_reader = PdfReader(file)

        # Get the total number of pages in the PDF
        total_pages = len(pdf_reader.pages)

        # Calculate the number of splits
        num_splits = (total_pages + pages_per_split - 1) // pages_per_split

        # Create a directory to store the split PDF files
        output_dir = 'split_pdfs'
        os.makedirs(output_dir, exist_ok=True)

        # List to store the paths of the split PDF files
        split_files = []

        # Iterate over the number of splits
        for i in range(num_splits):
            # Create a new PDF writer object for each split
            pdf_writer = PdfWriter()

            # Calculate the start and end page numbers for the current split
            start_page = i * pages_per_split
            end_page = min((i + 1) * pages_per_split, total_pages)

            # Add the pages to the PDF writer object
            for page in range(start_page, end_page):
                pdf_writer.add_page(pdf_reader.pages[page])

            # Generate the output file name for the current split
            output_file = os.path.join(output_dir, f'split_{i+1}.pdf')

            # Open the output file in write-binary mode
            with open(output_file, 'wb') as output:
                # Write the pages to the output file
                pdf_writer.write(output)

            # Add the path of the split PDF file to the list
            split_files.append(output_file)

        return split_files

# Gradio interface
def gradio_interface(input_file, pages_per_split):
    split_files = split_pdf(input_file, pages_per_split)
    return split_files

# Create a Gradio interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[gr.File(label="Input PDF"), gr.Number(label="Pages per Split")],
    outputs=gr.File(label="Split PDFs"),
    title="PDF Splitter",
    description="Split a large PDF file into smaller PDF files.",
)

# Launch the Gradio interface
iface.launch()