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 time
from threading import Thread

import gradio as gr
import requests
import torch
from PIL import Image
from transformers import (
    Qwen2_5_VLForConditionalGeneration,
    AutoProcessor,
    TextIteratorStreamer,
)
from qwen_vl_utils import process_vision_info

# Constants
MIN_PIXELS = 3136
MAX_PIXELS = 11289600
IMAGE_FACTOR = 28
MAX_INPUT_TOKEN_LENGTH = 2048
device = "cuda" if torch.cuda.is_available() else "cpu"

# Prompts
prompt = """Please output the layout information from the 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.
"""

# Load models
MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-062825"
processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)
model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_M,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to(device).eval()

MODEL_ID_T = "prithivMLmods/Megalodon-OCR-Sync-0713"
processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True)
model_t = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_T,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to(device).eval()

MODEL_ID_C = "nanonets/Nanonets-OCR-s"
processor_c = AutoProcessor.from_pretrained(MODEL_ID_C, trust_remote_code=True)
model_c = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_C,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to(device).eval()

MODEL_ID_G = "echo840/MonkeyOCR"
SUBFOLDER = "Recognition"
processor_g = AutoProcessor.from_pretrained(
    MODEL_ID_G,
    trust_remote_code=True,
    subfolder=SUBFOLDER
)
model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_G,
    trust_remote_code=True,
    subfolder=SUBFOLDER,
    torch_dtype=torch.float16
).to(device).eval()

# Utility functions
def round_by_factor(number: int, factor: int) -> int:
    return round(number / factor) * factor

def smart_resize(
    height: int,
    width: int,
    factor: int = 28,
    min_pixels: int = 3136,
    max_pixels: int = 11289600,
):
    if max(height, width) / min(height, width) > 200:
        raise ValueError(f"Aspect ratio too extreme: {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):
    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 or max_pixels:
        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 is_arabic_text(text: str) -> bool:
    if not text:
        return False
    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
        header_match = re.match(header_pattern, line, re.MULTILINE)
        if header_match:
            content_text.append(header_match.group(1))
            continue
        if re.match(paragraph_pattern, line, re.MULTILINE):
            content_text.append(line)
    if not content_text:
        return False
    combined_text = ' '.join(content_text)
    arabic_chars = 0
    total_chars = 0
    for char in combined_text:
        if char.isalpha():
            total_chars += 1
            if ('\u0600' <= char <= '\u06FF') or ('\u0750' <= char <= '\u077F') or ('\u08A0' <= char <= '\u08FF'):
                arabic_chars += 1
    return total_chars > 0 and (arabic_chars / total_chars) > 0.5

def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
    import base64
    from io import BytesIO
    markdown_lines = []
    try:
        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':
                if bbox and len(bbox) == 4:
                    try:
                        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 and y2 > y1:
                            cropped_img = image.crop((x1, y1, x2, y2))
                            buffer = BytesIO()
                            cropped_img.save(buffer, format='PNG')
                            img_data = base64.b64encode(buffer.getvalue()).decode()
                            markdown_lines.append(f"<image-card alt="Image" src="data:image/png;base64,{img_data}" ></image-card>\n")
                        else:
                            markdown_lines.append("<image-card alt="Image" src="Image region detected" ></image-card>\n")
                    except Exception as e:
                        print(f"Error processing image region: {e}")
                        markdown_lines.append("<image-card alt="Image" src="Image detected" ></image-card>\n")
                else:
                    markdown_lines.append("<image-card alt="Image" src="Image detected" ></image-card>\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.strip().startswith('<'):
                    markdown_lines.append(f"{text}\n")
                else:
                    markdown_lines.append(f"**Table:** {text}\n")
            elif category == 'Formula':
                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']:
                continue
            else:
                markdown_lines.append(f"{text}\n")
            markdown_lines.append("")
    except Exception as e:
        print(f"Error converting to markdown: {e}")
        return str(layout_data)
    return "\n".join(markdown_lines)

@spaces.GPU
def inference(model_name: str, image: Image.Image, text: str, max_new_tokens: int = 1024) -> str:
    try:
        if model_name == "Camel-Doc-OCR-062825":
            processor = processor_m
            model = model_m
        elif model_name == "Megalodon-OCR-Sync-0713":
            processor = processor_t
            model = model_t
        elif model_name == "Nanonets-OCR-s":
            processor = processor_c
            model = model_c
        elif model_name == "MonkeyOCR-Recognition":
            processor = processor_g
            model = model_g
        else:
            raise ValueError(f"Invalid model selected: {model_name}")

        if image is None:
            yield "Please upload an image.", "Please upload an image."
            return

        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": text},
            ]
        }]
        prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(
            text=[prompt_full],
            images=[image],
            return_tensors="pt",
            padding=True,
            truncation=False,
            max_length=MAX_INPUT_TOKEN_LENGTH
        ).to(device)
        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
        generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()
        buffer = ""
        for new_text in streamer:
            buffer += new_text
            buffer = buffer.replace("<|im_end|>", "")
            time.sleep(0.01)
            yield buffer, buffer
    except Exception as e:
        print(f"Error during inference: {e}")
        traceback.print_exc()
        yield f"Error during inference: {str(e)}", f"Error during inference: {str(e)}"

