import gradio as gr
import torch
import os
import sys
import tempfile
import shutil
import types
import re
import base64
import numpy as np
from typing import Iterable
from io import StringIO, BytesIO
from PIL import Image, ImageDraw, ImageFont, ImageOps
from transformers import AutoModel, AutoTokenizer
import torch.nn.functional as F
import spaces
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

try:
    import fitz  # PyMuPDF for PDF support
    PDF_SUPPORT = True
except ImportError:
    PDF_SUPPORT = False
    print("⚠️ PyMuPDF not installed. PDF support disabled.")

# --- Device Setup ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("=" * 50)
print("🔧 System Information")
print("=" * 50)
print(f"PyTorch: {torch.__version__} | CUDA: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"Device: {device}")
print("=" * 50)

# --- Custom Theme: Tech Blue/Cyan ---
colors.tech_blue = colors.Color(
    name="tech_blue",
    c50="#f3f7fb",
    c100="#e6eef7",
    c200="#c6d9ee",
    c300="#9bbbe0",
    c400="#6a96cf",
    c500="#3f6fb8",
    c600="#2f569d",
    c700="#26457f",
    c800="#213b69",
    c900="#1d3258",
    c950="#0f1b34",
)

colors.electric = colors.Color(
    name="electric",
    c50="#f6f4fb",
    c100="#ede9f8",
    c200="#d9d1f0",
    c300="#bba9e3",
    c400="#9a7ed3",
    c500="#7a56c2",
    c600="#6541aa",
    c700="#52348c",
    c800="#442c72",
    c900="#38265d",
    c950="#1e1436",
)


class TechTheme(Soft):
    def __init__(self):
        super().__init__(
            primary_hue=colors.tech_blue,
            secondary_hue=colors.electric,
            neutral_hue=colors.slate,
            text_size=sizes.text_md,
            font=(fonts.GoogleFont("Inter"), "system-ui", "sans-serif"),
            font_mono=(fonts.GoogleFont("JetBrains Mono"), "monospace"),
        )
        super().set(
            # Background - clean tech gradient
            body_background_fill="linear-gradient(135deg, #ecfeff 0%, #ffffff 35%, #eff6ff 65%, #f0f9ff 100%)",
            body_background_fill_dark="linear-gradient(135deg, #0f172a 0%, #020617 50%, #0c1929 100%)",
            # Buttons - cyan to blue gradient
            button_primary_background_fill="linear-gradient(135deg, #06b6d4 0%, #3b82f6 100%)",
            button_primary_background_fill_hover="linear-gradient(135deg, #0891b2 0%, #2563eb 100%)",
            button_primary_text_color="white",
            button_primary_shadow="0 4px 20px rgba(6, 182, 212, 0.35)",
            button_primary_shadow_hover="0 6px 24px rgba(6, 182, 212, 0.45)",
            # Blocks - glassmorphism
            block_background_fill="rgba(255, 255, 255, 0.85)",
            block_background_fill_dark="rgba(15, 23, 42, 0.7)",
            block_shadow="0 4px 24px rgba(6, 182, 212, 0.1)",
            block_border_width="1px",
            block_border_color="rgba(6, 182, 212, 0.12)",
            block_border_color_dark="rgba(6, 182, 212, 0.25)",
            block_title_text_weight="600",
            block_label_text_weight="500",
            # Inputs
            input_background_fill="rgba(255, 255, 255, 0.95)",
            input_background_fill_dark="rgba(15, 23, 42, 0.8)",
            input_border_color="rgba(6, 182, 212, 0.2)",
            input_border_color_focus="rgba(6, 182, 212, 0.5)",
            input_shadow="0 2px 8px rgba(6, 182, 212, 0.06)",
            input_shadow_focus="0 0 0 3px rgba(6, 182, 212, 0.12)",
            # Tabs
            border_color_accent="*primary_500",
        )


theme = TechTheme()

# --- Custom CSS ---
css = """
/* Main Title with tech gradient */
.main-title {
    text-align: center;
    padding: 1.5rem 0 1rem;
}
.main-title h1 {
    background: linear-gradient(135deg, #06b6d4 0%, #3b82f6 50%, #0ea5e9 100%);
    background-size: 200% 200%;
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-size: 2.4rem !important;
    font-weight: 700 !important;
    margin: 0 !important;
    letter-spacing: -0.02em;
    animation: gradient-shift 6s ease infinite;
}
@keyframes gradient-shift {
    0%, 100% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
}
.subtitle {
    color: #64748b;
    font-size: 0.95rem;
    margin-top: 0.4rem;
    font-weight: 400;
}

/* Extract Button */
#extract-btn {
    font-size: 1rem !important;
    font-weight: 600 !important;
    padding: 12px 28px !important;
    border-radius: 10px !important;
    transition: all 0.3s ease !important;
}
#extract-btn:hover {
    transform: translateY(-2px);
}

/* Output text styling */
.output-text textarea {
    font-family: 'SF Mono', 'Monaco', 'Inconsolata', 'Roboto Mono', 'Source Code Pro', 'Consolas', 'Liberation Mono', 'Courier New', monospace !important;
    font-size: 0.875rem !important;
    line-height: 1.6 !important;
    font-variant-ligatures: none !important;
    -webkit-font-feature-settings: "liga" 0 !important;
    font-feature-settings: "liga" 0 !important;
}

/* Tab styling */
.tab-nav {
    border-bottom: 2px solid rgba(6, 182, 212, 0.1) !important;
}
.tab-nav button {
    font-weight: 500 !important;
    padding: 10px 18px !important;
    border-radius: 8px 8px 0 0 !important;
    transition: all 0.2s ease !important;
}
.tab-nav button.selected {
    background: linear-gradient(135deg, rgba(6, 182, 212, 0.08), rgba(59, 130, 246, 0.08)) !important;
    border-bottom: 2px solid #06b6d4 !important;
}

/* Card hover effects */
.gradio-group {
    transition: all 0.3s ease;
    border-radius: 14px !important;
}
.gradio-group:hover {
    box-shadow: 0 8px 32px rgba(6, 182, 212, 0.12);
}

/* Dropdown styling */
.gradio-dropdown {
    border-radius: 8px !important;
}

/* Image container */
.gradio-image {
    border-radius: 10px !important;
    overflow: hidden;
}

/* Gallery styling */
.gradio-gallery {
    border-radius: 10px !important;
}
.gradio-gallery .gallery-item {
    border-radius: 8px !important;
    border: 1px solid rgba(6, 182, 212, 0.15) !important;
}

/* Accordion styling */
.gradio-accordion {
    border-radius: 10px !important;
    border: 1px solid rgba(6, 182, 212, 0.12) !important;
}
.gradio-accordion > .label-wrap {
    background: linear-gradient(135deg, rgba(6, 182, 212, 0.04), rgba(59, 130, 246, 0.04)) !important;
}

/* Scrollbar styling */
::-webkit-scrollbar {
    width: 6px;
    height: 6px;
}
::-webkit-scrollbar-track {
    background: rgba(6, 182, 212, 0.05);
    border-radius: 3px;
}
::-webkit-scrollbar-thumb {
    background: linear-gradient(135deg, #06b6d4, #3b82f6);
    border-radius: 3px;
}
::-webkit-scrollbar-thumb:hover {
    background: linear-gradient(135deg, #0891b2, #2563eb);
}

/* File upload area */
.gradio-file {
    border: 2px dashed rgba(6, 182, 212, 0.25) !important;
    border-radius: 10px !important;
    transition: all 0.3s ease;
}
.gradio-file:hover {
    border-color: rgba(6, 182, 212, 0.45) !important;
    background: rgba(6, 182, 212, 0.02) !important;
}

/* Examples section */
.gradio-examples {
    border-radius: 10px !important;
}
"""

# --- Model Configs ---
MODEL_NAME = "deepseek-ai/DeepSeek-OCR-2"

MODEL_CONFIGS = {
    "Gundam": {"base_size": 1024, "image_size": 768, "crop_mode": True},
    "Tiny": {"base_size": 512, "image_size": 512, "crop_mode": False},
    "Small": {"base_size": 640, "image_size": 640, "crop_mode": False},
    "Base": {"base_size": 1024, "image_size": 1024, "crop_mode": False},
    "Large": {"base_size": 1280, "image_size": 1280, "crop_mode": False},
}

TASK_PROMPTS = {
    "📋 Markdown": {"prompt": "<image>\n<|grounding|>Convert the document to markdown.", "grounding": True},
    "📝 Free OCR": {"prompt": "<image>\nFree OCR.", "grounding": False},
    "📊 Parse Figure": {"prompt": "<image>\nParse the figure.", "grounding": False},
    "📍 Locate": {"prompt": "<image>\nLocate <|ref|>{text}<|/ref|> in the image.", "grounding": True},
    "🔍 Describe": {"prompt": "<image>\nDescribe this image in detail.", "grounding": False},
    "✏️ Custom": {"prompt": "", "grounding": False},
}

# --- Load Model ---
print("\n🚀 Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModel.from_pretrained(
    MODEL_NAME,
    _attn_implementation="flash_attention_2",
    trust_remote_code=True,
    use_safetensors=True,
).to(device).eval()

model = model.to(torch.bfloat16)


def patch_model_forward(model):
    """Patches Qwen2Decoder2Encoder to handle dynamic resolutions."""
    for name, module in model.named_modules():
        if module.__class__.__name__ == "Qwen2Decoder2Encoder":
            print(f"🛠️ Patching: {name}")

            def fixed_forward(self, x: torch.Tensor) -> torch.Tensor:
                x = x.flatten(2).transpose(1, 2)
                bs, n_query, _ = x.shape

                if n_query == 144:
                    param_img = self.query_768.weight
                elif n_query == 256:
                    param_img = self.query_1024.weight
                else:
                    weight_ref = self.query_1024.weight.permute(1, 0).unsqueeze(0).to(x.dtype)
                    param_img = F.interpolate(weight_ref, size=n_query, mode="linear", align_corners=False)
                    param_img = param_img.squeeze(0).permute(1, 0)

                batch_query_imgs = param_img.unsqueeze(0).expand(bs, -1, -1)
                x_combined = torch.cat([x, batch_query_imgs], dim=1)
                token_type_ids = torch.cat([
                    torch.zeros(bs, n_query, dtype=torch.long, device=x.device),
                    torch.ones(bs, n_query, dtype=torch.long, device=x.device),
                ], dim=1)
                y = self.model(x_combined, token_type_ids)[0]
                return y[:, n_query:, :]

            module.forward = types.MethodType(fixed_forward, module)
            print("✅ Model patched!")
            return

    print("⚠️ Patch target not found.")


patch_model_forward(model)
print("✅ Model ready!\n")


# --- Helper Functions ---
def extract_grounding_refs(text):
    """Extract grounding references from model output."""
    pattern = r'(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)'
    return re.findall(pattern, text, re.DOTALL)


def draw_bounding_boxes(image, refs, extract_crops=False):
    """Draw bounding boxes on image with labels. Returns (annotated_image, crops_list)."""
    img_w, img_h = image.size
    img_draw = image.copy()
    draw = ImageDraw.Draw(img_draw)
    
    # Create overlay for semi-transparent fill
    overlay = Image.new('RGBA', img_draw.size, (0, 0, 0, 0))
    overlay_draw = ImageDraw.Draw(overlay)
    
    # Try to load font, fallback to default
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
    except:
        try:
            font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 18)
        except:
            font = ImageFont.load_default()
    
    crops = []  # List of tuples: (image, caption) for gallery
    color_map = {}
    
    # Tech blue color palette
    tech_colors = [
        (6, 182, 212),    # cyan
        (59, 130, 246),   # blue
        (14, 165, 233),   # sky
        (34, 211, 238),   # cyan light
        (99, 102, 241),   # indigo
        (56, 189, 248),   # sky light
    ]
    color_idx = 0

    for ref in refs:
        label = ref[1]
        if label not in color_map:
            color_map[label] = tech_colors[color_idx % len(tech_colors)]
            color_idx += 1

        color = color_map[label]
        try:
            coords = eval(ref[2])
        except:
            continue
            
        color_alpha = color + (40,)
        box_count = 0

        for box in coords:
            x1 = int(box[0] / 999 * img_w)
            y1 = int(box[1] / 999 * img_h)
            x2 = int(box[2] / 999 * img_w)
            y2 = int(box[3] / 999 * img_h)
            
            # Ensure valid crop area
            x1, x2 = max(0, min(x1, x2)), min(img_w, max(x1, x2))
            y1, y2 = max(0, min(y1, y2)), min(img_h, max(y1, y2))

            # Extract crops for all labels (not just 'image')
            if extract_crops and (x2 - x1) > 10 and (y2 - y1) > 10:
                crop_img = image.crop((x1, y1, x2, y2))
                box_count += 1
                caption = f"{label}" if box_count == 1 else f"{label} ({box_count})"
                crops.append((crop_img, caption))

            # Draw rectangle with tech-style
            width = 3 if label in ['title', 'heading', 'image', 'figure'] else 2
            draw.rectangle([x1, y1, x2, y2], outline=color, width=width)
            overlay_draw.rectangle([x1, y1, x2, y2], fill=color_alpha)

            # Draw label background with rounded feel
            text_bbox = draw.textbbox((0, 0), label, font=font)
            tw, th = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
            label_y = max(0, y1 - th - 6)
            draw.rectangle([x1, label_y, x1 + tw + 8, label_y + th + 6], fill=color)
            draw.text((x1 + 4, label_y + 3), label, font=font, fill=(255, 255, 255))

    img_draw = img_draw.convert('RGBA')
    img_draw = Image.alpha_composite(img_draw, overlay)
    return img_draw.convert('RGB'), crops


