import gradio as gr
from transformers import AutoModel, AutoTokenizer
import torch
import spaces
import os
import sys
import tempfile
import shutil
from PIL import Image, ImageDraw, ImageFont, ImageOps
import fitz
import re
import warnings
import numpy as np
import base64
from io import StringIO, BytesIO

MODEL_NAME = 'deepseek-ai/DeepSeek-OCR-2'

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL_NAME, _attn_implementation='flash_attention_2', torch_dtype=torch.bfloat16, trust_remote_code=True, use_safetensors=True)
model = model.eval().cuda()

MODEL_CONFIGS = {
    "Default": {"base_size": 1024, "image_size": 768, "crop_mode": True},
    "Quality": {"base_size": 1280, "image_size": 960, "crop_mode": True},
    "Fast": {"base_size": 1024, "image_size": 640, "crop_mode": True},
    "No Crop": {"base_size": 1024, "image_size": 768, "crop_mode": False},
    "Small": {"base_size": 768, "image_size": 512, "crop_mode": False},
}

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

def extract_grounding_references(text):
    pattern = r'(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)'
    return re.findall(pattern, text, re.DOTALL)

def draw_bounding_boxes(image, refs, extract_images=False):
    img_w, img_h = image.size
    img_draw = image.copy()
    draw = ImageDraw.Draw(img_draw)
    overlay = Image.new('RGBA', img_draw.size, (0, 0, 0, 0))
    draw2 = ImageDraw.Draw(overlay)
    font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 15)
    crops = []
    
    color_map = {}
    np.random.seed(42)

    for ref in refs:
        label = ref[1]
        if label not in color_map:
            color_map[label] = (np.random.randint(50, 255), np.random.randint(50, 255), np.random.randint(50, 255))

        color = color_map[label]
        coords = eval(ref[2])
        color_a = color + (60,)
        
        for box in coords:
            x1, y1, x2, y2 = int(box[0]/999*img_w), int(box[1]/999*img_h), int(box[2]/999*img_w), int(box[3]/999*img_h)
            
            if extract_images and label == 'image':
                crops.append(image.crop((x1, y1, x2, y2)))
            
            width = 5 if label == 'title' else 3
            draw.rectangle([x1, y1, x2, y2], outline=color, width=width)
            draw2.rectangle([x1, y1, x2, y2], fill=color_a)
            
            text_bbox = draw.textbbox((0, 0), label, font=font)
            tw, th = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
            ty = max(0, y1 - 20)
            draw.rectangle([x1, ty, x1 + tw + 4, ty + th + 4], fill=color)
            draw.text((x1 + 2, ty + 2), label, font=font, fill=(255, 255, 255))
    
    img_draw.paste(overlay, (0, 0), overlay)
    return img_draw, crops

def clean_output(text, include_images=False):
    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)
    
    text = text.replace('\\coloneqq', ':=').replace('\\eqqcolon', '=:')
    
    return text.strip()

