import spaces
import gradio as gr
import pyte
import math
import subprocess
import sys
import tempfile
import time
import os
import pty
import select
import fcntl
import struct
import termios
from pathlib import Path
from huggingface_hub import hf_hub_download
from PIL import Image
from ultralytics import YOLO
from preservation import PreservedWatermarkRemover
from ignored_preview import save_ignored_preview
from clean import load_clean, clean_image
from illustrious import load_illustrious, prepare_illustrious, infer_illustrious, save_illustrious
from watermark_remover import WatermarkRemover

INPUT_IMAGE_HEIGHT = 600
TEXT_WORKER = [sys.executable, "-m", "text_detection"]
TIME_GPU = 13
MODEL_DEVICE = "cuda"
PART_MODELS = {
    "eyes": ("killjoyelite/anime-eye-yolov8", "eye_yolov8n.pt", 0.25, 0),
    "hands": ("deepghs/anime_hand_detection", "hand_detect_v1.0_s/model.pt", 0.3, 0.12),
    "feet": ("Claquasse/foot_anime_yolo", "foot_anime_yolo11m_v3.pt", 0.25, 0.02),
}


def load_models():
    subprocess.run(TEXT_WORKER, check=True)
    removers = {}
    for model_id in (WatermarkRemover.DEFAULT_MODEL_ID, WatermarkRemover.CTRLREGEN_MODEL_ID):
        remover = PreservedWatermarkRemover((), model_id=model_id, device=MODEL_DEVICE)
        remover.preload()
        removers[model_id] = remover
    parts = {}
    for part, (repo, filename, _, _) in PART_MODELS.items():
        detector = YOLO(hf_hub_download(repo_id=repo, filename=filename))
        detector.fuse()
        parts[part] = detector.to(MODEL_DEVICE)
    return {
        "watermark": removers,
        "parts": parts,
        "clean": load_clean(MODEL_DEVICE),
        "illustrious": load_illustrious(MODEL_DEVICE),
    }


def set_terminal_size(fd, rows=24, cols=80):
    size = struct.pack('HHHH', rows, cols, 0, 0)
    fcntl.ioctl(fd, termios.TIOCSWINSZ, size)


def run_in_pty(cmd, env=None, timeout=600):
    screen = pyte.Screen(80, 30)
    stream = pyte.Stream(screen)

    master_fd, slave_fd = pty.openpty()
    set_terminal_size(master_fd, rows=30, cols=80)

    proc = subprocess.Popen(
        cmd, stdout=slave_fd, stderr=slave_fd, stdin=slave_fd,
        close_fds=True, env=env,
    )
    os.close(slave_fd)

    start_time = time.time()
    try:
        while True:
            if time.time() - start_time > timeout:
                proc.kill()
                yield get_screen_text(screen), True
                return

            ready, _, _ = select.select([master_fd], [], [], 0.3)
            if ready:
                try:
                    data = os.read(master_fd, 4096)
                    if not data:
                        break
                    stream.feed(data.decode("utf-8", errors="replace"))
                except OSError:
                    break

            if proc.poll() is not None:
                while True:
                    r, _, _ = select.select([master_fd], [], [], 0.1)
                    if r:
                        try:
                            d = os.read(master_fd, 4096)
                            if not d:
                                break
                            stream.feed(d.decode("utf-8", errors="replace"))
                        except OSError:
                            break
                    else:
                        break
                yield get_screen_text(screen), True
                return

            yield get_screen_text(screen), False
        
        # while loop exited via break (empty read or OSError)
        # yield final screen with is_done=True
        yield get_screen_text(screen), True
    finally:
        try:
            os.close(master_fd)
        except OSError:
            pass
        if proc.poll() is None:
            proc.kill()
            proc.wait()


def get_screen_text(screen):
    lines = []
    for row in range(screen.lines):
        line = ""
        for col in range(screen.columns):
            line += screen.buffer[row][col].data
        lines.append(line.rstrip())
    while lines and not lines[-1].strip():
        lines.pop()
    return '\n'.join(lines)


def escape_html(text):
    return text.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')


