app.py

import gradio as gr
from docling.document_converter import DocumentConverter
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions, TableFormerMode, RapidOcrOptions
from docling.document_converter import PdfFormatOption
from PIL import Image, ImageDraw, ImageFont
import json
import fitz  # PyMuPDF
import os
from dotenv import load_dotenv
import io
import numpy as np
import cv2
from typing import List, Tuple, Optional
from concurrent.futures import ThreadPoolExecutor
import threading

# Set CPU thread counts for better multi-core utilization
os.environ['OMP_NUM_THREADS'] = '2'
os.environ['OPENBLAS_NUM_THREADS'] = '2'
os.environ['MKL_NUM_THREADS'] = '2'
os.environ['NUMEXPR_NUM_THREADS'] = '2'

# Optional imports for signature detection
try:
    import supervision as sv
    from ultralytics import YOLO
    from huggingface_hub import hf_hub_download
    import onnxruntime as ort
except Exception:
    sv = None
    YOLO = None
    hf_hub_download = None
    ort = None

# Color mapping for different layout elements
COLORS = {
    "title": "#FF6B6B",
    "text": "#4ECDC4",
    "section_header": "#95E1D3",
    "table": "#F38181",
    "list": "#AA96DA",
    "figure": "#FCBAD3",
    "caption": "#A8D8EA",
    "formula": "#FFD93D",
    "footnote": "#6BCB77",
    "page_header": "#4D96FF",
    "page_footer": "#9D84B7",
    "picture": "#FF8C42",
    # Picture classifications
    "signature": "#9D4EDD",
    "qr_code": "#06FFA5",
    "bar_code": "#06FFA5",
    "logo": "#FFB627",
    "stamp": "#E63946",
    "icon": "#F4A261",
    "bar_chart": "#2A9D8F",
    "pie_chart": "#E76F51",
    "line_chart": "#264653",
    "flow_chart": "#8338EC",
    "map": "#3A86FF",
    "screenshot": "#FB5607",
    "other": "#CCCCCC",
}

# Load environment variables from .env if present (useful for HF_TOKEN)
try:
    load_dotenv()
except Exception:
    pass

# ------------- Signature Model Utilities -------------
_SIGNATURE_MODEL = None
_ONNX_SESSION = None


def load_signature_model() -> Optional["YOLO"]:
    """Load and cache the YOLOv8s signature model (ONNX format with OpenVINO).

    Returns None if dependencies are missing.
    """
    global _SIGNATURE_MODEL, _ONNX_SESSION
    if _SIGNATURE_MODEL is not None and _ONNX_SESSION is not None:
        return _SIGNATURE_MODEL
    if YOLO is None or hf_hub_download is None or ort is None:
        return None
    try:
        # Download ONNX model from Hugging Face
        onnx_path = hf_hub_download(
            repo_id="tech4humans/yolov8s-signature-detector",
            filename="yolov8s.onnx",
            token=os.environ.get("HF_TOKEN")
        )
        
        # Create ONNX Runtime session with OpenVINO execution provider
        # OpenVINO provides significant speedup on Intel CPUs
        providers = []
        
        # Try OpenVINO first (best for Intel CPUs)
        if 'OpenVINOExecutionProvider' in ort.get_available_providers():
            providers.append('OpenVINOExecutionProvider')
            print("✓ Using OpenVINO Execution Provider for ONNX Runtime")
        
        # Fallback to CPU provider
        providers.append('CPUExecutionProvider')
        
        # Configure session options for performance
        sess_options = ort.SessionOptions()
        sess_options.intra_op_num_threads = 2  # Use both CPU cores
        sess_options.inter_op_num_threads = 2
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        
        _ONNX_SESSION = ort.InferenceSession(
            onnx_path,
            sess_options=sess_options,
            providers=providers
        )
        
        print(f"✓ ONNX Runtime providers: {_ONNX_SESSION.get_providers()}")
        
        # Still load YOLO object for utility functions (but won't use for inference)
        pt_path = hf_hub_download(
            repo_id="tech4humans/yolov8s-signature-detector",
            filename="yolov8s.pt",
            token=os.environ.get("HF_TOKEN")
        )
        _SIGNATURE_MODEL = YOLO(pt_path)
        
        return _SIGNATURE_MODEL
    except Exception as e:
        print(f"Could not load signature model: {e}")
        return None