def embed_images(markdown, crops):
    if not crops:
        return markdown
    for i, img in enumerate(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

@spaces.GPU(duration=90)
def process_image(image, mode, task, custom_prompt):
    if image is None:
        return "Error: Upload an image", "", "", None, []
    if task in ["✏️ Custom", "📍 Locate"] and not custom_prompt.strip():
        return "Please enter a prompt", "", "", None, []
    
    if image.mode in ('RGBA', 'LA', 'P'):
        image = image.convert('RGB')
    image = ImageOps.exif_transpose(image)
    
    config = MODEL_CONFIGS[mode]
    
    if task == "✏️ Custom":
        prompt = f"<image>\n{custom_prompt.strip()}"
        has_grounding = '<|grounding|>' in custom_prompt
    elif task == "📍 Locate":
        prompt = f"<image>\nLocate <|ref|>{custom_prompt.strip()}<|/ref|> in the image."
        has_grounding = True
    else:
        prompt = TASK_PROMPTS[task]["prompt"]
        has_grounding = TASK_PROMPTS[task]["has_grounding"]
    
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
    image.save(tmp.name, 'JPEG', quality=95)
    tmp.close()
    out_dir = tempfile.mkdtemp()
    
    stdout = sys.stdout
    sys.stdout = StringIO()
    
    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=False
    )
    
    result = '\n'.join([l for l in sys.stdout.getvalue().split('\n') 
                        if not any(s in l for s in ['image:', 'other:', 'PATCHES', '====', 'BASE:', '%|', 'torch.Size'])]).strip()
    sys.stdout = stdout
    
    os.unlink(tmp.name)
    shutil.rmtree(out_dir, ignore_errors=True)
    
    if not result:
        return "No text detected", "", "", None, []
    
    cleaned = clean_output(result, False)
    markdown = clean_output(result, True)
    
    img_out = None
    crops = []
    
    if has_grounding and '<|ref|>' in result:
        refs = extract_grounding_references(result)
        if refs:
            img_out, crops = draw_bounding_boxes(image, refs, True)
    
    markdown = embed_images(markdown, crops)
    
    return cleaned, markdown, result, img_out, crops

@spaces.GPU(duration=90)
def process_pdf(path, mode, task, custom_prompt, page_num):
    doc = fitz.open(path)
    total_pages = len(doc)
    if page_num < 1 or page_num > total_pages:
        doc.close()
        return f"Invalid page number. PDF has {total_pages} pages.", "", "", None, []
    page = doc.load_page(page_num - 1)
    pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72), alpha=False)
    img = Image.open(BytesIO(pix.tobytes("png")))
    doc.close()
    
    return process_image(img, mode, task, custom_prompt)

def process_file(path, mode, task, custom_prompt, page_num):
    if not path:
        return "Error: Upload a file", "", "", None, []
    if path.lower().endswith('.pdf'):
        return process_pdf(path, mode, task, custom_prompt, page_num)
    else:
        return process_image(Image.open(path), mode, task, custom_prompt)

def toggle_prompt(task):
    if task == "✏️ Custom":
        return gr.update(visible=True, label="Custom Prompt", placeholder="Add <|grounding|> for bounding boxes")
    elif task == "📍 Locate":
        return gr.update(visible=True, label="Text to Locate", placeholder="Enter text to locate")
    return gr.update(visible=False)

def select_boxes(task):
    if task == "📍 Locate":
        return gr.update(selected="tab_boxes")
    return gr.update()

def get_pdf_page_count(file_path):
    if not file_path or not file_path.lower().endswith('.pdf'):
        return 1
    doc = fitz.open(file_path)
    count = len(doc)
    doc.close()
    return count

def load_image(file_path, page_num=1):
    if not file_path:
        return None
    if file_path.lower().endswith('.pdf'):
        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 update_page_selector(file_path):
    if not file_path:
        return gr.update(visible=False)
    if file_path.lower().endswith('.pdf'):
        page_count = get_pdf_page_count(file_path)
        return gr.update(visible=True, maximum=page_count, value=1, minimum=1,
                        label=f"Select Page (1-{page_count})")
    return gr.update(visible=False)