def make_terminal_html(content, cursor=True):
    cursor_html = "<span class='cursor'>█</span>" if cursor else ""
    return f"<div class='terminal'><pre>{content}{cursor_html}</pre></div>"


@spaces.GPU(duration=TIME_GPU)
def run_watermark(input_path, output_path, strength, steps, model_id, boxes, text_mask):
    selected = MODELS["watermark"][model_id or WatermarkRemover.DEFAULT_MODEL_ID]
    selected.boxes = boxes
    selected.text_mask = text_mask
    try:
        selected.remove_watermark(input_path, output_path, strength=strength, num_inference_steps=steps)
    finally:
        selected.boxes = ()
        selected.text_mask = None


@spaces.GPU(duration=TIME_GPU)
def run_illustrious(image, control, mask, strength, steps):
    return infer_illustrious(MODELS["illustrious"], image, control, mask, strength, steps)


@spaces.GPU(duration=TIME_GPU)
def run_clean(input_path, output_path):
    clean_image(MODELS["clean"], input_path, output_path)


@spaces.GPU(duration=TIME_GPU)
def detect_parts(input_path, part):
    _, _, confidence, padding = PART_MODELS[part]
    model = MODELS["parts"][part]
    result = model.predict(source=str(input_path), conf=confidence, device=MODEL_DEVICE, verbose=False)[0]
    width, height = result.orig_shape[1], result.orig_shape[0]
    boxes = []
    for x1, y1, x2, y2 in result.boxes.xyxy.cpu().tolist():
        dx, dy = (x2 - x1) * padding, (y2 - y1) * padding
        box = (
            max(0, math.floor(x1 - dx)), max(0, math.floor(y1 - dy)),
            min(width, math.ceil(x2 + dx)), min(height, math.ceil(y2 + dy)),
        )
        if box[2] > box[0] and box[3] > box[1]:
            boxes.append(box)
    return boxes


def detect_text(input_path):
    with tempfile.TemporaryDirectory() as directory:
        output_path = Path(directory) / "mask.png"
        subprocess.run([*TEXT_WORKER, str(input_path), str(output_path)], check=True)
        with Image.open(output_path) as mask:
            return mask.copy()


def restore_parts(input_path, output_path, boxes, text_mask=None):
    with Image.open(input_path) as original, Image.open(output_path) as cleaned:
        result = cleaned.copy()
        width, height = result.size
        sx, sy = width / original.width, height / original.height
        original = original.convert(result.mode)
        if original.size != result.size:
            original = original.resize(result.size, Image.Resampling.LANCZOS)
        for x1, y1, x2, y2 in boxes:
            box = (
                max(0, math.floor(x1 * sx)), max(0, math.floor(y1 * sy)),
                min(width, math.ceil(x2 * sx)), min(height, math.ceil(y2 * sy)),
            )
            if box[2] > box[0] and box[3] > box[1]:
                result.paste(original.crop(box), box)
        if text_mask is not None:
            result.paste(original, (0, 0), text_mask.resize(result.size, Image.Resampling.NEAREST))
        result.save(output_path)


