app.py

import spaces
import json
import math
import os
import traceback
from io import BytesIO
from typing import Any, Dict, List, Optional, Tuple
import re

import fitz  # PyMuPDF
import gradio as gr
import requests
import torch
from huggingface_hub import snapshot_download
from PIL import Image, ImageDraw, ImageFont
from qwen_vl_utils import process_vision_info
from transformers import AutoModelForCausalLM, AutoProcessor

# Constants
MIN_PIXELS = 3136
MAX_PIXELS = 11289600
IMAGE_FACTOR = 28

# Prompts
prompt = """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.

1. Bbox format: [x1, y1, x2, y2]

2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title'].

3. Text Extraction & Formatting Rules:
    - Picture: For the 'Picture' category, the text field should be omitted.
    - Formula: Format its text as LaTeX.
    - Table: Format its text as HTML.
    - All Others (Text, Title, etc.): Format their text as Markdown.

4. Constraints:
    - The output text must be the original text from the image, with no translation.
    - All layout elements must be sorted according to human reading order.

5. Final Output: The entire output must be a single JSON object.
"""

# Utility functions
def round_by_factor(number: int, factor: int) -> int:
    """Returns the closest integer to 'number' that is divisible by 'factor'."""
    return round(number / factor) * factor


def smart_resize(
    height: int,
    width: int,
    factor: int = 28,
    min_pixels: int = 3136,
    max_pixels: int = 11289600,
):
    """Rescales the image so that the following conditions are met:
    1. Both dimensions (height and width) are divisible by 'factor'.
    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
    3. The aspect ratio of the image is maintained as closely as possible.
    """
    if max(height, width) / min(height, width) > 200:
        raise ValueError(
            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
        )
    h_bar = max(factor, round_by_factor(height, factor))
    w_bar = max(factor, round_by_factor(width, factor))

    if h_bar * w_bar > max_pixels:
        beta = math.sqrt((height * width) / max_pixels)
        h_bar = round_by_factor(height / beta, factor)
        w_bar = round_by_factor(width / beta, factor)
    elif h_bar * w_bar < min_pixels:
        beta = math.sqrt(min_pixels / (height * width))
        h_bar = round_by_factor(height * beta, factor)
        w_bar = round_by_factor(width * beta, factor)
    return h_bar, w_bar


def fetch_image(image_input, min_pixels: int = None, max_pixels: int = None):
    """Fetch and process an image"""
    if isinstance(image_input, str):
        if image_input.startswith(("http://", "https://")):
            response = requests.get(image_input)
            image = Image.open(BytesIO(response.content)).convert('RGB')
        else:
            image = Image.open(image_input).convert('RGB')
    elif isinstance(image_input, Image.Image):
        image = image_input.convert('RGB')
    else:
        raise ValueError(f"Invalid image input type: {type(image_input)}")
    
    if min_pixels is not None or max_pixels is not None:
        min_pixels = min_pixels or MIN_PIXELS
        max_pixels = max_pixels or MAX_PIXELS
        height, width = smart_resize(
            image.height, 
            image.width, 
            factor=IMAGE_FACTOR,
            min_pixels=min_pixels,
            max_pixels=max_pixels
        )
        image = image.resize((width, height), Image.LANCZOS)
    
    return image


def load_images_from_pdf(pdf_path: str) -> List[Image.Image]:
    """Load images from PDF file"""
    images = []
    try:
        pdf_document = fitz.open(pdf_path)
        for page_num in range(len(pdf_document)):
            page = pdf_document.load_page(page_num)
            # Convert page to image
            mat = fitz.Matrix(2.0, 2.0)  # Increase resolution
            pix = page.get_pixmap(matrix=mat)
            img_data = pix.tobytes("ppm")
            image = Image.open(BytesIO(img_data)).convert('RGB')
            images.append(image)
        pdf_document.close()
    except Exception as e:
        print(f"Error loading PDF: {e}")
        return []
    return images