def yolo_detect_signatures(
    image_bgr: np.ndarray,
    imgsz: int = 640,  # Changed from 1280 to match training size (640x640)
    conf: float = 0.05,
    iou: float = 0.45,
    augment: bool = False,  # ONNX doesn't support augment
) -> List[Tuple[np.ndarray, float, int]]:
    """Run YOLO signature detection on a BGR image using ONNX Runtime.

    Returns list of (xyxy np.array[4], score float, class_idx int)
    """
    global _ONNX_SESSION
    model = load_signature_model()
    if model is None or _ONNX_SESSION is None:
        return []
    try:
        # Preprocess image for ONNX inference
        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        original_shape = image_rgb.shape[:2]  # (height, width)
        
        # Resize to model input size (640x640)
        img_resized = cv2.resize(image_rgb, (imgsz, imgsz))
        
        # Normalize and transpose to NCHW format
        img_normalized = img_resized.astype(np.float32) / 255.0
        img_transposed = np.transpose(img_normalized, (2, 0, 1))  # HWC to CHW
        img_batch = np.expand_dims(img_transposed, axis=0)  # Add batch dimension
        
        # Run ONNX inference
        input_name = _ONNX_SESSION.get_inputs()[0].name
        outputs = _ONNX_SESSION.run(None, {input_name: img_batch})
        
        # Post-process ONNX outputs (YOLOv8 format)
        # Output shape: [1, num_detections, 84] where 84 = 4 bbox coords + 80 class scores
        predictions = outputs[0][0]  # Remove batch dimension
        
        # Extract boxes and scores
        boxes = []
        for pred in predictions.T:  # Transpose to [num_detections, 84]
            # pred format: [cx, cy, w, h, class_scores...]
            if len(pred) < 5:
                continue
                
            # Get bbox coordinates (first 4 values)
            cx, cy, w, h = pred[:4]
            
            # Get max class score and index
            class_scores = pred[4:]
            max_score = np.max(class_scores)
            
            if max_score < conf:
                continue
                
            class_idx = np.argmax(class_scores)
            
            # Convert from center format to corner format
            x1 = (cx - w / 2) / imgsz * original_shape[1]
            y1 = (cy - h / 2) / imgsz * original_shape[0]
            x2 = (cx + w / 2) / imgsz * original_shape[1]
            y2 = (cy + h / 2) / imgsz * original_shape[0]
            
            boxes.append((np.array([x1, y1, x2, y2]), float(max_score), int(class_idx)))
        
        # Apply NMS
        if boxes:
            boxes = _apply_nms_to_detections(boxes, iou)
        
        return boxes
    except Exception as e:
        print(f"ONNX signature detection error: {e}")
        # Fallback to PyTorch if ONNX fails
        try:
            results = model(image_bgr, imgsz=imgsz, conf=conf, iou=iou, augment=False)
            r = results[0]
            boxes = []
            if hasattr(r, "boxes") and r.boxes is not None:
                xyxy = r.boxes.xyxy.cpu().numpy()
                scores = r.boxes.conf.cpu().numpy()
                classes = r.boxes.cls.cpu().numpy().astype(int)
                for b, s, c in zip(xyxy, scores, classes):
                    boxes.append((b, float(s), int(c)))
            return boxes
        except Exception as fallback_error:
            print(f"PyTorch fallback also failed: {fallback_error}")
            return []


def annotate_signature_boxes_on_pil(img_pil: Image.Image, boxes: List[Tuple[np.ndarray, float, int]]) -> Image.Image:
    """Draw signature boxes on a PIL image and return annotated copy."""
    if not boxes:
        return img_pil
    img = img_pil.copy()
    draw = ImageDraw.Draw(img)
    # Try fonts
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 16)
    except Exception:
        font = ImageFont.load_default()
    color = COLORS.get("signature", "#9D4EDD")
    for (xyxy, score, cls) in boxes:
        x1, y1, x2, y2 = map(int, xyxy)
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
        label = f"Signature {score*100:.0f}%"
        bbox_text = draw.textbbox((x1, y1 - 22), label, font=font)
        draw.rectangle([bbox_text[0] - 2, bbox_text[1] - 2, bbox_text[2] + 2, bbox_text[3] + 2], fill=color)
        draw.text((x1, y1 - 22), label, fill="white", font=font)
    return img

