import os
import shutil
import tempfile
import cv2
import numpy as np
import gradio as gr
import zipfile
import patoolib
import threading
from paddleocr import PaddleOCR
from PIL import Image

def delayed_cleanup(path, delay=30):
    def cleanup():
        import time
        time.sleep(delay)
        try:
            if os.path.isdir(path):
                shutil.rmtree(path)
            elif os.path.exists(path):
                os.remove(path)
        except Exception as e:
            print(f"Cleanup failed: {e}")
    threading.Thread(target=cleanup).start()

ocr = PaddleOCR(use_angle_cls=True, lang='en', det_model_dir='models/det', rec_model_dir='models/rec', cls_model_dir='models/cls')

def classify_background_color(avg_color, white_thresh=230, black_thresh=50, yellow_thresh=100):
    r, g, b = avg_color
    if r >= white_thresh and g >= white_thresh and b >= white_thresh:
        return (255, 255, 255)
    if r <= black_thresh and g <= black_thresh and b <= black_thresh:
        return (0, 0, 0)
    if r >= yellow_thresh and g >= yellow_thresh and b < yellow_thresh:
        return (255, 255, 0)
    return None

def sample_border_color(image, box, padding=2):
    h, w = image.shape[:2]
    x_min, y_min, x_max, y_max = box
    x_min = max(0, x_min - padding)
    x_max = min(w-1, x_max + padding)
    y_min = max(0, y_min - padding)
    y_max = min(h-1, y_max + padding)

    top = image[y_min:y_min+padding, x_min:x_max]
    bottom = image[y_max-padding:y_max, x_min:x_max]
    left = image[y_min:y_max, x_min:x_min+padding]
    right = image[y_min:y_max, x_max-padding:x_max]

    border_pixels = np.vstack((top.reshape(-1, 3), bottom.reshape(-1, 3),
                               left.reshape(-1, 3), right.reshape(-1, 3)))
    if border_pixels.size == 0:
        return (255, 255, 255)
    median_color = np.median(border_pixels, axis=0)
    return tuple(map(int, median_color))

def detect_text_boxes(image):
    results = ocr.ocr(image, cls=True)
    if not results or not results[0]:
        return []
    boxes = []
    for line in results[0]:
        box, (text, confidence) = line
        if text.strip():
            x_min = int(min(pt[0] for pt in box))
            x_max = int(max(pt[0] for pt in box))
            y_min = int(min(pt[1] for pt in box))
            y_max = int(max(pt[1] for pt in box))
            boxes.append(((x_min, y_min, x_max, y_max), text, confidence))
    return boxes

def remove_text_dynamic_fill(img_path, output_path):
    image = cv2.imread(img_path)
    if image is None:
        return
    if len(image.shape) == 2 or image.shape[2] == 1:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    else:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    boxes = detect_text_boxes(image)

    for (bbox, text, confidence) in boxes:
        if confidence < 0.4 or not text.strip():
            continue
        x_min, y_min, x_max, y_max = bbox
        height = y_max - y_min
        padding = 2 if height <= 30 else 4 if height <= 60 else 6
        x_min_p = max(0, x_min - padding)
        y_min_p = max(0, y_min - padding)
        x_max_p = min(image.shape[1]-1, x_max + padding)
        y_max_p = min(image.shape[0]-1, y_max + padding)
        sample_crop = image[y_min_p:y_max_p, x_min_p:x_max_p]
        avg_color = np.mean(sample_crop.reshape(-1, 3), axis=0)
        fill_color = classify_background_color(avg_color)
        if fill_color is None:
            fill_color = sample_border_color(image, (x_min, y_min, x_max, y_max))
        cv2.rectangle(image, (x_min_p, y_min_p), (x_max_p, y_max_p), fill_color, -1)

    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(output_path, image)

def extract_comic_archive(archive_path, extract_to):
    if archive_path.endswith(".cbz"):
        with zipfile.ZipFile(archive_path, 'r') as zip_ref:
            zip_ref.extractall(extract_to)
    elif archive_path.endswith(".cbr"):
        patoolib.extract_archive(archive_path, outdir=extract_to)

def process_cbz_cbr(file):
    temp_input = tempfile.mkdtemp()
    temp_output = tempfile.mkdtemp()
    extract_comic_archive(file.name, temp_input)
    for root, _, files in os.walk(temp_input):
        for fname in files:
            if fname.lower().endswith((".jpg", ".jpeg", ".png")):
                src = os.path.join(root, fname)
                dst = os.path.join(temp_output, fname)
                remove_text_dynamic_fill(src, dst)
    zip_path = shutil.make_archive(temp_output, 'zip', temp_output)
    delayed_cleanup(temp_input)
    delayed_cleanup(temp_output)
    delayed_cleanup(zip_path)
    return zip_path

demo = gr.Interface(
    fn=process_cbz_cbr,
    inputs=gr.File(file_types=[".cbz", ".cbr"], label="Upload Comic Archive (.cbz or .cbr)"),
    outputs=gr.File(label="Download Cleaned Zip"),
    title="Comic Text Cleaner (.cbz/.cbr)",
    description="Upload a .cbz or .cbr file and get a zip of cleaned comic images (text removed using PaddleOCR)."
)

demo.launch()