def clean_output(text, include_images=False):
    """Clean model output by removing grounding tags."""
    if not text:
        return ""
    
    pattern = r'(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)'
    matches = re.findall(pattern, text, re.DOTALL)
    img_num = 0

    for match in matches:
        if '<|ref|>image<|/ref|>' in match[0]:
            if include_images:
                text = text.replace(match[0], f'\n\n**[Figure {img_num + 1}]**\n\n', 1)
                img_num += 1
            else:
                text = text.replace(match[0], '', 1)
        else:
            text = re.sub(rf'(?m)^[^\n]*{re.escape(match[0])}[^\n]*\n?', '', text)

    return text.strip()


def embed_images_in_markdown(markdown, crops):
    """Embed cropped images as base64 in markdown."""
    if not crops:
        return markdown
    
    # Filter only 'image' type crops for embedding in markdown
    image_crops = []
    for item in crops:
        if isinstance(item, tuple):
            img, caption = item
            if 'image' in caption.lower() or 'figure' in caption.lower():
                image_crops.append(img)
        else:
            image_crops.append(item)
    
    for i, img in enumerate(image_crops):
        buf = BytesIO()
        img.save(buf, format="PNG")
        b64 = base64.b64encode(buf.getvalue()).decode()
        markdown = markdown.replace(
            f'**[Figure {i + 1}]**',
            f'\n\n![Figure {i + 1}](data:image/png;base64,{b64})\n\n',
            1
        )
    return markdown


