import gradio as gr
import spaces
from transformers import AutoModel, AutoTokenizer
import os
import base64
import io
import uuid
import time
import shutil
from pathlib import Path
import re
import easyocr

tokenizer = AutoTokenizer.from_pretrained('RufusRubin777/GOT-OCR2_0_CPU', trust_remote_code=True, device_map='cpu')
model = AutoModel.from_pretrained('RufusRubin777/GOT-OCR2_0_CPU', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cpu', use_safetensors=True)
model = model.eval().cpu()
reader = easyocr.Reader(['hi'])

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 run_GOT(image, language):
    unique_id = str(uuid.uuid4())
    image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.png")
    shutil.copy(image, image_path)
    try:
        english_extraction = model.chat(tokenizer, image_path, ocr_type='ocr')
        hindi_extraction = reader.readtext(image)
        hindi_extract = ''
        for x in hindi_extraction:
            hindi_extract += x[1] + '\n'
        return english_extraction + '\n' + hindi_extract
    except Exception as e:
        return f"Error: {str(e)}", None
    finally:
        if os.path.exists(image_path):
            os.remove(image_path)

def search_keyword(text, keyword):
    if not keyword:
        return "<p>Please enter a keyword to search.</p>"
    
    # Convert text and keyword to lowercase for case-insensitive search
    text_lower = text.lower()
    keyword_lower = keyword.lower()
    
    # Find all occurrences of the keyword
    matches = list(re.finditer(re.escape(keyword_lower), text_lower))
    
    if not matches:
        return f"<p>Keyword '{keyword}' not found in the text.</p>"
    
    # Highlight all occurrences of the keyword
    result = []
    last_end = 0
    for match in matches:
        start, end = match.span()
        result.append(text[last_end:start])
        result.append(f'<mark>{text[start:end]}</mark>')
        last_end = end
    result.append(text[last_end:])
    
    highlighted_text = ''.join(result)
    
    # Wrap the result in a scrollable div with some basic styling
    return f"""
    <div style="max-height: 300px; overflow-y: auto; border: 1px solid #ccc; padding: 10px; background-color: #f9f9f9;">
        <p>{highlighted_text}</p>
    </div>
    <p>Found {len(matches)} occurrence(s) of '{keyword}'</p>
    """

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

title_html = """
<h1><span class="gradient-text" id="text">IIT Roorkee (GOT ASSIGNMENT)</span></h1>
"""

with gr.Blocks() as scan_master_web_app:
    gr.HTML(title_html)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Upload your image")
            submit_button = gr.Button("Extract Text")
            ocr_result = gr.Textbox(label="Extracted Text")
                    

        with gr.Column():
            keyword = gr.Textbox(label="Search a keyword in the extracted text")
            search_button = gr.Button("Search Keyword")
            search_result = gr.HTML(label="Search result")
    
    submit_button.click(
        run_GOT,
        inputs=[image_input],
        outputs=[ocr_result]
    )
    
    search_button.click(
        search_keyword,
        inputs=[ocr_result, keyword],
        outputs=[search_result]
    )

if __name__ == "__main__":
    cleanup_old_files()
    scan_master_web_app.launch()