import streamlit as st
from PIL import Image
from fpdf import FPDF
import os

# Function to process a single image with user-provided sequence number
def process_image(image_file, sequence_number):
    try:
        image = Image.open(image_file)
        return image, sequence_number
    except Exception as e:
        st.error(f"Failed to open image: {e}")
        return None, None

# Function to generate the PDF using fpdf
def generate_pdf(images_with_sequence):
    pdf_name = "output.pdf"

    images_with_sequence.sort(key=lambda x: x[1])  # Sort by sequence number

    pdf = FPDF()
    for i, (image, _) in enumerate(images_with_sequence):
        if image.mode != 'RGB':
            image = image.convert('RGB')
        image.save(f"temp_image_{i}.jpg", "JPEG")
        pdf.add_page()
        pdf.image(f"temp_image_{i}.jpg", x=10, y=10, w=190)
    
    try:
        pdf.output(pdf_name)
    except Exception as e:
        st.error(f"Failed to generate PDF: {e}")
        return None
    return pdf_name

st.markdown("<h1 style='text-align: center; color: #00698f;'>Image to PDF Converter</h1>", unsafe_allow_html=True)
st.markdown("<h3 style='text-align: center; color: #0086b3;'>Supports multiple image upload at once and rearrangement of pages</h3>", unsafe_allow_html=True)

uploaded_images = st.file_uploader("Upload Images", type=["jpg", "jpeg", "png"], accept_multiple_files=True)

if uploaded_images:
    num_images = len(uploaded_images)
    if num_images > 0:  # Check if any image is uploaded
        st.write("**Uploaded Images:**")
        images_with_sequence = []

        for i, image_file in enumerate(uploaded_images):
            col1, col2 = st.columns([3, 1])
            with col1:
                st.image(image_file, width=250)
            with col2:
                sequence_options = [str(j) for j in range(1, num_images + 1)]
                default_index = i  # Set default selection of sequence as the order of upload
                sequence_number = st.selectbox(f"Sequence {i+1}", sequence_options, index=default_index)
                processed_image, processed_sequence = process_image(image_file, int(sequence_number))
                if processed_image:
                    images_with_sequence.append((processed_image, processed_sequence))

        if st.button("Generate PDF"):
            pdf_name = generate_pdf(images_with_sequence)
            if pdf_name:  # Check if PDF is generated successfully
                st.success(f"PDF generated! Download: {pdf_name}")
                with open(pdf_name, "rb") as pdf_file:
                    st.download_button("Download PDF", pdf_file, file_name=pdf_name)
                for i in range(len(images_with_sequence)):
                    if os.path.exists(f"temp_image_{i}.jpg"):
                        os.remove(f"temp_image_{i}.jpg")
                os.remove(pdf_name)  # Remove the file after download
else:
    st.write("Please upload images to continue")