from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
import gradio as gr
import pdf2image
from PIL import Image
import torch
import json

# Load model + processor with OCR enabled
processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=True, truncation=True)
model = LayoutLMv3ForTokenClassification.from_pretrained("microsoft/layoutlmv3-base")

def analyze_pdf_with_layoutlm(pdf_file):
    """
    Converts first page of PDF to image, lets LayoutLMv3 apply its OCR,
    and predicts token-level labels.
    """
    if pdf_file is None:
        return "Please upload a PDF file.", "{}"

    try:
        # Convert PDF to image (first page)
        images = pdf2image.convert_from_bytes(pdf_file, dpi=300)
        image = images[0]

        # Use LayoutLMv3 built-in OCR
        encoding = processor(image, return_tensors="pt")

        # ---- Defensive Checks ----
        if not hasattr(encoding, "words") or not callable(encoding.words):
            return "⚠️ This processor version doesn’t support .words(). Try updating transformers.", "{}"

        words = encoding.words()
        if not words or len(words) == 0:
            return "⚠️ No text detected on this page. Try a clearer PDF or higher DPI.", "{}"

        if "input_ids" not in encoding or encoding["input_ids"].shape[1] == 0:
            return "⚠️ Empty token sequence. OCR failed to detect text.", "{}"

        # ---- Run Model ----
        with torch.no_grad():
            outputs = model(**encoding)

        predictions = outputs.logits.argmax(-1).squeeze().tolist()
        labels = [model.config.id2label[p] for p in predictions]

        # Align safely
        length = min(len(words), len(labels))
        result_pairs = list(zip(words[:length], labels[:length]))

        # ---- Human-readable output ----
        formatted_output = "📄 **LayoutLMv3 Token Classification Results**\n\n"
        formatted_output += "Detected tokens and predicted entity labels:\n\n"
        for word, label in result_pairs[:30]:  # limit for readability
            formatted_output += f"- {word}: `{label}`\n"

        formatted_output += f"\nTotal tokens processed: {len(result_pairs)}"

        # ---- JSON formatted output ----
        json_output = json.dumps(
            [{"word": w, "label": l} for w, l in result_pairs],
            indent=2
        )

        return formatted_output, json_output

    except Exception as e:
        return f"❌ Error processing PDF: {str(e)}", "{}"

# --- Gradio Interface ---
with gr.Blocks(title="PDF LayoutLMv3 Analysis") as demo:
    gr.Markdown("# 🧠 PDF Analysis with LayoutLMv3 (Auto OCR)")
    gr.Markdown(
        "Upload a PDF and let LayoutLMv3 automatically extract text and layout information "
        "using its built-in OCR system."
    )

    with gr.Row():
        pdf_file = gr.File(label="Upload PDF File", file_types=["pdf"], type="binary")

    process_btn = gr.Button("Analyze PDF", variant="primary")

    with gr.Tabs():
        with gr.TabItem("📋 Model Results"):
            formatted_output = gr.Textbox(
                label="Model Output",
                lines=20,
                interactive=False,
                show_copy_button=True
            )

        with gr.TabItem("📄 JSON Data"):
            json_output = gr.Textbox(
                label="JSON Output",
                lines=20,
                interactive=False,
                show_copy_button=True
            )

    process_btn.click(
        fn=analyze_pdf_with_layoutlm,
        inputs=[pdf_file],
        outputs=[formatted_output, json_output]
    )

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