def draw_layout_boxes(image_path, layout_data, scale_x=1.0, scale_y=1.0):
    """Draw bounding boxes on the image based on layout predictions"""
    # Open the image
    if isinstance(image_path, str):
        img = Image.open(image_path).convert("RGB")
    else:
        img = image_path.convert("RGB")
    
    draw = ImageDraw.Draw(img)
    
    # Try to load a font, fallback to default if not available
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
        small_font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14)
    except:
        font = ImageFont.load_default()
        small_font = ImageFont.load_default()
    
    # Draw each cluster
    for cluster in layout_data:
        label = cluster.get("label", "unknown")
        bbox = cluster.get("bbox")
        classification = cluster.get("classification")
        
        if bbox:
            # bbox format: [x0, y0, x1, y1] from PDF coordinates
            # Scale to match rendered image dimensions
            x0, y0, x1, y1 = bbox
            x0 = x0 * scale_x
            y0 = y0 * scale_y
            x1 = x1 * scale_x
            y1 = y1 * scale_y
            
            # Get color for this label
            color = COLORS.get(label, "#999999")
            
            # Draw rectangle
            draw.rectangle([x0, y0, x1, y1], outline=color, width=3)
            
            # Draw label with classification confidence if available
            if classification:
                confidence_pct = classification['confidence'] * 100
                label_text = f"{label.replace('_', ' ').title()} ({confidence_pct:.0f}%)"
            else:
                label_text = label.replace("_", " ").title()
            
            bbox_text = draw.textbbox((x0, y0 - 25), label_text, font=small_font)
            draw.rectangle([bbox_text[0] - 2, bbox_text[1] - 2, bbox_text[2] + 2, bbox_text[3] + 2], 
                         fill=color)
            
            # Draw label text
            draw.text((x0, y0 - 25), label_text, fill="white", font=small_font)
    
    return img