def load_image_from_file(file_path, page_num=1):
    """Load image from file path (supports PDF)."""
    if not file_path:
        return None
    
    if file_path.lower().endswith('.pdf') and PDF_SUPPORT:
        doc = fitz.open(file_path)
        page_idx = max(0, min(int(page_num) - 1, len(doc) - 1))
        page = doc.load_page(page_idx)
        pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72), alpha=False)
        img = Image.open(BytesIO(pix.tobytes("png")))
        doc.close()
        return img
    else:
        return Image.open(file_path)


def get_pdf_pages(file_path):
    """Get number of pages in PDF."""
    if not file_path or not file_path.lower().endswith('.pdf') or not PDF_SUPPORT:
        return 1
    doc = fitz.open(file_path)
    count = len(doc)
    doc.close()
    return count


# --- Main Processing Function ---
@spaces.GPU(duration=120)
def process_ocr(image, mode, task, custom_prompt):
    """Process image with DeepSeek-OCR-2."""
    if image is None:
        return "⚠️ Please upload an image.", "", "", None, []

    # Validate input for tasks requiring prompt
    if task in ["📍 Locate", "✏️ Custom"] and not custom_prompt.strip():
        msg = "⚠️ Please enter text to locate." if task == "📍 Locate" else "⚠️ Please enter a custom prompt."
        return msg, "", "", None, []

    # Preprocess image
    if image.mode in ('RGBA', 'LA', 'P'):
        image = image.convert('RGB')
    image = ImageOps.exif_transpose(image)

    config = MODEL_CONFIGS[mode]
    task_info = TASK_PROMPTS[task]

    # Build prompt
    if task == "✏️ Custom":
        prompt = f"<image>\n{custom_prompt.strip()}"
        has_grounding = '<|grounding|>' in custom_prompt
    elif task == "📍 Locate":
        prompt = task_info["prompt"].format(text=custom_prompt.strip())
        has_grounding = True
    else:
        prompt = task_info["prompt"]
        has_grounding = task_info["grounding"]

    # Save temp image
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
    image.save(tmp.name, 'JPEG', quality=95)
    tmp.close()
    out_dir = tempfile.mkdtemp()

    # Capture stdout to get clean result
    stdout_backup = sys.stdout
    sys.stdout = StringIO()

    try:
        # Run inference - capture return value directly
        raw_result = model.infer(
            tokenizer=tokenizer,
            prompt=prompt,
            image_file=tmp.name,
            output_path=out_dir,
            base_size=config["base_size"],
            image_size=config["image_size"],
            crop_mode=config["crop_mode"],
            save_results=True,
            eval_mode=True,
            test_compress = True
        )
        
        # Extract text from result
        if isinstance(raw_result, str):
            result = raw_result
        elif isinstance(raw_result, list) and len(raw_result) > 0:
            result = str(raw_result[0])
        else:
            # Fallback: read from saved file
            result = ""
            for filename in os.listdir(out_dir):
                if filename.endswith(('.mmd', '.txt')):
                    with open(os.path.join(out_dir, filename), 'r', encoding='utf-8') as f:
                        result = f.read()
                    break
        
        if not result:
            result = "⚠️ No output generated"
            
    finally:
        os.unlink(tmp.name)
        shutil.rmtree(out_dir, ignore_errors=True)

    if not result:
        return "No text extracted.", "", "", None, []

    # Clean output
    text_clean = clean_output(result, include_images=False)
    markdown_clean = clean_output(result, include_images=True)

    # Process grounding
    img_result = None
    crops = []

    if has_grounding and '<|ref|>' in result:
        refs = extract_grounding_refs(result)
        if refs:
            img_result, crops = draw_bounding_boxes(image, refs, extract_crops=True)

    # Embed crops in markdown
    markdown_with_images = embed_images_in_markdown(markdown_clean, crops)

    return text_clean, markdown_with_images, result, img_result, crops


