import os
import subprocess

# Temporarily force reinstall transformers for debugging
subprocess.run(["pip", "install", "--force-reinstall", "transformers==4.44.2"], check=True)

import gradio as gr
import shutil
from PIL import Image

from utils.pdf_to_images import convert_pdf_to_images, convert_image_to_pil
from utils.hunyuan_ocr import HunyuanOCR
from utils.pdf_builder import build_searchable_pdf

# Get Hugging Face token from environment variables (for gated models)
hf_token = os.environ.get("HF_TOKEN")

# Initialize OCR model
ocr_model = HunyuanOCR(token=hf_token)

def process_document(file_obj):
    if file_obj is None:
        return [], None, "Please upload a PDF or image file."

    input_path = file_obj.name
    file_extension = os.path.splitext(input_path)[1].lower()
    
    images = []
    if file_extension == ".pdf":
        try:
            images = convert_pdf_to_images(input_path)
        except Exception as e:
            return [], None, f"Error converting PDF to images: {e}"
    elif file_extension in [".png", ".jpg", ".jpeg", ".tiff", ".tif"]:
        try:
            images = [convert_image_to_pil(input_path)]
        except Exception as e:
            return [], None, f"Error loading image: {e}"
    else:
        return [], None, "Unsupported file type. Please upload a PDF or image."

    ocr_results = []
    preview_images = []
    for i, image in enumerate(images):
        if max(image.width, image.height) > 10000:
            return [], None, f"Page {i+1} is too large ({image.width}x{image.height}px). Max 10000px allowed."
        
        try:
            ocr_output = ocr_model.extract_ocr(image)
            ocr_results.append(ocr_output)
            preview_images.append(ocr_model.visualize_ocr(image.copy(), ocr_output))
        except Exception as e:
            return [], None, f"Error during OCR processing for page {i+1}: {e}"

    output_pdf_path = "searchable_output.pdf"
    try:
        build_searchable_pdf(images, ocr_results, output_pdf_path)
    except Exception as e:
        return [], None, f"Error building searchable PDF: {e}"

    return preview_images, output_pdf_path, "Searchable PDF generated successfully!"

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Searchable PDF Generator with Hunyuan-DiT")
    gr.Markdown("Upload a PDF or image file to convert it into a layout-preserving, searchable PDF.")

    with gr.Row():
        file_upload = gr.File(label="Upload PDF or Image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".tif"])
        convert_button = gr.Button("Convert to Searchable PDF")

    status_message = gr.Textbox(label="Status", interactive=False)
    
    ocr_preview_gallery = gr.Gallery(label="OCR Preview (Bounding Boxes)", columns=2, object_fit="contain", height="auto")
    output_pdf_file = gr.File(label="Download Searchable PDF", interactive=False)

    convert_button.click(
        process_document,
        inputs=[file_upload],
        outputs=[ocr_preview_gallery, output_pdf_file, status_message]
    )

if __name__ == "__main__":
    demo.launch()