def process_document(file_path, mode, enable_ocr, enable_tables, run_signature_yolo=False, signature_conf=0.05):
    """Process document with Docling and return results"""
    try:
        # Configure pipeline options
        pipeline_options = PdfPipelineOptions()
        pipeline_options.do_table_structure = enable_tables
        
        if enable_tables:
            if mode == "Accurate":
                pipeline_options.table_structure_options.mode = TableFormerMode.ACCURATE
            else:
                pipeline_options.table_structure_options.mode = TableFormerMode.FAST
        
        pipeline_options.do_ocr = enable_ocr
        if enable_ocr:
            # Force RapidOCR with ONNX backend for fast & accurate CPU inference
            pipeline_options.ocr_options = RapidOcrOptions(
                backend="onnxruntime",
                force_full_page_ocr=True,
            )
        pipeline_options.generate_page_images = True
        pipeline_options.generate_picture_images = True
        pipeline_options.do_picture_classification = True  # Enable classification
        pipeline_options.images_scale = 3.0  # Higher resolution for better accuracy
        
        # Create converter
        converter = DocumentConverter(
            format_options={
                InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
                InputFormat.IMAGE: PdfFormatOption(pipeline_options=pipeline_options),
            }
        )
        
        # Convert document
        result = converter.convert(file_path)
        
        # Extract layout information
        layout_info = []
        total_clusters = 0
        table_count = 0
        
        # Get picture classifications for enrichment
        # We need to store by page number and use a more flexible matching
        picture_classifications_by_page = {}
        print(f"DEBUG: Total pictures found: {len(result.document.pictures)}")
        for picture in result.document.pictures:
            page_num = picture.prov[0].page_no
            bbox = picture.prov[0].bbox
            
            if page_num not in picture_classifications_by_page:
                picture_classifications_by_page[page_num] = []
            
            # Get classification if available
            for annotation in picture.annotations:
                if hasattr(annotation, 'predicted_classes') and annotation.predicted_classes:
                    top_pred = annotation.predicted_classes[0]
                    picture_classifications_by_page[page_num].append({
                        'bbox': bbox,
                        'class': top_pred.class_name,
                        'confidence': top_pred.confidence
                    })
                    print(f"DEBUG: Found classification - page: {page_num}, bbox: ({bbox.l:.2f}, {bbox.t:.2f}, {bbox.r:.2f}, {bbox.b:.2f}), class: {top_pred.class_name}")
                    break
        
        for page_no, page in enumerate(result.pages, 1):
            if page.predictions.layout:
                clusters = page.predictions.layout.clusters
                total_clusters += len(clusters)
                
                for cluster in clusters:
                    # Check if this is a picture with classification
                    label = cluster.label
                    classification = None
                    if cluster.label == "picture" and page_no in picture_classifications_by_page:
                        print(f"DEBUG: Picture cluster at page {page_no}: ({cluster.bbox.l:.2f}, {cluster.bbox.t:.2f}, {cluster.bbox.r:.2f}, {cluster.bbox.b:.2f})")
                        
                        # Find matching classification by comparing bounding boxes with tolerance
                        for pic_class in picture_classifications_by_page[page_no]:
                            pic_bbox = pic_class['bbox']
                            # Check if bboxes match with small tolerance (allowing for floating point differences)
                            # Compare left and right which should match exactly
                            if (abs(cluster.bbox.l - pic_bbox.l) < 1.0 and 
                                abs(cluster.bbox.r - pic_bbox.r) < 1.0):
                                # X coordinates match, this is likely the same picture
                                classification = {
                                    'class': pic_class['class'],
                                    'confidence': pic_class['confidence']
                                }
                                label = f"{classification['class']}"
                                print(f"DEBUG: Matched classification: {label} (conf: {classification['confidence']:.2%})")
                                break
                        
                        if not classification:
                            print(f"DEBUG: No classification match found")
                    
                    layout_info.append({
                        "page": page_no,
                        "label": label,
                        "bbox": [cluster.bbox.l, cluster.bbox.t, cluster.bbox.r, cluster.bbox.b],
                        "confidence": getattr(cluster, "confidence", None),
                        "classification": classification
                    })
            
            # Count tables
            if page.predictions.tablestructure and page.predictions.tablestructure.table_map:
                table_count += len(page.predictions.tablestructure.table_map)
        
        # Get markdown output
        markdown_output = result.document.export_to_markdown()
        
        # Create visualization for first page
        visualization = None
        first_page_base_image = None  # PIL image in pixel space used for overlays
        if result.pages and layout_info:
            # Draw boxes on first page only
            first_page_layout = [item for item in layout_info if item["page"] == 1]
            
            try:
                # Check if input is an image or PDF
                file_ext = file_path.lower().split('.')[-1]
                
                if file_ext in ['jpg', 'jpeg', 'png', 'tiff', 'bmp']:
                    # For images: Open directly, coordinates should match 1:1
                    first_page_image = Image.open(file_path).convert("RGB")
                    # No scaling needed for images - coordinates are already in pixels
                    first_page_base_image = first_page_image
                    visualization = draw_layout_boxes(first_page_image, first_page_layout, 
                                                     scale_x=1.0, scale_y=1.0)
                else:
                    # For PDFs: Render and calculate scale
                    doc = fitz.open(file_path)
                    page = doc[0]
                    
                    # Get page dimensions in PDF points
                    page_rect = page.rect
                    pdf_width = page_rect.width
                    pdf_height = page_rect.height
                    
                    # Render at 2x for better quality
                    zoom = 2.0
                    mat = fitz.Matrix(zoom, zoom)
                    pix = page.get_pixmap(matrix=mat)
                    first_page_image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    
                    # Calculate scale: rendered_pixels / pdf_points
                    scale_x = pix.width / pdf_width
                    scale_y = pix.height / pdf_height
                    
                    doc.close()
                    
                    first_page_base_image = first_page_image
                    # Draw boxes with calculated scale
                    visualization = draw_layout_boxes(first_page_image, first_page_layout, 
                                                     scale_x=scale_x, scale_y=scale_y)
            except Exception as e:
                print(f"Could not create visualization: {e}")
                import traceback
                traceback.print_exc()

        # Optionally run YOLO signature detection on the same first-page image and overlay
        if run_signature_yolo and first_page_base_image is not None:
            try:
                # Convert PIL RGB to BGR numpy for YOLO
                img_bgr = cv2.cvtColor(np.array(first_page_base_image), cv2.COLOR_RGB2BGR)
                sig_boxes = yolo_detect_signatures(
                    img_bgr,
                    imgsz=640,  # Changed to match training size for optimal performance
                    conf=float(signature_conf),
                    iou=0.45,
                    augment=False,  # ONNX doesn't support augment
                )
                if sig_boxes:
                    # Overlay signature boxes on top of visualization
                    base_for_overlay = visualization if visualization is not None else first_page_base_image
                    visualization = annotate_signature_boxes_on_pil(base_for_overlay, sig_boxes)
            except Exception as e:
                print(f"Signature overlay failed: {e}")
        
        # Create summary
        summary = f"""## Document Analysis Summary

📄 **Total Pages:** {len(result.document.pages)}
🏷️ **Layout Elements Detected:** {total_clusters}
📊 **Tables Found:** {table_count}

### Layout Elements by Type:
"""
        # Count elements by type
        element_counts = {}
        for item in layout_info:
            label = item["label"]
            element_counts[label] = element_counts.get(label, 0) + 1
        
        for label, count in sorted(element_counts.items()):
            summary += f"- **{label.replace('_', ' ').title()}**: {count}\n"
        
        # JSON output
        json_output = json.dumps(layout_info, indent=2)
        
        return visualization, summary, markdown_output, json_output
        
    except Exception as e:
        error_msg = f"Error processing document: {str(e)}"
        return None, error_msg, error_msg, error_msg