def draw_layout_on_image(image: Image.Image, layout_data: List[Dict]) -> Image.Image:
    """Draw layout bounding boxes on image"""
    img_copy = image.copy()
    draw = ImageDraw.Draw(img_copy)
    
    # Colors for different categories
    colors = {
        'Caption': '#FF6B6B',
        'Footnote': '#4ECDC4', 
        'Formula': '#45B7D1',
        'List-item': '#96CEB4',
        'Page-footer': '#FFEAA7',
        'Page-header': '#DDA0DD',
        'Picture': '#FFD93D',
        'Section-header': '#6C5CE7',
        'Table': '#FD79A8',
        'Text': '#74B9FF',
        'Title': '#E17055'
    }
    
    try:
        # Load a font
        try:
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 12)
        except Exception:
            font = ImageFont.load_default()
        
        for item in layout_data:
            if 'bbox' in item and 'category' in item:
                bbox = item['bbox']
                category = item['category']
                color = colors.get(category, '#000000')
                
                # Draw rectangle
                draw.rectangle(bbox, outline=color, width=2)
                
                # Draw label
                label = category
                label_bbox = draw.textbbox((0, 0), label, font=font)
                label_width = label_bbox[2] - label_bbox[0]
                label_height = label_bbox[3] - label_bbox[1]
                
                # Position label above the box
                label_x = bbox[0]
                label_y = max(0, bbox[1] - label_height - 2)
                
                # Draw background for label
                draw.rectangle(
                    [label_x, label_y, label_x + label_width + 4, label_y + label_height + 2],
                    fill=color
                )
                
                # Draw text
                draw.text((label_x + 2, label_y + 1), label, fill='white', font=font)
                
    except Exception as e:
        print(f"Error drawing layout: {e}")
    
    return img_copy


def is_arabic_text(text: str) -> bool:
    """Check if text in headers and paragraphs contains mostly Arabic characters"""
    if not text:
        return False
    
    # Extract text from headers and paragraphs only
    # Match markdown headers (# ## ###) and regular paragraph text
    header_pattern = r'^#{1,6}\s+(.+)$'
    paragraph_pattern = r'^(?!#{1,6}\s|!\[|```|\||\s*[-*+]\s|\s*\d+\.\s)(.+)$'
    
    content_text = []
    
    for line in text.split('\n'):
        line = line.strip()
        if not line:
            continue
            
        # Check for headers
        header_match = re.match(header_pattern, line, re.MULTILINE)
        if header_match:
            content_text.append(header_match.group(1))
            continue
            
        # Check for paragraph text (exclude lists, tables, code blocks, images)
        if re.match(paragraph_pattern, line, re.MULTILINE):
            content_text.append(line)
    
    if not content_text:
        return False
    
    # Join all content text and check for Arabic characters
    combined_text = ' '.join(content_text)
    
    # Arabic Unicode ranges
    arabic_chars = 0
    total_chars = 0
    
    for char in combined_text:
        if char.isalpha():
            total_chars += 1
            # Arabic script ranges
            if ('\u0600' <= char <= '\u06FF') or ('\u0750' <= char <= '\u077F') or ('\u08A0' <= char <= '\u08FF'):
                arabic_chars += 1
    
    if total_chars == 0:
        return False
    
    # Consider text as Arabic if more than 50% of alphabetic characters are Arabic
    return (arabic_chars / total_chars) > 0.5