def process_file(file_path, image, mode, task, custom_prompt, page_num):
    """Process uploaded file or image."""
    if file_path:
        img = load_image_from_file(file_path, int(page_num))
        return process_ocr(img, mode, task, custom_prompt)
    elif image is not None:
        return process_ocr(image, mode, task, custom_prompt)
    return "⚠️ Please upload a file or image.", "", "", None, []


# --- UI Event Handlers ---
def toggle_prompt_visibility(task):
    """Show/hide custom prompt input based on task."""
    if task == "✏️ Custom":
        return gr.update(visible=True, label="Custom Prompt", placeholder="Enter prompt (add <|grounding|> for boxes)")
    elif task == "📍 Locate":
        return gr.update(visible=True, label="Text to Locate", placeholder="Enter text or object to find")
    return gr.update(visible=False)


def update_page_selector(file_path):
    """Update page selector for PDF files."""
    if file_path and file_path.lower().endswith('.pdf') and PDF_SUPPORT:
        pages = get_pdf_pages(file_path)
        return gr.update(visible=True, maximum=pages, value=1, label=f"Page (1-{pages})")
    return gr.update(visible=False)


def switch_to_boxes_tab(task):
    """Switch to boxes tab for grounding tasks."""
    if task in ["📍 Locate"]:
        return gr.update(selected="boxes")
    return gr.update()


