import io
from PyPDF2 import PdfReader, PdfWriter, PageObject
import streamlit as st

def validate_pdf(file_stream):
    """
    Validates if the file is a proper PDF by trying to read it.
    """
    try:
        PdfReader(file_stream)
        return True
    except Exception:
        return False

def merge_two_pages_on_one(page1, page2):
    """
    Merges two PDF pages into a single PDF page (side by side).
    """
    width = page1.mediabox.width + page2.mediabox.width
    height = max(page1.mediabox.height, page2.mediabox.height)
    
    new_page = PageObject.create_blank_page(width=width, height=height)
    new_page.merge_page(page1)

    page2_content = PageObject.create_blank_page(width=width, height=height)
    page2_content.merge_page(page2)
    page2_content.add_transformation([1, 0, 0, 1, page1.mediabox.width, 0])  # x translation by page1 width

    new_page.merge_page(page2_content)

    return new_page

def reverse_pdf(input_pdf_stream):
    """
    Reverses the order of pages in a PDF and returns the reversed file stream.
    """
    reader = PdfReader(input_pdf_stream)
    writer = PdfWriter()

    for page_num in reversed(range(len(reader.pages))):
        writer.add_page(reader.pages[page_num])

    output_stream = io.BytesIO()
    writer.write(output_stream)
    output_stream.seek(0)
    return output_stream

def create_booklet_with_two_pages_per_sheet(pdf_file):
    """
    Processes the uploaded PDF to create a booklet with two pages per sheet.
    """
    input_pdf_stream = io.BytesIO(pdf_file.read())

    # Validate the PDF file
    if not validate_pdf(input_pdf_stream):
        raise ValueError("Uploaded file is not a valid PDF.")

    input_pdf_stream.seek(0)
    reader = PdfReader(input_pdf_stream)
    total_pages = len(reader.pages)

    writer = PdfWriter()
    for page in reader.pages:
        writer.add_page(page)

    remainder = total_pages % 4
    blank_pages_needed = (4 - remainder) if remainder > 0 else 0

    for _ in range(blank_pages_needed):
        writer.add_blank_page()

    temp_pdf_stream = io.BytesIO()
    writer.write(temp_pdf_stream)
    temp_pdf_stream.seek(0)

    modified_reader = PdfReader(temp_pdf_stream)
    total_pages = len(modified_reader.pages)

    front_writer = PdfWriter()
    back_writer = PdfWriter()

    # Create front pages
    for i in range(total_pages // 4):
        page1 = modified_reader.pages[total_pages - 1 - (2 * i)]
        page2 = modified_reader.pages[2 * i]
        merged_page = merge_two_pages_on_one(page1, page2)
        front_writer.add_page(merged_page)

    # Create back pages
    for i in reversed(range(total_pages // 4)):
        page1 = modified_reader.pages[2 * i + 1]
        page2 = modified_reader.pages[total_pages - 2 - (2 * i)]
        merged_page = merge_two_pages_on_one(page1, page2)
        back_writer.add_page(merged_page)

    # Prepare output streams
    front_pdf_stream = io.BytesIO()
    back_pdf_stream = io.BytesIO()
    reversed_front_pdf_stream = io.BytesIO()
    reversed_back_pdf_stream = io.BytesIO()

    front_writer.write(front_pdf_stream)
    back_writer.write(back_pdf_stream)

    front_pdf_stream.seek(0)
    back_pdf_stream.seek(0)

    # Reverse PDFs
    reversed_front_pdf_stream = reverse_pdf(front_pdf_stream)
    reversed_back_pdf_stream = reverse_pdf(back_pdf_stream)

    return front_pdf_stream, back_pdf_stream, reversed_front_pdf_stream, reversed_back_pdf_stream

# Streamlit UI
# st.title("Booklet Generator: Create Print-Ready PDFs")
st.title("Print-Ready PDF Booklet Generator")
st.write("""
This tool simplifies the process of creating booklet-style PDFs. It generates booklet-ready files designed for back-to-back printing. Additionally, it provides multiple versions, including reversed formats, to accommodate all printer configurations seamlessly.
         
Ideal for printing books, modules, and other documents, this tool specifically helps save paper while ensuring efficient, professional-quality results.
""")

# Initialize session state
if "pdf_outputs" not in st.session_state:
    st.session_state["pdf_outputs"] = None

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

if uploaded_file is not None and st.button("Generate Booklet"):
    with st.spinner("Processing..."):
        try:
            front_pdf, back_pdf, reversed_front_pdf, reversed_back_pdf = create_booklet_with_two_pages_per_sheet(uploaded_file)
            st.session_state["pdf_outputs"] = {
                "front_pdf": front_pdf,
                "back_pdf": back_pdf,
                "reversed_front_pdf": reversed_front_pdf,
                "reversed_back_pdf": reversed_back_pdf,
            }
            st.success("Booklet generated successfully!")
        except Exception as e:
            st.error(f"Error: {str(e)}")

# Add a separator before downloads
st.divider()  # Adds a horizontal line

# Display download buttons if PDFs are available
if st.session_state["pdf_outputs"] is not None:
    outputs = st.session_state["pdf_outputs"]
    st.download_button("Download Booklet Front PDF", outputs["front_pdf"], "booklet_front.pdf", "application/pdf")
    st.download_button("Download Booklet Back PDF", outputs["back_pdf"], "booklet_back.pdf", "application/pdf")
    st.download_button("Download Reversed Booklet Front PDF", outputs["reversed_front_pdf"], "reversed_booklet_front.pdf", "application/pdf")
    st.download_button("Download Reversed Booklet Back PDF", outputs["reversed_back_pdf"], "reversed_booklet_back.pdf", "application/pdf")


# Add footer
st.markdown("""
---
<div style="text-align: center; font-size: small;">
    Developed by <strong>Aspra, N.</strong> | © 2025 All Rights Reserved
</div>
""", unsafe_allow_html=True)

# streamlit run 241009_PDF_booklet/ui_PDF_booklet.py