def gradio_interface(file, mode, enable_ocr, enable_tables, run_signature_yolo=False, signature_conf=0.05):
    """Gradio interface function"""
    if file is None:
        return None, "Please upload a document", "", ""
    
    # Get file path - handle both direct path and gr.File object
    try:
        if hasattr(file, 'name'):
            file_path = file.name
        else:
            file_path = str(file)
        
        # Validate file exists and has valid extension
        if not os.path.exists(file_path):
            return None, f"File not found: {file_path}", "", ""
        
        ext = os.path.splitext(file_path)[1].lower()
        valid_exts = [".pdf", ".jpg", ".jpeg", ".png", ".tiff", ".bmp"]
        if ext not in valid_exts:
            return None, f"Invalid file format: {ext}. Supported: {', '.join(valid_exts)}", "", ""
        
        return process_document(file_path, mode, enable_ocr, enable_tables, run_signature_yolo, signature_conf)
    except Exception as e:
        error_msg = f"Error in gradio_interface: {str(e)}"
        return None, error_msg, error_msg, error_msg


# -------- Small preview helper (first page / image) --------
def preview_first_page(file: gr.File):
    """Return filepath for preview. For PDFs, extract first page as temp image."""
    if file is None:
        return None
    try:
        path = file.name
        ext = (os.path.splitext(path)[1] or "").lower()
        if ext in (".pdf",):
            # For PDF, render first page to temp image
            import tempfile
            doc = fitz.open(path)
            page = doc[0]
            pix = page.get_pixmap(matrix=fitz.Matrix(1.5, 1.5))
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            doc.close()
            # Use delete=True and return immediately - Gradio will handle the file
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
                img.save(tmp.name)
                return tmp.name
        else:
            # For images, return path directly
            return path
    except Exception:
        return None


def analyze_with_preview(file, mode, enable_ocr, enable_tables, run_signature_yolo=False, signature_conf=0.05):
    """Wrapper to also return an input preview for Examples clicks."""
    preview = preview_first_page(file)
    vis, summ, md, js = gradio_interface(file, mode, enable_ocr, enable_tables, run_signature_yolo, signature_conf)
    return preview, vis, summ, md, js


def signature_only_with_preview(file, try_scales, conf, iou, augment):
    """Wrapper to also return an input preview for Examples clicks."""
    preview = preview_first_page(file)
    img, summ, js = signature_only_infer(file, try_scales, conf, iou, augment)
    return preview, img, summ, js

