from flask import Flask, render_template, request
from model.got_model import perform_ocr  # Import the OCR function
import os
import re

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads/'

# Create the uploads folder if it doesn't exist
if not os.path.exists(app.config['UPLOAD_FOLDER']):
    os.makedirs(app.config['UPLOAD_FOLDER'])

# Global variable to store extracted text
extracted_text_global = ""

@app.route("/", methods=["GET", "POST"])
def upload_image():
    global extracted_text_global
    if request.method == "POST":
        file = request.files.get("file")  # Get the uploaded file
        
        if file:
            file_path = os.path.join(app.config["UPLOAD_FOLDER"], file.filename)
            file.save(file_path)  # Save the file to the uploads folder

            # Run the OCR model on the uploaded image
            extracted_text_global = perform_ocr(file_path)

            # Render the result template, passing the extracted text
            return render_template("result.html", extracted_text=extracted_text_global)
        else:
            # If no file was uploaded, display an error message
            return render_template("upload.html", error="Please upload a valid image file.")
    
    return render_template("upload.html")

@app.route("/search", methods=["POST"])
def search_text():
    global extracted_text_global
    keyword = request.form.get("keyword")
    
    if keyword:
        # Escape special characters in the keyword to avoid conflicts in HTML
        keyword = keyword.strip()

        # Check if the keyword is found in the extracted text
        if re.search(keyword, extracted_text_global, re.IGNORECASE):
            # Use re.sub to replace the keyword with <mark> for case-insensitive highlighting
            highlighted_text = re.sub(f"({re.escape(keyword)})", r"<mark>\1</mark>", extracted_text_global, flags=re.IGNORECASE)
            result_message = f"The keyword '{keyword}' was found and highlighted in the text."
        else:
            highlighted_text = extracted_text_global
            result_message = f"The keyword '{keyword}' was not found in the extracted text."
        
        # Render the result page with the highlighted text and search result message
        return render_template("result.html", extracted_text=highlighted_text, search_result=result_message)
    
    return render_template("result.html", extracted_text=extracted_text_global, search_result="Please enter a keyword.")
if __name__ == "__main__":
    app.run(debug=True)