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


model_name = 'ucaslcl/GOT-OCR2_0'

device = 'cuda' if torch.cuda.is_available() else 'cpu'

tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
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)
model = model.eval().to(device)
model.config.pad_token_id = tokenizer.eos_token_id

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()


@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

        if task == "Plain Text OCR":
            res = model.chat(tokenizer, image_path, ocr_type='ocr')
            return res, None, unique_id
        else:
            if task == "Format Text OCR":
                res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
            elif task == "Fine-grained OCR (Box)":
                res = model.chat(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 = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
            elif task == "Multi-crop OCR":
                res = model.chat_crop(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
            elif task == "Render Formatted OCR":
                res = model.chat(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 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):
    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)