import streamlit as st
from PyPDF2 import PdfReader, PdfWriter

def split_pdf(input_file, output_files_with_ranges):
    """
    Splits a PDF into multiple PDFs based on specified page ranges.

    Parameters:
        input_file (str): Path to the input PDF file.
        output_files_with_ranges (list of tuples): Each tuple contains the output filename and the page range (start, end).
            Example: [("output1.pdf", (1, 3)), ("output2.pdf", (4, 6))]

    Returns:
        None
    """
    try:
        reader = PdfReader(input_file)

        for output_file, page_range in output_files_with_ranges:
            writer = PdfWriter()
            start, end = page_range

            # Add specified page range to writer
            for page_num in range(start - 1, end):
                if 0 <= page_num < len(reader.pages):
                    writer.add_page(reader.pages[page_num])
                else:
                    raise ValueError(f"Page number {page_num + 1} out of range for input file.")

            # Write the output file
            with open(output_file, "wb") as f:
                writer.write(f)

        st.success("PDF split successfully!")
    except Exception as e:
        st.error(f"An error occurred: {e}")

# Streamlit app
st.title("PDF Splitter App")

uploaded_file = st.file_uploader("Upload a PDF file", type="pdf")
if uploaded_file:
    input_pdf_path = f"uploaded_{uploaded_file.name}"
    with open(input_pdf_path, "wb") as f:
        f.write(uploaded_file.read())

    reader = PdfReader(input_pdf_path)
    total_pages = len(reader.pages)
    st.write(f"The uploaded PDF has {total_pages} pages.")

    num_splits = st.number_input("How many splits do you want?", min_value=1, max_value=total_pages, step=1, value=1)

    output_files_with_ranges = []
    download_links = []
    for i in range(num_splits):
        st.write(f"### Split {i + 1}")
        start = st.number_input(f"Start page for split {i + 1}", min_value=1, max_value=total_pages, step=1, value=1, key=f"start_{i}")
        end = st.number_input(f"End page for split {i + 1}", min_value=1, max_value=total_pages, step=1, value=1, key=f"end_{i}")
        output_filename = st.text_input(f"Output filename for split {i + 1}", value=f"PDF_Part{i + 1}.pdf", key=f"filename_{i}")

        if start <= end:
            output_files_with_ranges.append((output_filename, (start, end)))
        else:
            st.warning(f"End page must be greater than or equal to start page for split {i + 1}.")

    if st.button("Split PDF"):
        split_pdf(input_pdf_path, output_files_with_ranges)

        for output_file, _ in output_files_with_ranges:
            with open(output_file, "rb") as f:
                download_links.append((output_file, f.read()))

        for filename, data in download_links:
            st.download_button(label=f"Download {filename}", data=data, file_name=filename, mime="application/pdf")