import gradio as gr
import torch
from transformers import AutoModel, AutoTokenizer, AutoConfig
import os
import base64
import spaces
import io
from PIL import Image
import numpy as np
import yaml
from pathlib import Path
from globe import title, description, modelinfor, joinus, howto
import uuid
import tempfile
import time
import shutil
import cv2
import re
import warnings

# Try to import spaces module for ZeroGPU compatibility
try:
    import spaces
    SPACES_AVAILABLE = True
except ImportError:
    SPACES_AVAILABLE = False
    # Create a dummy decorator for local development
    def dummy_gpu_decorator(func):
        return func
    spaces = type('spaces', (), {'GPU': dummy_gpu_decorator})()

# Suppress specific warnings that are known issues with GOT-OCR
warnings.filterwarnings("ignore", message="The attention mask and the pad token id were not set")
warnings.filterwarnings("ignore", message="Setting `pad_token_id` to `eos_token_id`")
warnings.filterwarnings("ignore", message="The attention mask is not set and cannot be inferred")
warnings.filterwarnings("ignore", message="The `seen_tokens` attribute is deprecated")


def initialize_model_safely():
    """
    Safely initialize the GOT-OCR model with proper error handling for ZeroGPU
    """
    model_name = 'ucaslcl/GOT-OCR2_0'
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    try:
        # Initialize tokenizer with proper settings
        tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
        
        # Set pad token properly
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        
        config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
        
        model = AutoModel.from_pretrained(
            'ucaslcl/GOT-OCR2_0', 
            trust_remote_code=True, 
            low_cpu_mem_usage=True,  
            device_map=device, 
            use_safetensors=True,
            pad_token_id=tokenizer.eos_token_id,
            use_cache=True,
            torch_dtype=torch.float16 if device == 'cuda' else torch.float32
        )
        
        model = model.eval().to(device)
        model.config.pad_token_id = tokenizer.eos_token_id
        
        # Ensure the model has proper tokenizer settings
        if hasattr(model, 'config'):
            model.config.pad_token_id = tokenizer.eos_token_id
            model.config.eos_token_id = tokenizer.eos_token_id
            
        return model, tokenizer
        
    except Exception as e:
        print(f"Error initializing model: {str(e)}")
        # Fallback initialization
        try:
            tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token
                
            model = AutoModel.from_pretrained(
                'ucaslcl/GOT-OCR2_0', 
                trust_remote_code=True, 
                low_cpu_mem_usage=True,  
                device_map=device, 
                use_safetensors=True
            )
            model = model.eval().to(device)
            return model, tokenizer
        except Exception as fallback_error:
            raise Exception(f"Failed to initialize model: {str(e)}. Fallback also failed: {str(fallback_error)}")

model, tokenizer = initialize_model_safely()

UPLOAD_FOLDER = "./uploads"
RESULTS_FOLDER = "./results"

for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
    if not os.path.exists(folder):
        os.makedirs(folder)

def image_to_base64(image):
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

def safe_model_chat(model, tokenizer, image_path, **kwargs):
    """
    Safe wrapper for model.chat to handle DynamicCache and other compatibility issues
    Optimized for ZeroGPU environments
    """
    try:
        # First attempt: normal call
        return model.chat(tokenizer, image_path, **kwargs)
    except AttributeError as e:
        if "get_max_length" in str(e):
            # Try to fix the cache issue by clearing it
            try:
                if hasattr(model, 'clear_cache'):
                    model.clear_cache()
                # Retry the call
                return model.chat(tokenizer, image_path, **kwargs)
            except:
                # If still failing, try with different parameters
                try:
                    # Remove any cache-related parameters
                    kwargs_copy = kwargs.copy()
                    if 'use_cache' in kwargs_copy:
                        del kwargs_copy['use_cache']
                    return model.chat(tokenizer, image_path, **kwargs_copy)
                except:
                    raise Exception("Model compatibility issue: DynamicCache error. Please try again.")
        else:
            raise e
    except Exception as e:
        # Handle other potential issues
        if "attention_mask" in str(e).lower():
            # Try to handle attention mask issues
            try:
                return model.chat(tokenizer, image_path, **kwargs)
            except:
                raise Exception(f"Attention mask error: {str(e)}")
        else:
            raise e