def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
    """Convert layout JSON to markdown format"""
    import base64
    from io import BytesIO
    
    markdown_lines = []
    
    try:
        # Sort items by reading order (top to bottom, left to right)
        sorted_items = sorted(layout_data, key=lambda x: (x.get('bbox', [0, 0, 0, 0])[1], x.get('bbox', [0, 0, 0, 0])[0]))
        
        for item in sorted_items:
            category = item.get('category', '')
            text = item.get(text_key, '')
            bbox = item.get('bbox', [])
            
            if category == 'Picture':
                # Extract image region and embed it
                if bbox and len(bbox) == 4:
                    try:
                        # Extract the image region
                        x1, y1, x2, y2 = bbox
                        # Ensure coordinates are within image bounds
                        x1, y1 = max(0, int(x1)), max(0, int(y1))
                        x2, y2 = min(image.width, int(x2)), min(image.height, int(y2))
                        
                        if x2 > x1 and y2 > y1:
                            cropped_img = image.crop((x1, y1, x2, y2))
                            
                            # Convert to base64 for embedding
                            buffer = BytesIO()
                            cropped_img.save(buffer, format='PNG')
                            img_data = base64.b64encode(buffer.getvalue()).decode()
                            
                            # Add as markdown image
                            markdown_lines.append(f"![Image](data:image/png;base64,{img_data})\n")
                        else:
                            markdown_lines.append("![Image](Image region detected)\n")
                    except Exception as e:
                        print(f"Error processing image region: {e}")
                        markdown_lines.append("![Image](Image detected)\n")
                else:
                    markdown_lines.append("![Image](Image detected)\n")
            elif not text:
                continue
            elif category == 'Title':
                markdown_lines.append(f"# {text}\n")
            elif category == 'Section-header':
                markdown_lines.append(f"## {text}\n")
            elif category == 'Text':
                markdown_lines.append(f"{text}\n")
            elif category == 'List-item':
                markdown_lines.append(f"- {text}\n")
            elif category == 'Table':
                # If text is already HTML, keep it as is
                if text.strip().startswith('<'):
                    markdown_lines.append(f"{text}\n")
                else:
                    markdown_lines.append(f"**Table:** {text}\n")
            elif category == 'Formula':
                # If text is LaTeX, format it properly
                if text.strip().startswith('$') or '\\' in text:
                    markdown_lines.append(f"$$\n{text}\n$$\n")
                else:
                    markdown_lines.append(f"**Formula:** {text}\n")
            elif category == 'Caption':
                markdown_lines.append(f"*{text}*\n")
            elif category == 'Footnote':
                markdown_lines.append(f"^{text}^\n")
            elif category in ['Page-header', 'Page-footer']:
                # Skip headers and footers in main content
                continue
            else:
                markdown_lines.append(f"{text}\n")
            
            markdown_lines.append("")  # Add spacing
            
    except Exception as e:
        print(f"Error converting to markdown: {e}")
        return str(layout_data)
    
    return "\n".join(markdown_lines)

# Initialize model/processor lazily inside GPU context
model_id = "rednote-hilab/dots.ocr"
model_path = "./models/dots-ocr-local"
model = None
processor = None

def ensure_model_loaded():
    """Lazily download and load model/processor using eager attention (no FlashAttention)."""
    global model, processor
    if model is not None and processor is not None:
        return

    # Always use eager attention
    attn_impl = "eager"
    # Use GPU if available, otherwise CPU
    if torch.cuda.is_available():
        dtype = torch.bfloat16  # Use bfloat16 on GPU for consistency
        device_map = "auto"
    else:
        dtype = torch.float32
        device_map = "cpu"

    # Download snapshot locally (idempotent)
    snapshot_download(
        repo_id=model_id,
        local_dir=model_path,
        local_dir_use_symlinks=False,
    )

    # Load model/processor
    loaded_model = AutoModelForCausalLM.from_pretrained(
        model_path,
        attn_implementation=attn_impl,
        torch_dtype=dtype,
        device_map=device_map,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )
    loaded_processor = AutoProcessor.from_pretrained(
        model_path,
        trust_remote_code=True,
    )

    model = loaded_model
    processor = loaded_processor

# Global state variables
device = "cuda" if torch.cuda.is_available() else "cpu"