# -------- Signature-only utilities (full-image, no ROI) --------
def _apply_nms_to_detections(boxes, iou_threshold=0.5):
    """Apply Non-Maximum Suppression to remove duplicate detections.
    
    Used for ONNX post-processing.
    """
    if not boxes:
        return []
    
    # Convert to format for NMS: [x1, y1, x2, y2, score]
    boxes_array = np.array([[b[0][0], b[0][1], b[0][2], b[0][3], b[1]] for b in boxes])
    
    # Sort by score
    indices = np.argsort(boxes_array[:, 4])[::-1]
    keep = []
    
    while len(indices) > 0:
        # Pick the box with highest score
        current = indices[0]
        keep.append(current)
        
        if len(indices) == 1:
            break
            
        # Compute IoU with remaining boxes
        current_box = boxes_array[current, :4]
        other_boxes = boxes_array[indices[1:], :4]
        
        # Calculate intersection
        x1 = np.maximum(current_box[0], other_boxes[:, 0])
        y1 = np.maximum(current_box[1], other_boxes[:, 1])
        x2 = np.minimum(current_box[2], other_boxes[:, 2])
        y2 = np.minimum(current_box[3], other_boxes[:, 3])
        
        intersection = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
        
        # Calculate union
        current_area = (current_box[2] - current_box[0]) * (current_box[3] - current_box[1])
        other_areas = (other_boxes[:, 2] - other_boxes[:, 0]) * (other_boxes[:, 3] - other_boxes[:, 1])
        union = current_area + other_areas - intersection
        
        # Calculate IoU
        iou = intersection / (union + 1e-6)
        
        # Keep boxes with IoU less than threshold
        indices = indices[1:][iou < iou_threshold]
    
    # Return filtered boxes
    return [boxes[i] for i in keep]


def _apply_nms(boxes, iou_threshold=0.5):
    """Apply Non-Maximum Suppression to remove duplicate detections.
    
    Used for multi-scale signature detection.
    """
    return _apply_nms_to_detections(boxes, iou_threshold)


def _process_single_scale(base_bgr, s, rw, rh, conf, iou, augment):
    """Process a single scale - used for parallel execution."""
    tw, th = int(rw * s), int(rh * s)
    resized = cv2.resize(base_bgr, (tw, th), interpolation=cv2.INTER_CUBIC)
    boxes = yolo_detect_signatures(resized, imgsz=640, conf=conf, iou=iou, augment=augment)
    if not boxes:
        return []
    sx, sy = rw / max(1, tw), rh / max(1, th)
    mapped_boxes = []
    for (xyxy, score, cls) in boxes:
        xb1, yb1, xb2, yb2 = xyxy
        # Map back to original image coords
        x1o = xb1 * sx
        y1o = yb1 * sy
        x2o = xb2 * sx
        y2o = yb2 * sy
        mapped = (np.array([x1o, y1o, x2o, y2o]), float(score), int(cls))
        mapped_boxes.append(mapped)
    return mapped_boxes


def signature_only_infer(
    file: gr.File,
    try_scales: bool,
    conf: float,
    iou: float,
    augment: bool,
):
    if file is None:
        return None, "Upload an image or PDF", "[]"

    try:
        # Get file path - handle both direct path and gr.File object
        if hasattr(file, 'name'):
            path = file.name
        else:
            path = str(file)
        
        # Validate file exists
        if not os.path.exists(path):
            return None, f"File not found: {path}", "[]"
        
        ext = (os.path.splitext(path)[1] or "").lower()
        valid_exts = [".pdf", ".jpg", ".jpeg", ".png", ".tiff", ".bmp"]
        if ext not in valid_exts:
            return None, f"Invalid file format: {ext}. Supported: {', '.join(valid_exts)}", "[]"
    except Exception as e:
        return None, f"Error validating file: {str(e)}", "[]"
    
    # Load source image (first page for PDFs)
    ext = (os.path.splitext(path)[1] or "").lower()
    if ext in (".pdf",):
        doc = fitz.open(path)
        page = doc[0]
        pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
        base_rgb = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        doc.close()
    else:
        base_rgb = Image.open(path).convert("RGB")

    base_bgr = cv2.cvtColor(np.array(base_rgb), cv2.COLOR_RGB2BGR)

    scales = [1.0, 1.5, 2.0] if try_scales else [1.0]
    best = None
    all_boxes_mapped = []
    rh, rw = base_bgr.shape[:2]

    # Process scales in parallel if multiple scales
    if len(scales) > 1 and try_scales:
        with ThreadPoolExecutor(max_workers=2) as executor:
            futures = [
                executor.submit(_process_single_scale, base_bgr, s, rw, rh, conf, iou, augment)
                for s in scales
            ]
            for future in futures:
                boxes = future.result()
                all_boxes_mapped.extend(boxes)
    else:
        # Single scale - no threading overhead
        boxes = _process_single_scale(base_bgr, scales[0], rw, rh, conf, iou, augment)
        all_boxes_mapped.extend(boxes)

    # Apply NMS to remove duplicate detections from different scales
    if len(all_boxes_mapped) > 1:
        all_boxes_mapped = _apply_nms(all_boxes_mapped, iou_threshold=0.5)

    # Find best detection
    for box in all_boxes_mapped:
        if best is None or box[1] > best[1]:
            best = box

    # Annotate and prepare outputs
    annotated = annotate_signature_boxes_on_pil(base_rgb, all_boxes_mapped)
    det_json = [
        {
            "bbox": list(map(lambda v: float(v), xyxy.tolist() if hasattr(xyxy, "tolist") else list(xyxy))),
            "score": float(score),
            "class": int(cls)
        }
        for (xyxy, score, cls) in all_boxes_mapped
    ]
    summary = (
        f"Detections: {len(all_boxes_mapped)}" +
        (f" | Best score: {best[1]:.3f}" if best else " | No detections above threshold")
    )
    return annotated, summary, json.dumps(det_json, indent=2)

