import streamlit as st
from pypdf import PdfWriter, PdfReader
from PIL import Image
import io

# --- PAGE CONFIGURATION ---
st.set_page_config(
    page_title="PDF & Image Toolkit",
    page_icon="📄",
    layout="centered",
)

# --- HELPER FUNCTIONS ---
def merge_pdfs(uploaded_files):
    """Merges multiple PDF files into one."""
    merger = PdfWriter()
    for pdf_file in uploaded_files:
        # pypdf needs a file object, so we wrap the uploaded bytes in BytesIO
        merger.append(io.BytesIO(pdf_file.getvalue()))
    
    # Write the merged PDF to a BytesIO object in memory
    output_buffer = io.BytesIO()
    merger.write(output_buffer)
    merger.close()
    return output_buffer.getvalue()

def compress_pdf(uploaded_file):
    """Performs basic compression on a PDF file."""
    reader = PdfReader(io.BytesIO(uploaded_file.getvalue()))
    writer = PdfWriter()

    for page in reader.pages:
        # This is a basic form of compression
        page.compress_content_streams()
        writer.add_page(page)

    # Write the compressed PDF to a BytesIO object
    output_buffer = io.BytesIO()
    writer.write(output_buffer)
    writer.close()
    return output_buffer.getvalue()

def images_to_pdf(uploaded_files):
    """Converts a list of images to a single PDF."""
    images = []
    for img_file in uploaded_files:
        img = Image.open(io.BytesIO(img_file.getvalue()))
        # Convert to RGB to avoid issues with different image modes (e.g., RGBA, P)
        if img.mode == 'RGBA' or img.mode == 'P':
            img = img.convert('RGB')
        images.append(img)
    
    if not images:
        return None

    # Save the images as a PDF to a BytesIO object
    output_buffer = io.BytesIO()
    images[0].save(
        output_buffer, 
        "PDF" ,
        resolution=100.0, 
        save_all=True, 
        append_images=images[1:]
    )
    return output_buffer.getvalue()


# --- STREAMLIT UI ---

st.title("📄 PDF & Image Toolkit")
st.write("A simple web app to manage your documents. All processing is done in your browser and on the server, and your files are not saved.")

# --- SIDEBAR FOR TOOL SELECTION ---
with st.sidebar:
    st.header("Tools")
    tool_choice = st.radio(
        "Select a tool:",
        ("Merge PDFs", "Compress PDF", "Convert Images to PDF")
    )

# --- MAIN PAGE CONTENT ---

if tool_choice == "Merge PDFs":
    st.subheader("Merge Multiple PDF Files")
    st.write("Upload two or more PDF files to combine them into a single document.")
    
    uploaded_pdfs = st.file_uploader(
        "Upload PDF files", 
        type="pdf", 
        accept_multiple_files=True,
        key="pdf_merger"
    )

    if st.button("Merge PDFs") and uploaded_pdfs:
        if len(uploaded_pdfs) < 2:
            st.warning("Please upload at least two PDF files to merge.")
        else:
            with st.spinner("Merging..."):
                merged_pdf_bytes = merge_pdfs(uploaded_pdfs)
                st.success("PDFs merged successfully!")
                st.download_button(
                    label="Download Merged PDF",
                    data=merged_pdf_bytes,
                    file_name="merged_document.pdf",
                    mime="application/pdf",
                )

elif tool_choice == "Compress PDF":
    st.subheader("Compress a PDF File")
    st.info("This tool performs basic optimization. It may not significantly reduce the size of all PDFs, especially those that are already optimized or text-heavy.")

    uploaded_pdf = st.file_uploader(
        "Upload a PDF file", 
        type="pdf", 
        accept_multiple_files=False,
        key="pdf_compressor"
    )

    if st.button("Compress PDF") and uploaded_pdf:
        with st.spinner("Compressing..."):
            compressed_pdf_bytes = compress_pdf(uploaded_pdf)
            original_size = uploaded_pdf.size
            compressed_size = len(compressed_pdf_bytes)
            reduction = ((original_size - compressed_size) / original_size) * 100
            
            st.success(f"Compression complete! Size reduced by {reduction:.2f}%.")
            st.download_button(
                label="Download Compressed PDF",
                data=compressed_pdf_bytes,
                file_name="compressed_document.pdf",
                mime="application/pdf",
            )


elif tool_choice == "Convert Images to PDF":
    st.subheader("Convert Images to a Single PDF")
    st.write("Upload one or more images (JPG, PNG) to combine them into a PDF.")

    uploaded_images = st.file_uploader(
        "Upload image files",
        type=["png", "jpg", "jpeg"],
        accept_multiple_files=True,
        key="img_converter"
    )

    if st.button("Convert to PDF") and uploaded_images:
        with st.spinner("Converting..."):
            pdf_from_images_bytes = images_to_pdf(uploaded_images)
            st.success("Images converted to PDF successfully!")
            st.download_button(
                label="Download PDF",
                data=pdf_from_images_bytes,
                file_name="images_converted.pdf",
                mime="application/pdf",
            )

st.sidebar.markdown("---")
st.sidebar.info("App built with [Streamlit](https://streamlit.io) and hosted on [Hugging Face Spaces](https://huggingface.co/spaces).")