# PDF handling state
pdf_cache = {
    "images": [],
    "current_page": 0,
    "total_pages": 0,
    "file_type": None,
    "is_parsed": False,
    "results": []
}
@spaces.GPU(duration=300)
def inference(image: Image.Image, prompt: str, max_new_tokens: int = 24000) -> str:
    """Run inference on an image with the given prompt"""
    try:
        ensure_model_loaded()
        if model is None or processor is None:
            raise RuntimeError("Model not loaded. Please check model initialization.")
        
        # Prepare messages in the expected format
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": image
                    },
                    {"type": "text", "text": prompt}
                ]
            }
        ]
        
        # Apply chat template
        text = processor.apply_chat_template(
            messages, 
            tokenize=False, 
            add_generation_prompt=True
        )
        
        # Process vision information
        image_inputs, video_inputs = process_vision_info(messages)
        
        # Prepare inputs
        inputs = processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        
        # Move to the model's primary device (works with device_map as well)
        primary_device = next(model.parameters()).device
        inputs = inputs.to(primary_device)
        
        # Generate output
        with torch.no_grad():
            generated_ids = model.generate(
                **inputs, 
                max_new_tokens=max_new_tokens,
                do_sample=False,
                temperature=0.1
            )
        
        # Decode output
        generated_ids_trimmed = [
            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        
        output_text = processor.batch_decode(
            generated_ids_trimmed, 
            skip_special_tokens=True, 
            clean_up_tokenization_spaces=False
        )
        
        return output_text[0] if output_text else ""
        
    except Exception as e:
        print(f"Error during inference: {e}")
        traceback.print_exc()
        return f"Error during inference: {str(e)}"


@spaces.GPU(duration=300)
def _generate_text_and_confidence_for_crop(
    image: Image.Image,
    max_new_tokens: int = 128,
    temperature: float = 0.1,
) -> Tuple[str, float]:
    """Generate text for a cropped region and compute average per-token confidence from model scores.

    Returns (generated_text, average_confidence_percent).
    """
    try:
        ensure_model_loaded()
        # Prepare a concise extraction prompt for the crop
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {
                        "type": "text",
                        "text": "Extract the exact text content from this image region. Output text only without translation or additional words.",
                    },
                ],
            }
        ]

        # Apply chat template
        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        # Process vision information
        image_inputs, video_inputs = process_vision_info(messages)

        # Prepare inputs
        inputs = processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        primary_device = next(model.parameters()).device
        inputs = inputs.to(primary_device)

        # Generate with scores
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                temperature=temperature,
                output_scores=True,
                return_dict_in_generate=True,
            )

        sequences = outputs.sequences  # [batch, seq_len]
        input_len = inputs.input_ids.shape[1]
        # Trim input prompt ids to isolate generated tokens
        generated_ids = sequences[:, input_len:]
        generated_text = processor.batch_decode(
            generated_ids,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )[0].strip()

        # Compute average probability of chosen tokens
        confidences: List[float] = []
        for step, step_scores in enumerate(outputs.scores or []):
            # step_scores: [batch, vocab]
            probs = torch.nn.functional.softmax(step_scores, dim=-1)
            # token id chosen at this step
            if input_len + step < sequences.shape[1]:
                chosen_ids = sequences[:, input_len + step].unsqueeze(-1)
                chosen_probs = probs.gather(dim=-1, index=chosen_ids)  # [batch, 1]
                confidences.append(float(chosen_probs[0, 0].item()))

        avg_conf_percent = (sum(confidences) / len(confidences) * 100.0) if confidences else 0.0
        return generated_text, avg_conf_percent
    except Exception as e:
        print(f"Error generating crop confidence: {e}")
        traceback.print_exc()
        return "", 0.0


def process_image(
    image: Image.Image,
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None,
    max_new_tokens: int = 24000,
) -> Dict[str, Any]:
    """Process a single image with the specified prompt mode"""
    try:
        # Resize image if needed
        if min_pixels is not None or max_pixels is not None:
            image = fetch_image(image, min_pixels=min_pixels, max_pixels=max_pixels)
        
        # Run inference with the default prompt
        raw_output = inference(image, prompt, max_new_tokens=max_new_tokens)
        
        # Process results based on prompt mode
        result = {
            'original_image': image,
            'raw_output': raw_output,
            'processed_image': image,
            'layout_result': None,
            'markdown_content': None
        }
        
        # Try to parse JSON and create visualizations (since we're doing layout analysis)
        try:
            # Try to parse JSON output
            layout_data = json.loads(raw_output)

            # Compute per-region confidence using the model on each cropped region
            for idx, item in enumerate(layout_data):
                try:
                    bbox = item.get('bbox', [])
                    text_content = item.get('text', '')
                    category = item.get('category', '')
                    if (not text_content) or category == 'Picture' or not bbox or len(bbox) != 4:
                        continue
                    x1, y1, x2, y2 = bbox
                    x1, y1 = max(0, int(x1)), max(0, int(y1))
                    x2, y2 = min(image.width, int(x2)), min(image.height, int(y2))
                    if x2 <= x1 or y2 <= y1:
                        continue
                    crop_img = image.crop((x1, y1, x2, y2))
                    # Generate and score text for this crop; we only keep the confidence
                    _, region_conf = _generate_text_and_confidence_for_crop(crop_img)
                    item['confidence'] = region_conf
                except Exception as e:
                    print(f"Error scoring region {idx}: {e}")
                    # Leave confidence absent if scoring fails

            result['layout_result'] = layout_data
            
            # Create visualization with bounding boxes
            try:
                processed_image = draw_layout_on_image(image, layout_data)
                result['processed_image'] = processed_image
            except Exception as e:
                print(f"Error drawing layout: {e}")
                result['processed_image'] = image
            
            # Generate markdown from layout data
            try:
                markdown_content = layoutjson2md(image, layout_data, text_key='text')
                result['markdown_content'] = markdown_content
            except Exception as e:
                print(f"Error generating markdown: {e}")
                result['markdown_content'] = raw_output
            
        except json.JSONDecodeError:
            print("Failed to parse JSON output, using raw output")
            result['markdown_content'] = raw_output
        
        return result
        
    except Exception as e:
        print(f"Error processing image: {e}")
        traceback.print_exc()
        return {
            'original_image': image,
            'raw_output': f"Error processing image: {str(e)}",
            'processed_image': image,
            'layout_result': None,
            'markdown_content': f"Error processing image: {str(e)}"
        }