def safe_model_chat_crop(model, tokenizer, image_path, **kwargs):
    """
    Safe wrapper for model.chat_crop to handle DynamicCache and other compatibility issues
    Optimized for ZeroGPU environments
    """
    try:
        # First attempt: normal call
        return model.chat_crop(tokenizer, image_path, **kwargs)
    except AttributeError as e:
        if "get_max_length" in str(e):
            # Try to fix the cache issue by clearing it
            try:
                if hasattr(model, 'clear_cache'):
                    model.clear_cache()
                # Retry the call
                return model.chat_crop(tokenizer, image_path, **kwargs)
            except:
                # If still failing, try with different parameters
                try:
                    # Remove any cache-related parameters
                    kwargs_copy = kwargs.copy()
                    if 'use_cache' in kwargs_copy:
                        del kwargs_copy['use_cache']
                    return model.chat_crop(tokenizer, image_path, **kwargs_copy)
                except:
                    raise Exception("Model compatibility issue: DynamicCache error. Please try again.")
        else:
            raise e
    except Exception as e:
        # Handle other potential issues
        if "attention_mask" in str(e).lower():
            # Try to handle attention mask issues
            try:
                return model.chat_crop(tokenizer, image_path, **kwargs)
            except:
                raise Exception(f"Attention mask error: {str(e)}")
        else:
            raise e