def remove_watermark(file_obj, strength, steps, model, ignore_eyes=True, ignore_hands=True,
                     ignore_feet=True, clean_after=True, ignore_text=True):
    if file_obj is None:
        raise gr.Error("Please upload an image first.")

    steps = max(int(steps), math.ceil(1 / strength))
    input_path = Path(file_obj)
    tmp_dir = tempfile.mkdtemp()
    output_path = Path(tmp_dir) / "cleaned.png"
    preview_path = Path(tmp_dir) / "ignored.png"
    work_path = Path(tmp_dir) / "result.png"
    enabled = {"eyes": ignore_eyes, "hands": ignore_hands, "feet": ignore_feet}
    regions = {}
    text_mask = None

    if model == "CtrlRegen Illustrious (SDXL)":
        cmd_display = f"CtrlRegen-like Illustrious: {input_path.name} --strength {strength} --steps {int(steps)} -o cleaned.png"
    else:
        profile = " --model-profile ctrlregen" if model == "CtrlRegen" else ""
        cmd_display = f"$ noai-watermark {input_path.name}{profile} --strength {strength} --steps {int(steps)} -o cleaned.png"
    cmd_html = f"<span style='color:#8be9fd'>{escape_html(cmd_display)}</span>"

    try:
        for part, ignore in enabled.items():
            if ignore:
                regions[part] = detect_parts(input_path, part)
        if ignore_text:
            text_mask = detect_text(input_path)
            if text_mask.getbbox() is None:
                text_mask = None
        boxes = [box for found in regions.values() for box in found]
        if model == "CtrlRegen Illustrious (SDXL)":
            image, control, mask, crop, size = prepare_illustrious(input_path, boxes, text_mask)
            result = run_illustrious(image, control, mask, strength, int(steps))
            save_illustrious(result, crop, size, work_path)
        else:
            model_id = WatermarkRemover.CTRLREGEN_MODEL_ID if model == "CtrlRegen" else None
            run_watermark(input_path, work_path, strength, int(steps), model_id, boxes, text_mask)
        with Image.open(work_path) as image:
            image.save(output_path)
        if clean_after:
            run_clean(output_path, output_path)
        if boxes or text_mask is not None:
            restore_parts(input_path, output_path, boxes, text_mask)
        save_ignored_preview(input_path, preview_path, regions, text_mask)
    except Exception as error:
        return make_terminal_html(cmd_html + "\n\n" + escape_html(str(error)), cursor=False), None, None

    preserved = "".join(
        f" Preserved {len(boxes)} detected {part}." if boxes else f" No {part} detected; none preserved."
        for part, boxes in regions.items()
    )
    if ignore_text:
        preserved += " Preserved text character mask." if text_mask is not None else " No text characters segmented; none preserved."
    cleaned = " Applied Clean after (1x Archivist Soft)." if clean_after else ""
    return make_terminal_html(cmd_html + f"\n\nDone!{preserved}{cleaned}", cursor=False), str(output_path), str(preview_path)


def check_ai_metadata(file_obj):
    if file_obj is None:
        raise gr.Error("Please upload an image first.")

    input_path = Path(file_obj)

    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    env["TERM"] = "xterm-256color"
    env["FORCE_COLOR"] = "1"
    env["CLICOLOR_FORCE"] = "1"

    final_text = ""
    for screen_text, is_done in run_in_pty(
        ["noai-watermark", str(input_path), "--check-ai"],
        env=env, timeout=30
    ):
        final_text = screen_text

    display = escape_html(final_text)
    cmd_display = f"$ noai-watermark {input_path.name} --check-ai"
    return make_terminal_html(f"<span style='color:#8be9fd'>{escape_html(cmd_display)}</span>\n\n" + display, cursor=False)


# --- CSS ---
terminal_css = """
.terminal {
    background: #1a1b26;
    color: #a9b1d6;
    border-radius: 10px;
    padding: 20px 24px;
    font-family: 'JetBrains Mono', 'Fira Code', 'SF Mono', 'Cascadia Code', 'Consolas', monospace;
    font-size: 13px;
    line-height: 1.7;
    overflow-x: auto;
    border: 1px solid #33354a;
    max-height: 500px;
    overflow-y: auto;
    box-shadow: 0 4px 16px rgba(0,0,0,0.5);
}
.terminal pre {
    margin: 0;
    white-space: pre-wrap;
    word-wrap: break-word;
    color: #a9b1d6;
    font-family: inherit;
}
.terminal .cursor {
    color: #7aa2f7;
    animation: blink 1s step-end infinite;
}
@keyframes blink {
    50% { opacity: 0; }
}
"""

MODELS = load_models()