def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
    """Load file for preview (supports PDF and images)"""
    global pdf_cache
    
    if not file_path or not os.path.exists(file_path):
        return None, "No file selected"
    
    file_ext = os.path.splitext(file_path)[1].lower()
    
    try:
        if file_ext == '.pdf':
            # Load PDF pages
            images = load_images_from_pdf(file_path)
            if not images:
                return None, "Failed to load PDF"
            
            pdf_cache.update({
                "images": images,
                "current_page": 0,
                "total_pages": len(images),
                "file_type": "pdf",
                "is_parsed": False,
                "results": []
            })
            
            return images[0], f"Page 1 / {len(images)}"
            
        elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
            # Load single image
            image = Image.open(file_path).convert('RGB')
            
            pdf_cache.update({
                "images": [image],
                "current_page": 0,
                "total_pages": 1,
                "file_type": "image",
                "is_parsed": False,
                "results": []
            })
            
            return image, "Page 1 / 1"
        else:
            return None, f"Unsupported file format: {file_ext}"
            
    except Exception as e:
        print(f"Error loading file: {e}")
        return None, f"Error loading file: {str(e)}"


def turn_page(direction: str) -> Tuple[Optional[Image.Image], str, List, Any, Optional[Image.Image], Optional[Dict]]:
    """Navigate through PDF pages and update all relevant outputs."""
    global pdf_cache

    if not pdf_cache["images"]:
        return None, '<div class="page-info">No file loaded</div>', [], "No results yet", None, None

    if direction == "prev":
        pdf_cache["current_page"] = max(0, pdf_cache["current_page"] - 1)
    elif direction == "next":
        pdf_cache["current_page"] = min(
            pdf_cache["total_pages"] - 1,
            pdf_cache["current_page"] + 1
        )

    index = pdf_cache["current_page"]
    current_image_preview = pdf_cache["images"][index]
    page_info_html = f'<div class="page-info">Page {index + 1} / {pdf_cache["total_pages"]}</div>'

    # Initialize default result values
    markdown_content = "Page not processed yet"
    processed_img = None
    layout_json = None
    ocr_table_data = []

    # Get results for current page if available
    if (pdf_cache["is_parsed"] and
        index < len(pdf_cache["results"]) and
        pdf_cache["results"][index]):

        result = pdf_cache["results"][index]
        markdown_content = result.get('markdown_content') or result.get('raw_output', 'No content available')
        processed_img = result.get('processed_image', None) # Get the processed image
        layout_json = result.get('layout_result', None) # Get the layout JSON
        
        # Generate OCR table for current page
        if layout_json and result.get('original_image'):
            # Need to import the helper here or move it outside
            import base64
            from io import BytesIO
            
            for idx, item in enumerate(layout_json):
                bbox = item.get('bbox', [])
                text = item.get('text', '')
                category = item.get('category', '')
                
                if not text or category == 'Picture':
                    continue
                
                img_html = ""
                if bbox and len(bbox) == 4:
                    try:
                        x1, y1, x2, y2 = bbox
                        orig_img = result['original_image']
                        x1, y1 = max(0, int(x1)), max(0, int(y1))
                        x2, y2 = min(orig_img.width, int(x2)), min(orig_img.height, int(y2))
                        
                        if x2 > x1 and y2 > y1:
                            cropped_img = orig_img.crop((x1, y1, x2, y2))
                            buffer = BytesIO()
                            cropped_img.save(buffer, format='PNG')
                            img_data = base64.b64encode(buffer.getvalue()).decode()
                            img_html = f'<img src="data:image/png;base64,{img_data}" style="max-width:200px; max-height:100px; object-fit:contain;" />'
                    except Exception as e:
                        print(f"Error cropping region {idx}: {e}")
                        img_html = f"<div>Region {idx+1}</div>"
                else:
                    img_html = f"<div>Region {idx+1}</div>"
                
                # Extract confidence from item if available, otherwise N/A
                confidence = item.get('confidence', 'N/A')
                if isinstance(confidence, (int, float)):
                    confidence = f"{confidence:.1f}%"
                elif confidence != 'N/A':
                    confidence = str(confidence)
                    
                ocr_table_data.append([img_html, text, confidence])

    # Check for Arabic text to set RTL property
    if is_arabic_text(markdown_content):
        markdown_update = gr.update(value=markdown_content, rtl=True)
    else:
        markdown_update = markdown_content

    return current_image_preview, page_info_html, ocr_table_data, markdown_update, processed_img, layout_json