with gr.Blocks(title="DeepSeek-OCR-2") as demo:
    gr.Markdown("""
    # 🚀 DeepSeek-OCR-2 Demo
    **Convert documents to markdown, extract text, parse figures, and locate specific content with bounding boxes.**
    
    Powered by **DeepEncoder V2** - a novel LLM-style vision encoder that dynamically reorders visual tokens based on semantic understanding, mimicking human reading patterns instead of rigid left-to-right scanning. Achieves **91.09%** on OmniDocBench (+3.73% over v1).
    
    **If this tool was helpful, please consider giving it a like ❤️!**
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            file_in = gr.File(label="Upload Image or PDF", file_types=["image", ".pdf"], type="filepath")
            input_img = gr.Image(label="Input Image", type="pil", height=300)
            page_selector = gr.Number(label="Select Page", value=1, minimum=1, step=1, visible=False)
            mode = gr.Dropdown(list(MODEL_CONFIGS.keys()), value="Default", label="Mode")
            task = gr.Dropdown(list(TASK_PROMPTS.keys()), value="📋 Markdown", label="Task")
            prompt = gr.Textbox(label="Prompt", lines=2, visible=False)
            btn = gr.Button("Extract", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            with gr.Tabs() as tabs:
                with gr.Tab("Text", id="tab_text"):
                    text_out = gr.Textbox(lines=20, buttons=["copy"], show_label=False)
                with gr.Tab("Markdown Preview", id="tab_markdown"):
                    md_out = gr.Markdown("")
                with gr.Tab("Boxes", id="tab_boxes"):
                    img_out = gr.Image(type="pil", height=500, show_label=False)
                with gr.Tab("Cropped Images", id="tab_crops"):
                    gallery = gr.Gallery(show_label=False, columns=3, height=400)
                with gr.Tab("Raw Text", id="tab_raw"):
                    raw_out = gr.Textbox(lines=20, buttons=["copy"], show_label=False)
    
    gr.Examples(
        examples=[
            ["examples/ocr.jpg", "Default", "📋 Markdown", ""],
            ["examples/reachy-mini.jpg", "Default", "📍 Locate", "Robot"]
        ],
        inputs=[input_img, mode, task, prompt],
        cache_examples=False
    )
    
    with gr.Accordion("ℹ️ Info", open=False):
        gr.Markdown("""
        ### Modes
        - **Default**: 1024 base + 768 tiles with cropping - Recommended for most use cases
        - **Quality**: 1280 base + 960 tiles with cropping - Higher quality, slower
        - **Fast**: 1024 base + 640 tiles with cropping - Faster processing
        - **No Crop**: 1024 base + 768 tiles without cropping - Single image processing
        - **Small**: 768 base + 512 tiles without cropping - Fastest, lower quality
        
        ### Tasks
        - **Markdown**: Convert document to structured markdown with layout detection (grounding ✅)
        - **Free OCR**: Simple text extraction without layout
        - **OCR Image**: OCR for general images with grounding (grounding ✅)
        - **Parse Figure**: Parse figures and charts in documents
        - **Locate**: Find and highlight specific text/elements in image (grounding ✅)
        - **Describe**: General image description
        - **Custom**: Your own prompt (add `<|grounding|>` for bounding boxes)
        
        Document: <image>\\n<|grounding|>Convert the document to markdown.
        Free OCR: <image>\\nFree OCR.
        Other Image: <image>\\n<|grounding|>OCR this image.
        Parse Figure: <image>\\nParse the figure.
        Describe: <image>\\nDescribe this image in detail.
        Locate: <image>\\nLocate <|ref|>text<|/ref|> in the image.
        ```
        """)
    
    file_in.change(load_image, [file_in, page_selector], [input_img])
    file_in.change(update_page_selector, [file_in], [page_selector])
    page_selector.change(load_image, [file_in, page_selector], [input_img])
    task.change(toggle_prompt, [task], [prompt])
    task.change(select_boxes, [task], [tabs])
    
    def run(image, file_path, mode, task, custom_prompt, page_num):
        if file_path:
            return process_file(file_path, mode, task, custom_prompt, int(page_num))
        if image is not None:
            return process_image(image, mode, task, custom_prompt)
        return "Error: Upload a file or image", "", "", None, []

    submit_event = btn.click(run, [input_img, file_in, mode, task, prompt, page_selector],
                             [text_out, md_out, raw_out, img_out, gallery])
    submit_event.then(select_boxes, [task], [tabs])

if __name__ == "__main__":
    demo.queue(max_size=20).launch(theme=gr.themes.Soft())