def process_image(
    model_name: str,
    image: Image.Image,
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None,
    max_new_tokens: int = 1024
):
    try:
        if min_pixels or max_pixels:
            image = fetch_image(image, min_pixels=min_pixels, max_pixels=max_pixels)
        buffer = ""
        for raw_output, _ in inference(model_name, image, prompt, max_new_tokens):
            buffer = raw_output
            yield buffer, None  # Yield raw OCR stream and None for JSON during processing
        try:
            json_match = re.search(r'```json
            json_str = json_match.group(1) if json_match else buffer
            layout_data = json.loads(json_str)
            yield buffer, layout_data  # Final yield with raw OCR and parsed JSON
        except json.JSONDecodeError:
            print("Failed to parse JSON output, using raw output")
            yield buffer, None  # If JSON parsing fails, yield raw OCR with no JSON
    except Exception as e:
        print(f"Error processing image: {e}")
        traceback.print_exc()
        yield f"Error processing image: {str(e)}", None

def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
    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 in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
            image = Image.open(file_path).convert('RGB')
            return image, "Image loaded"
        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 create_gradio_interface():
    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; 
        background-color: blue !important;} 
    .process-button:hover { 
        background-color: darkblue !important;
        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="bethecloud/storj_theme", css=css) as demo:
        gr.HTML("""
        <div class="title" style="text-align: center">
            <h1>Dot<span style="color: red;">●</span><strong></strong>OCR Comparator</h1>
            <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
                Advanced vision-language model for image to markdown document processing
            </p>
        </div>
        """)
        with gr.Row():
            with gr.Column(scale=1):
                model_choice = gr.Radio(
                    choices=["Camel-Doc-OCR-062825", "MonkeyOCR-Recognition", "Nanonets-OCR-s", "Megalodon-OCR-Sync-0713"],
                    label="Select Model",
                    value="Camel-Doc-OCR-062825"
                )
                file_input = gr.File(
                    label="Upload Image",
                    file_types  =[".jpg", ".jpeg", ".png", ".bmp", ".tiff"],
                    type="filepath"
                )
                image_preview = gr.Image(label="Preview", type="pil", interactive=False, height=300)
                with gr.Accordion("Advanced Settings", open=False):
                    max_new_tokens = gr.Slider(minimum=1000, maximum=32000, value=24000, step=1000, label="Max New Tokens")
                    min_pixels = gr.Number(value=MIN_PIXELS, label="Min Pixels")
                    max_pixels = gr.Number(value=MAX_PIXELS, label="Max Pixels")
                process_btn = gr.Button("🚀 Process Document", variant="primary", elem_classes=["process-button"], size="lg")
                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
            with gr.Column(scale=2):
                with gr.Tabs():
                    with gr.Tab("📝 Extracted Content"):
                        output = gr.Textbox(label="Raw OCR Stream", interactive=False, lines=10, show_copy_button=True)
                    with gr.Tab("📋 Layout Analysis Results"):
                        json_output = gr.JSON(label="Layout Analysis Results", value=None)
        def process_document(model_name, file_path, max_tokens, min_pix, max_pix):
            try:
                if not file_path:
                    return "Please upload an image.", None
                image, status = load_file_for_preview(file_path)
                if image is None:
                    return status, None
                for raw_output, layout_result in process_image(model_name, image, min_pixels=int(min_pix) if min_pix else None, max_pixels=int(max_pix) if max_pix else None, max_new_tokens=max_tokens):
                    yield raw_output, layout_result
            except Exception as e:
                error_msg = f"Error processing document: {str(e)}"
                print(error_msg)
                traceback.print_exc()
                yield error_msg, None
        def handle_file_upload(file_path):
            if not file_path:
                return None, "No file loaded"
            image, page_info = load_file_for_preview(file_path)
            return image, page_info
        def clear_all():
            return None, None, "No file loaded", None
        file_input.change(handle_file_upload, inputs=[file_input], outputs=[image_preview, output])
        process_btn.click(
            process_document,
            inputs=[model_choice, file_input, max_new_tokens, min_pixels, max_pixels],
            outputs=[output, json_output]
        )
        clear_btn.click(
            clear_all,
            outputs=[file_input, image_preview, output, json_output]
        )
    return demo

if __name__ == "__main__":
    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)