import os
import random
import shutil
import tempfile
import zipfile

import gradio as gr
from huggingface_hub import HfApi
from pdf2image import convert_from_path
from PyPDF2 import PdfReader


def pdf_to_images(pdf_files, sample_size, temp_dir, progress=gr.Progress()):
    if not os.path.exists(temp_dir):
        os.makedirs(temp_dir)
    progress(0, desc="Starting conversion")
    all_images = []
    for pdf_file in progress.tqdm(pdf_files, desc="Converting PDFs"):
        pdf_path = pdf_file.name
        pdf = PdfReader(pdf_path)
        total_pages = len(pdf.pages)

        # Determine the number of pages to convert
        pages_to_convert = (
            total_pages if sample_size == 0 else min(sample_size, total_pages)
        )

        # Select random pages if sampling
        if sample_size > 0 and sample_size < total_pages:
            selected_pages = sorted(
                random.sample(range(1, total_pages + 1), pages_to_convert)
            )
        else:
            selected_pages = range(1, total_pages + 1)

        # Convert selected PDF pages to images
        for page_num in selected_pages:
            images = convert_from_path(
                pdf_path, first_page=page_num, last_page=page_num
            )
            for image in images:
                image_path = os.path.join(
                    temp_dir, f"{os.path.basename(pdf_path)}_page_{page_num}.jpg"
                )
                image.save(image_path, "JPEG")
                all_images.append(image_path)

    return all_images, f"Saved {len(all_images)} images to temporary directory"


def process_pdfs(
    pdf_files,
    sample_size,
    hf_repo,
    oauth_token: gr.OAuthToken | None,
    progress=gr.Progress(),
):
    if not pdf_files:
        return (
            None,
            None,
            gr.Markdown(
                "⚠️ No PDF files uploaded. Please upload at least one PDF file."
            ),
        )

    if oauth_token is None:
        return (
            None,
            None,
            gr.Markdown(
                "⚠️ Not logged in to Hugging Face. Please log in to upload to a Hugging Face dataset."
            ),
        )

    try:
        temp_dir = tempfile.mkdtemp()
        images_dir = os.path.join(temp_dir, "images")
        os.makedirs(images_dir)

        progress(0, desc="Starting PDF processing")
        images, message = pdf_to_images(pdf_files, sample_size, images_dir)

        # Create a zip file of the images
        zip_path = os.path.join(temp_dir, "converted_images.zip")
        with zipfile.ZipFile(zip_path, "w") as zipf:
            progress(0, desc="Zipping images")
            for image in progress.tqdm(images, desc="Zipping images"):
                zipf.write(image, os.path.basename(image))

        if hf_repo:
            try:
                hf_api = HfApi(token=oauth_token.token)
                hf_api.create_repo(
                    hf_repo,
                    repo_type="dataset",
                )
                hf_api.upload_folder(
                    folder_path=images_dir,
                    repo_id=hf_repo,
                    repo_type="dataset",
                    path_in_repo="images",
                )
                message += f"\nUploaded images to Hugging Face repo: {hf_repo}/images"
            except Exception as e:
                message += f"\nFailed to upload to Hugging Face: {str(e)}"

        return images, zip_path, message
    except Exception as e:
        if "temp_dir" in locals():
            shutil.rmtree(temp_dir)
        return None, None, f"An error occurred: {str(e)}"


# Define the Gradio interface
with gr.Blocks() as demo:
    gr.HTML(
        """<h1 style='text-align: center;'> PDFs to Page Images Converter</h1>
        <center><i> &#128193; Convert PDFs to an image dataset &#128193; </i></center>"""
    )
    gr.Markdown(
        "Upload PDF(s), convert pages to images, and optionally upload them to a Hugging Face repo. If a sample size is specified, random pages will be selected."
    )

    with gr.Row():
        gr.LoginButton(size="sm")

    with gr.Row():
        pdf_files = gr.File(
            file_count="multiple", label="Upload PDF(s)", file_types=["*.pdf"]
        )
    with gr.Row():
        sample_size = gr.Number(
            value=None,
            label="Pages per PDF (0 for all pages)",
            info="Specify how many pages to convert from each PDF. Use 0 to convert all pages.",
        )
        hf_repo = gr.Textbox(
            label="Hugging Face Repo",
            placeholder="username/repo-name",
            info="Enter the Hugging Face repository name in the format 'username/repo-name'",
        )
    with gr.Accordion("View converted images", open=False):
        output_gallery = gr.Gallery(label="Converted Images")
    status_text = gr.Markdown(label="Status")
    download_button = gr.File(label="Download Converted Images")

    submit_button = gr.Button("Convert PDFs to page images")
    submit_button.click(
        process_pdfs,
        inputs=[pdf_files, sample_size, hf_repo],
        outputs=[output_gallery, download_button, status_text],
    )

# Launch the app
demo.launch(debug=True)