def create_gradio_interface():
    """Create the Gradio interface"""
    
    # Custom CSS
    css = """
    .main-container {
        max-width: 1400px;
        margin: 0 auto;
    }
    
    .header-text {
        text-align: center;
        color: #2c3e50;
        margin-bottom: 20px;
    }
    
    .process-button {
        border: none !important;
        color: white !important;
        font-weight: bold !important;
    }
    
    .process-button:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important;
    }
    
    .info-box {
        border: 1px solid #dee2e6;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    
    .page-info {
        text-align: center;
        padding: 8px 16px;
        border-radius: 20px;
        font-weight: bold;
        margin: 10px 0;
    }
    
    .model-status {
        padding: 10px;
        border-radius: 8px;
        margin: 10px 0;
        text-align: center;
        font-weight: bold;
    }
    
    .status-ready {
        background: #d1edff;
        color: #0c5460;
        border: 1px solid #b8daff;
    }
    """
    
    with gr.Blocks(theme=gr.themes.Soft(), css=css, title="Arabic OCR - Document Text Extraction") as demo:
        
        # Header
        gr.HTML("""
        <div class="title" style="text-align: center">
            <h1>🔍 Arabic OCR - Professional Document Text Extraction</h1>
            <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
                Advanced AI-powered OCR solution for Arabic documents with high accuracy layout detection and text extraction
            </p>
        </div>
        """)
        
        # Main interface
        with gr.Row():
            # Left column - Input and controls
            with gr.Column(scale=1):
                
                # File input
                file_input = gr.File(
                    label="Upload Image or PDF",
                    file_types=[".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".pdf"],
                    type="filepath"
                )
                
                # Image preview
                image_preview = gr.Image(
                    label="Preview",
                    type="pil",
                    interactive=False,
                    height=300
                )
                
                # Page navigation for PDFs
                with gr.Row():
                    prev_page_btn = gr.Button("◀ Previous", size="md")
                    page_info = gr.HTML('<div class="page-info">No file loaded</div>')
                    next_page_btn = gr.Button("Next ▶", size="md")
                
                # Advanced settings
                with gr.Accordion("Advanced Settings", open=False):
                    max_new_tokens = gr.Slider(
                        minimum=1000,
                        maximum=32000,
                        value=24000,
                        step=1000,
                        label="Max New Tokens",
                        info="Maximum number of tokens to generate"
                    )
                    
                    min_pixels = gr.Number(
                        value=MIN_PIXELS,
                        label="Min Pixels",
                        info="Minimum image resolution"
                    )
                    
                    max_pixels = gr.Number(
                        value=MAX_PIXELS,
                        label="Max Pixels", 
                        info="Maximum image resolution"
                    )
                
                # Process button
                process_btn = gr.Button(
                    "🚀 Process Document",
                    variant="primary",
                    elem_classes=["process-button"],
                    size="lg"
                )
                
                # Clear button
                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
            
            # Right column - Results
            with gr.Column(scale=2):
                
                # Results tabs
                with gr.Tabs():
                    # Processed image tab
                    with gr.Tab("🖼️ Processed Image"):
                        processed_image = gr.Image(
                            label="Image with Layout Detection",
                            type="pil",
                            interactive=False,
                            height=500
                        )
                    # Editable OCR Results Table
                    with gr.Tab("📊 OCR Results Table"):
                        gr.Markdown("### Editable OCR Results\nReview and edit the extracted text for each detected region")
                        ocr_table = gr.Dataframe(
                            headers=["Region Image", "Extracted Text", "Confidence"],
                            datatype=["html", "str", "str"],
                            label="OCR Results",
                            interactive=True,
                            wrap=True
                        )
                    # Markdown output tab  
                    with gr.Tab("📝 Extracted Content"):
                        markdown_output = gr.Markdown(
                            value="Click 'Process Document' to see extracted content...",
                            height=500
                        )
                    # JSON layout tab
                    with gr.Tab("📋 Layout JSON"):
                        json_output = gr.JSON(
                            label="Layout Analysis Results",
                            value=None
                        )
        
        # Helper function to create OCR table
        def create_ocr_table(image: Image.Image, layout_data: List[Dict]) -> List[List[str]]:
            """Create table data from layout results with cropped images"""
            import base64
            from io import BytesIO
            
            if not layout_data:
                return []
            
            table_data = []
            
            for idx, item in enumerate(layout_data):
                bbox = item.get('bbox', [])
                text = item.get('text', '')
                category = item.get('category', '')
                
                # Skip items without text or Picture category
                if not text or category == 'Picture':
                    continue
                
                # Crop the image region
                img_html = ""
                if bbox and len(bbox) == 4:
                    try:
                        x1, y1, x2, y2 = bbox
                        # Ensure coordinates are within image bounds
                        x1, y1 = max(0, int(x1)), max(0, int(y1))
                        x2, y2 = min(image.width, int(x2)), min(image.height, int(y2))
                        
                        if x2 > x1 and y2 > y1:
                            cropped_img = image.crop((x1, y1, x2, y2))
                            
                            # Convert to base64 for HTML display
                            buffer = BytesIO()
                            cropped_img.save(buffer, format='PNG')
                            img_data = base64.b64encode(buffer.getvalue()).decode()
                            
                            # Create HTML img tag
                            img_html = f'<img src="data:image/png;base64,{img_data}" style="max-width:200px; max-height:100px; object-fit:contain;" />'
                    except Exception as e:
                        print(f"Error cropping region {idx}: {e}")
                        img_html = f"<div>Region {idx+1}</div>"
                else:
                    img_html = f"<div>Region {idx+1}</div>"
                
                # Add confidence score - extract from item if available, otherwise N/A
                confidence = item.get('confidence', 'N/A')
                if isinstance(confidence, (int, float)):
                    confidence = f"{confidence:.1f}%"
                elif confidence != 'N/A':
                    confidence = str(confidence)
                
                # Add row to table
                table_data.append([img_html, text, confidence])
            
            return table_data
        
        # Event handlers
        @spaces.GPU(duration=240)
        def process_document(file_path, max_tokens, min_pix, max_pix):
            """Process the uploaded document"""
            global pdf_cache
            
            try:
                # Ensure model/processor are loaded within GPU context
                ensure_model_loaded()
                if not file_path:
                    return None, [], "Please upload a file first.", None
                
                if model is None:
                    return None, [], "Model not loaded. Please refresh the page and try again.", None
                
                # Load and preview file
                image, page_info = load_file_for_preview(file_path)
                if image is None:
                    return None, [], page_info, None
                
                # Process the image(s)
                if pdf_cache["file_type"] == "pdf":
                    # Process all pages for PDF
                    all_results = []
                    all_markdown = []
                    
                    for i, img in enumerate(pdf_cache["images"]):
                        result = process_image(
                            img,
                            min_pixels=int(min_pix) if min_pix else None,
                            max_pixels=int(max_pix) if max_pix else None,
                            max_new_tokens=int(max_tokens) if max_tokens else 24000,
                        )
                        all_results.append(result)
                        if result.get('markdown_content'):
                            all_markdown.append(f"## Page {i+1}\n\n{result['markdown_content']}")
                    
                    pdf_cache["results"] = all_results
                    pdf_cache["is_parsed"] = True
                    
                    # Show results for first page
                    first_result = all_results[0]
                    combined_markdown = "\n\n---\n\n".join(all_markdown)
                    
                    # Check if the combined markdown contains mostly Arabic text
                    if is_arabic_text(combined_markdown):
                        markdown_update = gr.update(value=combined_markdown, rtl=True)
                    else:
                        markdown_update = combined_markdown
                    
                    # Create OCR table for first page
                    ocr_table_data = []
                    if first_result['layout_result']:
                        ocr_table_data = create_ocr_table(
                            first_result['original_image'],
                            first_result['layout_result']
                        )
                    
                    return (
                        first_result['processed_image'],
                        ocr_table_data,
                        markdown_update,
                        first_result['layout_result']
                    )
                else:
                    # Process single image
                    result = process_image(
                        image,
                        min_pixels=int(min_pix) if min_pix else None,
                        max_pixels=int(max_pix) if max_pix else None,
                        max_new_tokens=int(max_tokens) if max_tokens else 24000,
                    )
                    
                    pdf_cache["results"] = [result]
                    pdf_cache["is_parsed"] = True
                    
                    # Check if the content contains mostly Arabic text
                    content = result['markdown_content'] or "No content extracted"
                    if is_arabic_text(content):
                        markdown_update = gr.update(value=content, rtl=True)
                    else:
                        markdown_update = content
                    
                    # Create OCR table
                    ocr_table_data = []
                    if result['layout_result']:
                        ocr_table_data = create_ocr_table(
                            result['original_image'],
                            result['layout_result']
                        )
                    
                    return (
                        result['processed_image'],
                        ocr_table_data,
                        markdown_update,
                        result['layout_result']
                    )
                    
            except Exception as e:
                error_msg = f"Error processing document: {str(e)}"
                print(error_msg)
                traceback.print_exc()
                return None, [], error_msg, None
        
        def handle_file_upload(file_path):
            """Handle file upload and show preview"""
            if not file_path:
                return None, "No file loaded"
            
            image, page_info = load_file_for_preview(file_path)
            return image, page_info
        
        def handle_page_turn(direction):
            """Handle page navigation"""
            image, page_info, result = turn_page(direction)
            return image, page_info, result
        
        def clear_all():
            """Clear all data and reset interface"""
            global pdf_cache

            pdf_cache = {
                "images": [], "current_page": 0, "total_pages": 0,
                "file_type": None, "is_parsed": False, "results": []
            }

            return (
                None,  # file_input
                None,  # image_preview
                '<div class="page-info">No file loaded</div>',  # page_info
                None,  # processed_image
                [],  # ocr_table
                "Click 'Process Document' to see extracted content...",  # markdown_output
                None,  # json_output
            )
        
        # Wire up event handlers
        file_input.change(
            handle_file_upload,
            inputs=[file_input],
            outputs=[image_preview, page_info]
        )
        
        # The outputs list is now updated to include all components that need to change
        prev_page_btn.click(
            lambda: turn_page("prev"),
            outputs=[image_preview, page_info, ocr_table, markdown_output, processed_image, json_output]
        )

        next_page_btn.click(
            lambda: turn_page("next"),
            outputs=[image_preview, page_info, ocr_table, markdown_output, processed_image, json_output]
        )
        
        process_btn.click(
            process_document,
            inputs=[file_input, max_new_tokens, min_pixels, max_pixels],
            outputs=[processed_image, ocr_table, markdown_output, json_output]
        )
        
        # The outputs list for the clear button is now correct
        clear_btn.click(
            clear_all,
            outputs=[
                file_input, image_preview, page_info, processed_image,
                ocr_table, markdown_output, json_output
            ]
        )
    
    return demo


if __name__ == "__main__":
    # Create and launch the interface
    demo = create_gradio_interface()
    demo.queue(max_size=10).launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True,
        show_error=True
    )