import gradio as gr
import easyocr
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import math
from sklearn.cluster import KMeans
import urllib.request
import os

# Initialize EasyOCR reader (Set to CPU for free web hosting)
print("Initializing EasyOCR...")
reader = easyocr.Reader(['en'], gpu=False) 

# Download a default scalable font if it doesn't exist
FONT_PATH = "Roboto-Regular.ttf"
if not os.path.exists(FONT_PATH):
    print("Downloading default font...")
    urllib.request.urlretrieve(
        "https://github.com/googlefonts/roboto/raw/main/src/hinted/Roboto-Regular.ttf", 
        FONT_PATH
    )

def get_text_color(crop_img, mask):
    pixels = crop_img.reshape((-1, 3))
    mask_pixels = mask.reshape((-1))
    text_pixels = pixels[mask_pixels > 0]
    
    if len(text_pixels) == 0:
        return (0, 0, 0) 
        
    kmeans = KMeans(n_clusters=1, n_init=10)
    kmeans.fit(text_pixels)
    dominant_color = kmeans.cluster_centers_[0]
    return tuple(map(int, dominant_color))

def process_image(image, old_text, new_text):
    if image is None or not old_text or not new_text:
        return image, "Please provide an image and both text fields."

    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    img_h, img_w = img_cv.shape[:2]

    results = reader.readtext(img_cv)
    
    target_bbox = None
    for (bbox, text, prob) in results:
        if old_text.lower() in text.lower():
            target_bbox = bbox
            break
            
    if not target_bbox:
        return image, f"Text '{old_text}' not found in the image."

    (tl, tr, br, bl) = target_bbox
    tl, tr, br, bl = (int(tl[0]), int(tl[1])), (int(tr[0]), int(tr[1])), (int(br[0]), int(br[1])), (int(bl[0]), int(bl[1]))

    dx = tr[0] - tl[0]
    dy = tr[1] - tl[1]
    angle = math.degrees(math.atan2(dy, dx))
    
    width = int(math.hypot(dx, dy))
    height = int(math.hypot(tl[0] - bl[0], tl[1] - bl[1]))

    full_mask = np.zeros((img_h, img_w), dtype=np.uint8)
    pts = np.array([tl, tr, br, bl], dtype=np.int32)
    cv2.fillPoly(full_mask, [pts], 255)

    x_coords, y_coords = [p[0] for p in pts], [p[1] for p in pts]
    x_min, x_max = max(0, min(x_coords)), min(img_w, max(x_coords))
    y_min, y_max = max(0, min(y_coords)), min(img_h, max(y_coords))
    
    crop_img = img_cv[y_min:y_max, x_min:x_max]
    gray_crop = cv2.cvtColor(crop_img, cv2.COLOR_BGR2GRAY)
    _, text_mask_crop = cv2.threshold(gray_crop, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    text_color_bgr = get_text_color(crop_img, text_mask_crop)
    text_color_rgb = (text_color_bgr[2], text_color_bgr[1], text_color_bgr[0])

    kernel = np.ones((3,3), np.uint8)
    text_mask_crop = cv2.dilate(text_mask_crop, kernel, iterations=1)

    precise_full_mask = np.zeros((img_h, img_w), dtype=np.uint8)
    precise_full_mask[y_min:y_max, x_min:x_max] = text_mask_crop
    precise_full_mask = cv2.bitwise_and(precise_full_mask, full_mask)

    inpainted_img_cv = cv2.inpaint(img_cv, precise_full_mask, inpaintRadius=3, flags=cv2.INPAINT_NS)
    inpainted_img_rgb = cv2.cvtColor(inpainted_img_cv, cv2.COLOR_BGR2RGB)
    final_image = Image.fromarray(inpainted_img_rgb)

    txt_overlay = Image.new('RGBA', final_image.size, (255, 255, 255, 0))
    draw = ImageDraw.Draw(txt_overlay)

    font_size = 10
    font = ImageFont.truetype(FONT_PATH, font_size)
    while True:
        bbox = draw.textbbox((0, 0), new_text, font=font)
        if bbox[3] - bbox[1] >= height * 0.8: 
            break
        font_size += 1
        font = ImageFont.truetype(FONT_PATH, font_size)

    text_bbox = draw.textbbox((0, 0), new_text, font=font)
    text_w = text_bbox[2] - text_bbox[0]
    text_h = text_bbox[3] - text_bbox[1]
    
    text_img = Image.new('RGBA', (text_w, text_h), (255, 255, 255, 0))
    text_draw = ImageDraw.Draw(text_img)
    text_draw.text((0, 0), new_text, font=font, fill=text_color_rgb)

    rotated_text_img = text_img.rotate(-angle, expand=True, resample=Image.BICUBIC)
    
    center_x, center_y = sum(x_coords) // 4, sum(y_coords) // 4
    paste_x = center_x - (rotated_text_img.width // 2)
    paste_y = center_y - (rotated_text_img.height // 2)

    txt_overlay.paste(rotated_text_img, (paste_x, paste_y), rotated_text_img)
    final_image = Image.alpha_composite(final_image.convert('RGBA'), txt_overlay).convert('RGB')

    return final_image, "Success!"

with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🪄 Seamless Text-in-Image Replacement AI")
    gr.Markdown("Upload an image, specify the text you want to replace, and enter the new text. The AI will erase the old text, reconstruct the background, and render the new text matching the original color, size, and angle.")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Image")
            old_text_input = gr.Textbox(label="Old Text (to remove)", placeholder="e.g., Trisoy")
            new_text_input = gr.Textbox(label="New Text (to insert)", placeholder="e.g., Ridoy")
            submit_btn = gr.Button("Replace Text", variant="primary")
            
        with gr.Column():
            image_output = gr.Image(type="pil", label="Result Image")
            status_output = gr.Textbox(label="Status", interactive=False)

    submit_btn.click(
        fn=process_image,
        inputs=[image_input, old_text_input, new_text_input],
        outputs=[image_output, status_output]
    )

if __name__ == "__main__":
    app.launch()