import os
import shutil
import tempfile
import zipfile
import rarfile
import gradio as gr
import cv2
import numpy as np
from paddleocr import PaddleOCR
from PIL import Image
import rarfile
rarfile.UNRAR_TOOL = "unrar"
import psutil
import time

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 detect_text_boxes(image):
    results = ocr.ocr(image, cls=True)
    boxes = []
    if results and results[0]:
        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))
    else:
        print("No text detected in the image.")
    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:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif 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 process_cbz_cbr(files):
    final_output = tempfile.mkdtemp()
    wait_for_cpu()
    for file_path in files:
        if file_path.endswith(".cbz"):
            with zipfile.ZipFile(file_path, 'r') as archive:
                extract_dir = tempfile.mkdtemp()
                archive.extractall(extract_dir)
        elif file_path.endswith(".cbr"):
            with rarfile.RarFile(file_path,'r') as archive:
                extract_dir = tempfile.mkdtemp()
                archive.extractall(extract_dir)
        else:
            continue

        for root, _, imgs in os.walk(extract_dir):
            for img in imgs:
                if img.lower().endswith(('.jpg', '.jpeg', '.png')):
                    input_path = os.path.join(root, img)
                    output_path = os.path.join(final_output, os.path.basename(img))
                    remove_text_dynamic_fill(input_path, output_path)

    # Create output zip
    zip_buffer = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
    with zipfile.ZipFile(zip_buffer.name, 'w', zipfile.ZIP_DEFLATED) as zf:
        for root, _, files in os.walk(final_output):
            for file in files:
                zf.write(os.path.join(root, file), arcname=file)

    return zip_buffer.name

import os
import zipfile
import rarfile
import tempfile
import shutil

def convert_cbr_to_cbz(cbr_path):
    temp_dir = tempfile.mkdtemp()
    cbz_path = cbr_path.replace('.cbr', '.cbz')
    return cbz_path


def wait_for_cpu(threshold=90, interval=3, timeout=30):
    start = time.time()
    while psutil.cpu_percent(interval=1) > threshold:
        print("High CPU usage detected, waiting...")
        time.sleep(interval)
        if time.time() - start > timeout:
            print("Timed out waiting for CPU to cool down.")
            break


demo = gr.Interface(
    fn=process_cbz_cbr,
    inputs=gr.File(file_types=[".cbz"], file_count="multiple", label="Upload only .cbz Comic Files"),
    outputs=gr.File(label="Download Cleaned Zip"),
    concurrency_limit=1,
    title="Comic Cleaner from .CBZ",
    description="Upload .cbz comics. The app extracts, cleans (removes text), and gives back a zip of cleaned images."
)

demo.launch()