with gr.Blocks(title="noai-watermark") as demo:

    gr.HTML("""
        <h1 style="text-align:center;">🛡️ noai-watermark</h1>
        <p style="text-align:center; color:#888;">
            Remove invisible AI watermarks (SynthID, StableSignature, TreeRing)
            and strip AI metadata from images
        </p>
    """)

    with gr.Tabs():

        with gr.TabItem("🔧 Watermark Removal"):
            upload_wm = gr.Image(label="Upload Image (PNG or JPEG)", sources=["upload"],
                                 type="filepath", image_mode=None, format="png",
                                 height=INPUT_IMAGE_HEIGHT)
            model = gr.Radio(
                ["CtrlRegen Illustrious (SDXL)", "CtrlRegen", "Default"],
                value="CtrlRegen", label="Model",
            )
            with gr.Row():
                ignore_eyes = gr.Checkbox(value=True, label="Ignore eyes")
                ignore_hands = gr.Checkbox(value=True, label="Ignore hands")
                ignore_feet = gr.Checkbox(value=True, label="Ignore feet")
                ignore_text = gr.Checkbox(value=True, label="Ignore text")
                clean_after = gr.Checkbox(value=True, label="Clean after")

            with gr.Row():
                strength = gr.Slider(
                    minimum=0.01, maximum=1.0, value=0.02, step=0.01,
                    label="Strength",
                    info="Lower = less change. Increase if watermark persists.",
                )
                steps = gr.Slider(
                    minimum=10, maximum=100, value=50, step=5,
                    label="Denoising Steps",
                    info="More steps = better quality but slower.",
                )

            with gr.Row():
                btn_super_minimal = gr.Button("Super Minimal (0.02/50)", size="sm")
                btn_minimal = gr.Button("Minimal (0.04/50)", size="sm")
                btn_balanced = gr.Button("Balanced (0.15/50)", size="sm")
                btn_aggressive = gr.Button("Aggressive (0.35/60)", size="sm")
                btn_max = gr.Button("Maximum (0.70/60)", size="sm")

            btn_super_minimal.click(lambda: (0.02, 50), outputs=[strength, steps])
            btn_minimal.click(lambda: (0.04, 50), outputs=[strength, steps])
            btn_balanced.click(lambda: (0.15, 50), outputs=[strength, steps])
            btn_aggressive.click(lambda: (0.35, 60), outputs=[strength, steps])
            btn_max.click(lambda: (0.70, 60), outputs=[strength, steps])

            run_btn = gr.Button("🚀 Remove Watermark", variant="primary", size="lg")
            wm_terminal = gr.HTML(value=make_terminal_html("<span style='color:#565f89'>Waiting for input...</span>"))

            with gr.Row():
                output_image = gr.Image(
                    type="filepath", image_mode=None, format="png", interactive=False,
                    show_label=False, height=INPUT_IMAGE_HEIGHT, buttons=["download", "fullscreen"],
                )
                ignored_image = gr.Image(
                    type="filepath", image_mode=None, format="png", interactive=False,
                    show_label=False, height=INPUT_IMAGE_HEIGHT, buttons=["download", "fullscreen"],
                )

            run_btn.click(
                fn=remove_watermark,
                inputs=[upload_wm, strength, steps, model, ignore_eyes, ignore_hands, ignore_feet, clean_after, ignore_text],
                outputs=[wm_terminal, output_image, ignored_image],
                concurrency_limit=1,
            )

        with gr.TabItem("🔍 Check AI Metadata"):
            upload_meta = gr.Image(label="Upload Image to Check (PNG or JPEG)", sources=["upload"],
                                   type="filepath", image_mode=None,
                                   height=INPUT_IMAGE_HEIGHT)
            check_btn = gr.Button("🔍 Check Metadata", variant="secondary", size="lg")
            meta_terminal = gr.HTML(value=make_terminal_html("<span style='color:#565f89'>Waiting for input...</span>"))

            check_btn.click(
                fn=check_ai_metadata,
                inputs=[upload_meta],
                outputs=[meta_terminal],
            )

    gr.Markdown("""
    ---
    <div style="text-align:center; color:#888; font-size:0.85em;">
        <a href="https://github.com/mertizci/noai-watermark" target="_blank">GitHub</a> ·
        <a href="https://pypi.org/project/noai-watermark/" target="_blank">PyPI</a> ·
        Built by <a href="https://github.com/mertizci" target="_blank">Mert IZCI</a>
        <br>For security research and defensive analysis only.
    </div>
    """)

if __name__ == "__main__":
    demo.launch(ssr_mode=False, css=terminal_css, allowed_paths=["/tmp"])