# --- Build Interface ---
with gr.Blocks(css=css, theme=theme, title="DeepSeek-OCR-2") as demo:
    # Header
    gr.HTML(
        """
        <div class="main-header">
            <h1>🔮 DeepSeek-OCR-2</h1>
            <p class="subtitle">Powerful document understanding with advanced OCR capabilities</p>
        </div>
    """)

    with gr.Row():
        # Left Column - Input
        with gr.Column(scale=1):
            file_input = gr.File(
                label="📁 Upload File",
                file_types=["image", ".pdf"] if PDF_SUPPORT else ["image"],
                type="filepath",
            )
            
            preview_image = gr.Image(
                label="Preview",
                type="pil",
                height=250,
                sources=["upload", "clipboard"],
            )
            
            page_selector = gr.Number(
                label="Page",
                value=1,
                minimum=1,
                step=1,
                visible=False,
            )

            mode = gr.Dropdown(
                choices=list(MODEL_CONFIGS.keys()),
                value="Gundam",
                label="🎯 Mode",
                info="Gundam: best balance of speed & quality",
            )

            task = gr.Dropdown(
                choices=list(TASK_PROMPTS.keys()),
                value="📋 Markdown",
                label="📌 Task",
            )

            custom_prompt = gr.Textbox(
                label="Prompt",
                lines=2,
                visible=False,
            )

            extract_btn = gr.Button(
                "🚀 Extract",
                variant="primary",
                elem_id="extract-btn",
            )

        # Right Column - Output
        with gr.Column(scale=2):
            with gr.Tabs() as output_tabs:
                with gr.Tab("📝 Text", id="text"):
                    text_output = gr.Textbox(
                        lines=18,
                        show_copy_button=True,
                        show_label=False,
                        elem_classes=["output-text"],
                        placeholder="Extracted text will appear here...",
                    )

                with gr.Tab("📄 Markdown", id="markdown"):
                    markdown_output = gr.Markdown()

                with gr.Tab("🎯 Layout", id="boxes"):
                    boxes_output = gr.Image(
                        type="pil",
                        height=450,
                        show_label=False,
                    )

                with gr.Tab("🖼️ Crops", id="crops"):
                    crops_gallery = gr.Gallery(
                        show_label=False,
                        columns=3,
                        height=400,
                    )

                with gr.Tab("📜 Raw", id="raw"):
                    raw_output = gr.Textbox(
                        lines=18,
                        show_copy_button=True,
                        show_label=False,
                        elem_classes=["output-text"],
                    )

    # Examples section - full width
    gr.Markdown("### 📂 Examples")
    gr.Examples(
        examples=[
            ["examples/1.jpg", "Gundam", "📋 Markdown", ""],
            ["examples/2.jpg", "Gundam", "📍 Locate", "Robot"],
            ["examples/3.jpg", "Base", "📝 Free OCR", ""],
        ],
        inputs=[preview_image, mode, task, custom_prompt],
        label="",
        cache_examples=False,
        examples_per_page=6,
    )

    # Info section
    with gr.Accordion("ℹ️ Info & Tips", open=False):
        gr.Markdown("""
### 🐳 DeepSeek-OCR-2: Visual Causal Flow

**DeepEncoder V2** với causal reasoning cho document phức tạp, tables, formulas.

### 🎯 Strategies

- **Gundam (Recommended):** base_size=1024, image_size=640, crop_mode=True - Best balance of speed & quality
- **Global View (Fast):** base_size=1024, image_size=1024, crop_mode=False - Fast, text sparse
- **Tiny:** base_size=512, image_size=512, crop_mode=False - Quick previews
- **Small:** base_size=640, image_size=640, crop_mode=False - Simple documents
- **Base:** base_size=1024, image_size=1024, crop_mode=False - Standard quality
- **Large:** base_size=1280, image_size=1280, crop_mode=False - High detail

### 📌 Tasks

- **📋 Markdown**: Convert to structured markdown with layout detection
- **📝 Free OCR**: Simple text extraction
- **📊 Parse Figure**: Extract data from charts/graphs
- **📍 Locate**: Find specific text/objects (shows bounding boxes)
- **🔍 Describe**: General image description
- **✏️ Custom**: Your own prompt (add `<|grounding|>` for boxes)

### 💡 Tips

- Use **Gundam** mode for most documents (best balance)
- Check **Layout** tab to see detected regions with bounding boxes
- **Crops** tab shows extracted images/figures from the document
- For complex documents with tables/formulas, Gundam mode works best

**Please give a like ❤️ if this is helpful!**
        """)

    # Event handlers
    file_input.change(load_image_from_file, [file_input, page_selector], [preview_image])
    file_input.change(update_page_selector, [file_input], [page_selector])
    page_selector.change(load_image_from_file, [file_input, page_selector], [preview_image])
    task.change(toggle_prompt_visibility, [task], [custom_prompt])
    task.change(switch_to_boxes_tab, [task], [output_tabs])

    extract_btn.click(
        process_file,
        inputs=[file_input, preview_image, mode, task, custom_prompt, page_selector],
        outputs=[text_output, markdown_output, raw_output, boxes_output, crops_gallery],
    )

# --- Launch ---
if __name__ == "__main__":
    os.makedirs("examples", exist_ok=True)
    demo.queue(max_size=20).launch(
        share=True,
        mcp_server=True,
        ssr_mode=False,
    )