import gradio as gr
import fitz
from PIL import Image
import io

def pdf_page_to_image(pdf_path, page_number):
    # Open the PDF file
    doc = fitz.open(pdf_path)
    page_number -= 1

    # Check if the page number is valid
    if page_number < 0 or page_number >= len(doc):
        raise gr.Error("Page number out of range")

    # Select the specified page
    page = doc.load_page(page_number)

    # Render the page to an image
    pix = page.get_pixmap()

    # Convert the image to a PIL Image
    img_data = pix.tobytes("png")  # Get the PNG image data
    img = Image.open(io.BytesIO(img_data))

    # Close the PDF document
    doc.close()

    return img

def get_pdf_length(pdf_path):
    return len(fitz.open(pdf_path))

def update_number(pdf_path):
    num_pages = get_pdf_length(pdf_path)
    return gr.Number(info=f"Enter a number between 1 and {num_pages}", maximum=num_pages)

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            f = gr.File(label="Upload files", file_types=[".pdf"])
            n = gr.Number(1, minimum=1, label="Page number")
            b = gr.Button("Submit", variant="primary")
        
        with gr.Column():
            i = gr.Image()
            o = gr.Textbox(label="Prediction")
            e = gr.Textbox(label="Edited prediction")
        
    f.upload(update_number, f, n)
    n.submit(pdf_page_to_image, [f, n], i)
    b.click(pdf_page_to_image, [f, n], i)
    
demo.launch()