@spaces.GPU()
def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
    if image is None:
        return "Error: No image provided", None, None
    
    unique_id = str(uuid.uuid4())
    image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.png")
    result_path = os.path.join(RESULTS_FOLDER, f"{unique_id}.html")
    
    try:
        if isinstance(image, dict):
            composite_image = image.get("composite")
            if composite_image is not None:
                if isinstance(composite_image, np.ndarray):
                    cv2.imwrite(image_path, cv2.cvtColor(composite_image, cv2.COLOR_RGB2BGR))
                elif isinstance(composite_image, Image.Image):
                    composite_image.save(image_path)
                else:
                    return "Error: Unsupported image format from ImageEditor", None, None
            else:
                return "Error: No composite image found in ImageEditor output", None, None
        elif isinstance(image, np.ndarray):
            cv2.imwrite(image_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
        elif isinstance(image, str):
            shutil.copy(image, image_path)
        else:
            return "Error: Unsupported image format", None, None

        # Wrap model calls in try-except to handle DynamicCache errors
        try:
            if task == "Plain Text OCR":
                res = safe_model_chat(model, tokenizer, image_path, ocr_type='ocr')
                return res, None, unique_id
            else:
                if task == "Format Text OCR":
                    res = safe_model_chat(model, tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
                elif task == "Fine-grained OCR (Box)":
                    res = safe_model_chat(model, tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
                elif task == "Fine-grained OCR (Color)":
                    res = safe_model_chat(model, tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
                elif task == "Multi-crop OCR":
                    res = safe_model_chat_crop(model, tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
                elif task == "Render Formatted OCR":
                    res = safe_model_chat(model, tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
                
                if os.path.exists(result_path):
                    with open(result_path, 'r') as f:
                        html_content = f.read()
                    return res, html_content, unique_id
                else:
                    return res, None, unique_id
        except AttributeError as e:
            if "get_max_length" in str(e):
                # Handle DynamicCache compatibility issue
                return "Error: Model compatibility issue detected. Please try again or contact support.", None, None
            else:
                raise e
                
    except Exception as e:
        return f"Error: {str(e)}", None, None
    finally:
        if os.path.exists(image_path):
            os.remove(image_path)
            
def update_image_input(task):
    if task == "Fine-grained OCR (Color)":
        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
    else:
        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)

def update_inputs(task):
    if task in ["Plain Text OCR", "Format Text OCR", "Multi-crop OCR", "Render Formatted OCR"]:
        return [
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=False)
        ]
    elif task == "Fine-grained OCR (Box)":
        return [
            gr.update(visible=True, choices=["ocr", "format"]),
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=False)
        ]
    elif task == "Fine-grained OCR (Color)":
        return [
            gr.update(visible=True, choices=["ocr", "format"]),
            gr.update(visible=False),
            gr.update(visible=True, choices=["red", "green", "blue"]),
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=True)
        ]
    
def parse_latex_output(res):
    # Split the input, preserving newlines and empty lines
    lines = re.split(r'(\$\$.*?\$\$)', res, flags=re.DOTALL)
    parsed_lines = []
    in_latex = False
    latex_buffer = []

    for line in lines:
        if line == '\n':
            if in_latex:
                latex_buffer.append(line)
            else:
                parsed_lines.append(line)
            continue

        line = line.strip()
        
        latex_patterns = [r'\{', r'\}', r'\[', r'\]', r'\\', r'\$', r'_', r'^', r'"']
        contains_latex = any(re.search(pattern, line) for pattern in latex_patterns)
        
        if contains_latex:
            if not in_latex:
                in_latex = True
                latex_buffer = ['$$']
            latex_buffer.append(line)
        else:
            if in_latex:
                latex_buffer.append('$$')
                parsed_lines.extend(latex_buffer)
                in_latex = False
                latex_buffer = []
            parsed_lines.append(line)

    if in_latex:
        latex_buffer.append('$$')
        parsed_lines.extend(latex_buffer)

    return '$$\\$$\n'.join(parsed_lines)
                         

def ocr_demo(image, task, ocr_type, ocr_box, ocr_color):
    """
    Main OCR demonstration function that processes images and returns results.
    
    Args:
        image (Union[dict, np.ndarray, str, PIL.Image]): Input image in one of these formats: Image component state with keys: path: str | None (Path to local file) url: str | None (Public URL or base64 image) size: int | None (Image size in bytes) orig_name: str | None (Original filename) mime_type: str | None (Image MIME type) is_stream: bool (Always False) meta: dict(str, Any) OR  dict: ImageEditor component state with keys: background: filepath | None layers: list[filepath] composite: filepath | None id: str | None OR np.ndarray: Raw image array str: Path to image file PIL.Image: PIL Image object
        task (Literal['Plain Text OCR', 'Format Text OCR', 'Fine-grained OCR (Box)', 'Fine-grained OCR (Color)', 'Multi-crop OCR', 'Render Formatted OCR'], default: "Plain Text OCR"): The type of OCR processing to perform: "Plain Text OCR": Basic text extraction without formatting, "Format Text OCR": Text extraction with preserved formatting, "Fine-grained OCR (Box)": Text extraction from specific bounding box regions, "Fine-grained OCR (Color)": Text extraction from regions marked with specific colors, "Multi-crop OCR": Text extraction from multiple cropped regions, "Render Formatted OCR": Text extraction with HTML rendering of formatting
        ocr_type (Literal['ocr', 'format'], default: "ocr"):The type of OCR processing to apply: "ocr": Basic text extraction without formatting "format": Text extraction with preserved formatting and structure
        ocr_box (str): Bounding box coordinates specifying the region for fine-grained OCR. Format: "x1,y1,x2,y2" where: x1,y1: Top-left corner coordinates ; x2,y2: Bottom-right corner coordinates Example: "100,100,300,200" for a box starting at (100,100) and ending at (300,200)
        ocr_color (Literal['red', 'green', 'blue'], default: "red"): Color specification for fine-grained OCR when using color-based region selection: "red": Extract text from regions marked in red "green": Extract text from regions marked in green "blue": Extract text from regions marked in blue
        
    Returns:
        tuple: (formatted_result, html_output)
            - formatted_result (str): Formatted OCR result text
            - html_output (str): HTML visualization if applicable
    """
    res, html_content, unique_id = process_image(image, task, ocr_type, ocr_box, ocr_color)
    
    if isinstance(res, str) and res.startswith("Error:"):
        return res, None

    res = res.replace("\\title", "\\title ")
    formatted_res = res
    # formatted_res = parse_latex_output(res)
    
    if html_content:
        encoded_html = base64.b64encode(html_content.encode('utf-8')).decode('utf-8')
        iframe_src = f"data:text/html;base64,{encoded_html}"
        iframe = f'<iframe src="{iframe_src}" width="100%" height="600px"></iframe>'
        download_link = f'<a href="data:text/html;base64,{encoded_html}" download="result_{unique_id}.html">Download Full Result</a>'
        return formatted_res, f"{download_link}<br>{iframe}"
    return formatted_res, None

def cleanup_old_files():
    current_time = time.time()
    for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
        for file_path in Path(folder).glob('*'):
            if current_time - file_path.stat().st_mtime > 3600:  # 1 hour
                file_path.unlink()

with gr.Blocks(theme=gr.themes.Base()) as demo:
    with gr.Row():
        gr.Markdown(title)
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group():                    
                gr.Markdown(description)
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown(modelinfor)
                gr.Markdown(joinus)
    with gr.Row():
        with gr.Accordion("How to use Fine-grained OCR (Color)", open=False):
            with gr.Row():
                gr.Image("res/image/howto_1.png", label="Select the Following Parameters")
                gr.Image("res/image/howto_2.png", label="Click on Paintbrush in the Image Editor")
                gr.Image("res/image/howto_3.png", label="Select your Brush Color (Red)")
                gr.Image("res/image/howto_4.png", label="Make a Box Around The Text")
            with gr.Row():
                with gr.Group():
                    gr.Markdown(howto)

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group():
                image_input = gr.Image(type="filepath", label="Input Image")
                image_editor = gr.ImageEditor(label="Image Editor", type="pil", visible=False)
                task_dropdown = gr.Dropdown(
                    choices=[
                        "Plain Text OCR",
                        "Format Text OCR",
                        "Fine-grained OCR (Box)",
                        "Fine-grained OCR (Color)",
                        "Multi-crop OCR",
                        "Render Formatted OCR"
                    ],
                    label="Select Task",
                    value="Plain Text OCR"
                )
                ocr_type_dropdown = gr.Dropdown(
                    choices=["ocr", "format"],
                    label="OCR Type",
                    visible=False
                )
                ocr_box_input = gr.Textbox(
                    label="OCR Box (x1,y1,x2,y2)",
                    placeholder="[100,100,200,200]",
                    visible=False
                )
                ocr_color_dropdown = gr.Dropdown(
                    choices=["red", "green", "blue"],
                    label="OCR Color",
                    visible=False
                )
                # with gr.Row():
                    # max_new_tokens_slider = gr.Slider(50, 500, step=10, value=150, label="Max New Tokens")
                    # no_repeat_ngram_size_slider = gr.Slider(1, 10, step=1, value=2, label="No Repeat N-gram Size")

                submit_button = gr.Button("Process")
                editor_submit_button = gr.Button("Process Edited Image", visible=False)

        with gr.Column(scale=1):
            with gr.Group():
                output_markdown = gr.Textbox(label="🫴🏻📸GOT-OCR")
                output_html = gr.HTML(label="🫴🏻📸GOT-OCR")

    task_dropdown.change(
        update_inputs,
        inputs=[task_dropdown],
        outputs=[ocr_type_dropdown, ocr_box_input, ocr_color_dropdown, image_input, image_editor, submit_button, editor_submit_button]
    )
    
    task_dropdown.change(
        update_image_input,
        inputs=[task_dropdown],
        outputs=[image_input, image_editor, editor_submit_button]
    )
    
    submit_button.click(
        ocr_demo,
        inputs=[image_input, task_dropdown, ocr_type_dropdown, ocr_box_input, ocr_color_dropdown],
        outputs=[output_markdown, output_html]
    )
    editor_submit_button.click(
        ocr_demo,
        inputs=[image_editor, task_dropdown, ocr_type_dropdown, ocr_box_input, ocr_color_dropdown],
        outputs=[output_markdown, output_html]
    )

if __name__ == "__main__":
    cleanup_old_files()
    demo.launch(ssr_mode = False, mcp_server=True)