import gradio as gr
from PIL import Image
import tempfile
import os

def combine_two_images_to_pdf(file1, file2):
    """
    Takes two specific image files, converts them to RGB, 
    and saves them as a single PDF.
    """
    # Create a list of the files that were actually uploaded (ignore empty ones)
    image_paths = [f for f in [file1, file2] if f is not None]

    if not image_paths:
        return None

    pil_images = []
    
    try:
        # Loop through the uploaded files
        for file_path in image_paths:
            try:
                img = Image.open(file_path)
                # Convert RGBA (transparent) to RGB to ensure PDF compatibility
                if img.mode == 'RGBA':
                    img = img.convert('RGB')
                pil_images.append(img)
            except Exception as e:
                print(f"Error processing file {file_path}: {e}")
                continue

        if not pil_images:
            return None

        # Create a temporary file for the PDF output
        temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
        output_path = temp_pdf.name
        
        # Save the first image and append the rest (if any)
        if len(pil_images) > 1:
            pil_images[0].save(
                output_path, 
                save_all=True, 
                append_images=pil_images[1:]
            )
        else:
            # If only one image was uploaded, just save that one
            pil_images[0].save(output_path)
        
        return output_path

    except Exception as e:
        print(f"Error creating PDF: {e}")
        return None

# Define the Gradio Interface
with gr.Blocks(title="2 Images to PDF") as app:
    gr.Markdown("# 📄 2 Images to PDF")
    gr.Markdown("Upload two specific images below to merge them into one PDF.")
    
    with gr.Row():
        with gr.Column():
            # Slot 1
            img_input1 = gr.File(
                label="Upload Image 1",
                file_count="single",
                file_types=["image"]
            )
            # Slot 2
            img_input2 = gr.File(
                label="Upload Image 2",
                file_count="single",
                file_types=["image"]
            )
            convert_btn = gr.Button("Convert to PDF", variant="primary")
        
        with gr.Column():
            pdf_output = gr.File(label="Download PDF")

    # Connect both inputs to the function
    convert_btn.click(
        fn=combine_two_images_to_pdf, 
        inputs=[img_input1, img_input2], 
        outputs=pdf_output
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0")