# Create Gradio interface
with gr.Blocks(title="Document Layout Detection", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 📄 Document Layout & Structure Detection
    
    Upload a document (PDF, image, etc.) to automatically detect its layout structure including text, tables, figures, and more!
    
    **Features:**
    - **AI-Powered Layout Detection**: Automatically identifies document elements
    - **Table Structure Extraction**: Recognizes and extracts table data
    - **OCR Support**: Reads text from scanned documents and images
    """)
    
    # Top-level tabs: Analyze and Signature Detection
    with gr.Tabs() as top_tabs:
        with gr.Tab("📄 Analyze"):
            with gr.Row():
                with gr.Column(scale=1):
                    file_input = gr.File(
                        label="Upload Document", 
                        file_types=[".pdf", ".jpg", ".jpeg", ".png", ".tiff", ".bmp"]
                    )
                    input_preview = gr.Image(label="Input Preview", type="filepath", height=240, interactive=False, show_label=True)
                    
                    mode_dropdown = gr.Dropdown(
                        choices=["Fast", "Accurate"],
                        value="Fast",
                        label="Processing Mode",
                        info="Accurate mode is slower but better for complex tables"
                    )
                    
                    ocr_checkbox = gr.Checkbox(
                        label="Enable OCR",
                        value=True,
                        info="Use OCR for scanned documents and images"
                    )
                    
                    tables_checkbox = gr.Checkbox(
                        label="Enable Table Detection",
                        value=True,
                        info="Detect and extract table structures"
                    )
                    
                    process_btn = gr.Button("🚀 Process Document", variant="primary", size="lg")
                    run_sig_chk = gr.Checkbox(label="Also detect signatures (Finetuned Signature Model)", value=False)
                    sig_conf_slider = gr.Slider(minimum=0.01, maximum=0.5, step=0.01, value=0.05, label="Signature confidence")
                
                with gr.Column(scale=2):
                    visualization_output = gr.Image(label="Layout Visualization (First Page)")
                    summary_output = gr.Markdown(label="Summary")

            with gr.Tabs():
                with gr.Tab("📝 Markdown Output"):
                    markdown_output = gr.Textbox(
                        label="Extracted Content (Markdown)",
                        lines=20,
                        max_lines=30
                    )
                
                with gr.Tab("🔧 JSON Layout Data"):
                    json_output = gr.Code(
                        label="Layout Predictions (JSON)",
                        language="json",
                        lines=20
                    )

            gr.Markdown("""
            ### Legend
            Different colors represent different document elements:
            
            **Layout Elements:**
            - 🔴 Title • 🔵 Text • 🟢 Section Header • 🟠 Table • 🟣 List/Figure/Formula
            
            **Picture Classifications (AI-detected):**
            - 🟣 Signature • 🟢 QR Code • 🟢 Barcode • 🟡 Logo • 🔴 Stamp
            - 🟦 Charts (Bar/Pie/Line) • 🟣 Flow Chart • 🟠 Screenshot • ⚪ Other
            
            ### How to Use
            1. Upload your document (PDF or image of ID card, invoice, report, etc.)
            2. Choose processing options (Fast mode recommended for quick results)
            3. Click "Process Document"
            4. View the visualization with bounding boxes and explore the outputs
            
            ### 💡 Try Examples Below!
            Click on any example document to see instant results on different document types.
            """)
            
            # Add examples; clicking a row will trigger file_input.change automatically
            with gr.Row():
                gr.Examples(
                    examples=[
                        ["sample/Screenshot 2025-10-13 114010.png", "Fast", True, True, False, 0.05],
                        ["sample/Screenshot 2025-10-13 114606.png", "Fast", True, True, False, 0.05],
                        ["sample/Screenshot 2025-10-15 191615.png", "Fast", True, True, False, 0.05],
                    ],
                    inputs=[file_input, mode_dropdown, ocr_checkbox, tables_checkbox, run_sig_chk, sig_conf_slider],
                    label="📚 Example Documents",
                    examples_per_page=3
                )

            # Connect the button
            process_btn.click(
                fn=gradio_interface,
                inputs=[file_input, mode_dropdown, ocr_checkbox, tables_checkbox, run_sig_chk, sig_conf_slider],
                outputs=[visualization_output, summary_output, markdown_output, json_output]
            )
            
            # Preview on file selection and auto-process
            file_input.change(
                fn=analyze_with_preview,
                inputs=[file_input, mode_dropdown, ocr_checkbox, tables_checkbox, run_sig_chk, sig_conf_slider],
                outputs=[input_preview, visualization_output, summary_output, markdown_output, json_output]
            )

        with gr.Tab("✍️ Signature Detection (Only)"):
            gr.Markdown("""
            Run the finetuned signature model on an image or the first page of a PDF. Simple controls, no ROI.
            """)
            with gr.Row():
                with gr.Column(scale=1):
                    sig_file_input = gr.File(
                        label="Upload Image or PDF (first page processed)", 
                        file_types=[".pdf", ".jpg", ".jpeg", ".png", ".tiff", ".bmp"]
                    )
                    sig_input_preview = gr.Image(label="Input Preview", type="filepath", height=240, interactive=False, show_label=True)
                    try_scales = gr.Checkbox(label="Try multiscale (1.0, 1.5, 2.0)", value=True)
                    sig_only_conf = gr.Slider(0.01, 0.5, value=0.03, step=0.01, label="Confidence")
                    sig_only_iou = gr.Slider(0.1, 0.9, value=0.45, step=0.05, label="IoU")
                    sig_only_aug = gr.Checkbox(label="Augment (slower, more recall)", value=True)
                    sig_run_btn = gr.Button("🔎 Detect Signatures", variant="primary")
                with gr.Column(scale=2):
                    sig_only_image = gr.Image(label="Annotated Signatures")
                    sig_only_summary = gr.Markdown(label="Signature Summary")
                    sig_only_json = gr.Code(label="Detections JSON", language="json", lines=16)

            gr.Examples(
                examples=[
                    ["sample_signature/X_074.jpeg", True, 0.03, 0.45, True],
                    ["sample_signature/X_014.jpeg", True, 0.03, 0.45, True],
                    ["sample_signature/X_081.jpeg", True, 0.03, 0.45, True]
                ],
                inputs=[sig_file_input, try_scales, sig_only_conf, sig_only_iou, sig_only_aug],
                label="✍️ Signature Examples",
                cache_examples=False
            )

            # Wire signature-only button
            sig_run_btn.click(
                fn=signature_only_infer,
                inputs=[sig_file_input, try_scales, sig_only_conf, sig_only_iou, sig_only_aug],
                outputs=[sig_only_image, sig_only_summary, sig_only_json]
            )

            # Preview for signature-only selection
            sig_file_input.change(
                fn=preview_first_page,
                inputs=[sig_file_input],
                outputs=[sig_input_preview]
            )
    
    # Events are now scoped within tabs above

# Launch the app
if __name__ == "__main__":
    # Queue with up to 2 concurrent workers (fits Spaces CPU with 2 cores)
    # Optional: pre-load signature model to reduce first-run latency (requires HF access)
    try:
        load_signature_model()
    except Exception:
        pass
    # Gradio v5 uses default_concurrency_limit; increase to 4 for better resource utilization
    # With 18GB RAM and 2 CPU cores, we can handle more concurrent requests
    try:
        demo.queue(default_concurrency_limit=4)
    except TypeError:
        demo.queue(concurrency_count=4)
    demo.launch()