import html
import queue
import threading
import uuid
from dataclasses import dataclass
from functools import wraps
from typing import Any, Callable

import gradio as gr

from shared import magic_mask


MAGIC_WAND_LABEL = "\U0001FA84"
_ABORT_EVENTS: dict[str, threading.Event] = {}
_ORIGINAL_IMAGE_EDITOR = None


class MagicMaskAbort(Exception):
    pass


def magic_mask_button_updates(image_mode, video_prompt_type):
    mask_visible = "V" in video_prompt_type and "A" in video_prompt_type and "U" not in video_prompt_type
    image_outputs = image_mode > 0
    return gr.update(visible=mask_visible and image_outputs), gr.update(visible=mask_visible and not image_outputs)


def _status_html(message, tone="info"):
    if not message:
        return ""
    tone_class = "is-error" if tone == "error" else ""
    return f"<div class='wangp-magic-mask-message {tone_class}'>{html.escape(str(message))}</div>"


def _progress_html(message, percent=0):
    if not message:
        return ""
    percent = max(0, min(100, int(percent)))
    return (
        "<div class='wangp-magic-mask-progress'>"
        f"<div class='wangp-magic-mask-progress-label'>{html.escape(str(message))}</div>"
        "<div class='wangp-magic-mask-progress-track'>"
        f"<div class='wangp-magic-mask-progress-bar' style='width:{percent}%;'></div>"
        "</div></div>"
    )


def _abort_event(token):
    token = str(token or "")
    if len(token) == 0:
        token = "default"
    if token not in _ABORT_EVENTS:
        _ABORT_EVENTS[token] = threading.Event()
    return _ABORT_EVENTS[token]


def _open_panel():
    return gr.update(visible=True), "", "", gr.update(visible=True, interactive=True), gr.update(visible=False, interactive=False), None, None


def _close_panel():
    return (
        gr.update(visible=False),
        "",
        "",
        gr.update(visible=False, interactive=False),
        None,
        None,
    )


def _abort_magic_mask(abort_token):
    _abort_event(abort_token).set()
    return _status_html("Aborting Magic Mask..."), gr.update(visible=True, interactive=False)


def _exit_button_running():
    return gr.update(visible=False, interactive=False)


def _exit_button_idle():
    return gr.update(visible=True, interactive=True)


def _abort_button_running():
    return gr.update(visible=True, interactive=True)


def _abort_button_idle():
    return gr.update(visible=False, interactive=False)


def _raise_if_aborted(abort_event):
    if abort_event.is_set():
        raise MagicMaskAbort()


def _image_source(image_mask_guide, image_guide):
    if isinstance(image_mask_guide, dict) and image_mask_guide.get("background") is not None:
        return image_mask_guide["background"]
    return image_guide


def _keywords_processed_html(processed, total):
    return _status_html(f"Masks generated: {processed}/{total}")


def _mask_progress_html(keyword, done, total):
    percent = int(done * 100 / max(int(total), 1))
    return _progress_html(f'Generating Mask "{keyword}"', percent)


def _current_keyword_progress(keywords, done, total):
    keyword_count = max(len(keywords), 1)
    total_steps = max(float(total), 1.0)
    done = max(0.0, min(float(done), total_steps))
    current_index = min(keyword_count - 1, int(done * keyword_count / total_steps))
    keyword_steps = total_steps / keyword_count
    keyword_done = max(0.0, min(keyword_steps, done - current_index * keyword_steps))
    return keywords[current_index], keyword_done, keyword_steps


def _run_keyword_mask(video, keywords, abort_event):
    progress_events = queue.Queue()

    def progress_callback(done, total):
        if abort_event.is_set():
            raise MagicMaskAbort()
        progress_events.put(("progress", int(done), int(total)))

    def worker():
        try:
            if abort_event.is_set():
                raise MagicMaskAbort()
            progress_events.put(("done", magic_mask.generate_keyword_masks(video, keywords, progress_callback=progress_callback)))
        except MagicMaskAbort as exc:
            progress_events.put(("abort", exc))
        except Exception as exc:
            progress_events.put(("error", exc))

    thread = threading.Thread(target=worker, daemon=True)
    thread.start()
    pending_event = None
    while True:
        event = pending_event or progress_events.get()
        pending_event = None
        if event[0] == "progress":
            latest_event = event
            while True:
                try:
                    next_event = progress_events.get_nowait()
                except queue.Empty:
                    break
                if next_event[0] == "progress":
                    latest_event = next_event
                else:
                    pending_event = next_event
                    break
            event = latest_event
            yield event[1], event[2]
        elif event[0] == "done":
            thread.join()
            return event[1]
        elif event[0] == "abort":
            thread.join()
            raise event[1]
        else:
            thread.join()
            raise event[1]


def _generate_magic_mask(
    state,
    keywords_text,
    negative_mask,
    image_mode,
    video_guide,
    image_mask_guide,
    image_guide,
    abort_token,
    *,
    download_assets: Callable[[dict[str, Any]], Any],
    acquire_gpu: Callable[[Any, str, str], Any],
    release_gpu: Callable[[Any, str], Any],
    get_model_settings: Callable[[Any], dict],
):
    source_image = None
    if image_mode > 0:
        source_image = _image_source(image_mask_guide, image_guide)
        if source_image is None:
            yield gr.update(), gr.update(), gr.update(), _status_html("Magic Mask needs a control image.", "error"), gr.update(visible=True), "", _exit_button_idle(), _abort_button_idle(), None, None
            return
    elif video_guide is None:
        yield gr.update(), gr.update(), gr.update(), _status_html("Magic Mask needs a control video.", "error"), gr.update(visible=True), "", _exit_button_idle(), _abort_button_idle(), None, None
        return

    keywords = magic_mask.parse_keywords(keywords_text)
    if len(keywords) == 0:
        yield gr.update(), gr.update(), gr.update(), _status_html("Enter at least one keyword.", "error"), gr.update(visible=True), "", _exit_button_idle(), _abort_button_idle(), None, None
        return
    keywords_label = ", ".join(keywords)
    abort_event = _abort_event(abort_token)
    abort_event.clear()
    acquired = False
    try:
        yield gr.update(), gr.update(), gr.update(), _status_html("Initializing Magic Mask"), gr.update(visible=True), _progress_html("Preparing files", 0), _exit_button_running(), _abort_button_running(), None, None
        download_assets(magic_mask.query_download_def())
        _raise_if_aborted(abort_event)
        yield gr.update(), gr.update(), gr.update(), _status_html("Initializing Magic Mask"), gr.update(visible=True), _progress_html("Initializing", 0), _exit_button_running(), _abort_button_running(), None, None
        acquire_gpu(state, magic_mask.PROCESS_ID, magic_mask.PROCESS_NAME)
        _raise_if_aborted(abort_event)
        acquired = True
        ui_settings = get_model_settings(state)
        if image_mode > 0:
            _raise_if_aborted(abort_event)
            background, video = magic_mask.prepare_image_mask_input(source_image)
            total = len(keywords)
            mask_generator = _run_keyword_mask(video, keywords, abort_event)
            progress_started = False
            try:
                while True:
                    done, frame_total = next(mask_generator)
                    if done <= 0:
                        continue
                    progress_started = True
                    processed = min(total, int(done * total / max(frame_total, 1)))
                    current_keyword, keyword_done, keyword_total = _current_keyword_progress(keywords, done, frame_total)
                    yield gr.update(), gr.update(), gr.update(), _keywords_processed_html(processed, total), gr.update(visible=True), _mask_progress_html(current_keyword, keyword_done, keyword_total), _exit_button_running(), _abort_button_running(), None, None
            except StopIteration as stop:
                merged_mask = stop.value[0]
            if not progress_started:
                yield gr.update(), gr.update(), gr.update(), _keywords_processed_html(0, total), gr.update(visible=True), _mask_progress_html(keywords[0], 1, 1), _exit_button_running(), _abort_button_running(), None, None
            yield gr.update(), gr.update(), gr.update(), _keywords_processed_html(total, total), gr.update(visible=True), _mask_progress_html(keywords[-1], 1, 1), _exit_button_running(), _abort_button_running(), None, None
            yield gr.update(), gr.update(), gr.update(), _status_html("Saving Image Mask..."), gr.update(visible=True), "", _exit_button_running(), _abort_button_running(), None, None
            _raise_if_aborted(abort_event)
            mask_image = magic_mask.mask_to_image(magic_mask.finalize_masks(merged_mask, negative_mask=negative_mask))
            image_mask_guide_value = magic_mask.build_image_editor_value(background, mask_image)
            if isinstance(ui_settings, dict):
                ui_settings["image_guide"] = background
                ui_settings["image_mask"] = mask_image
            gr.Info(f"Magic Mask generated {'a negative ' if negative_mask else 'an '}image mask for: {keywords_label}.")
            yield gr.update(value=image_mask_guide_value), gr.update(value=mask_image), gr.update(), "", gr.update(visible=False), "", _exit_button_idle(), _abort_button_idle(), None, None
            return
        _raise_if_aborted(abort_event)
        video_path, video, fps = magic_mask.prepare_video_mask_input(video_guide)
        total = len(keywords)
        mask_generator = _run_keyword_mask(video, keywords, abort_event)
        progress_started = False
        try:
            while True:
                done, frame_total = next(mask_generator)
                if done <= 0:
                    continue
                progress_started = True
                processed = min(total, int(done * total / max(frame_total, 1)))
                current_keyword, keyword_done, keyword_total = _current_keyword_progress(keywords, done, frame_total)
                yield gr.update(), gr.update(), gr.update(), _keywords_processed_html(processed, total), gr.update(visible=True), _mask_progress_html(current_keyword, keyword_done, keyword_total), _exit_button_running(), _abort_button_running(), None, None
        except StopIteration as stop:
            merged_mask = stop.value
        if not progress_started:
            yield gr.update(), gr.update(), gr.update(), _keywords_processed_html(0, total), gr.update(visible=True), _mask_progress_html(keywords[0], 1, 1), _exit_button_running(), _abort_button_running(), None, None
        yield gr.update(), gr.update(), gr.update(), _keywords_processed_html(total, total), gr.update(visible=True), _mask_progress_html(keywords[-1], 1, 1), _exit_button_running(), _abort_button_running(), None, None
        yield gr.update(), gr.update(), gr.update(), _status_html("Saving Video Mask..."), gr.update(visible=True), "", _exit_button_running(), _abort_button_running(), None, None
        mask_path = magic_mask.save_mask_video(video_path, magic_mask.finalize_masks(merged_mask, negative_mask=negative_mask), fps, keywords, abort_callback=lambda: _raise_if_aborted(abort_event))
        if isinstance(ui_settings, dict):
            ui_settings["video_mask"] = mask_path
        gr.Info(f"Magic Mask generated {'a negative ' if negative_mask else 'a '}video mask for: {keywords_label}.")
        yield gr.update(), gr.update(), gr.update(value=mask_path), "", gr.update(visible=False), "", _exit_button_idle(), _abort_button_idle(), None, None
    except MagicMaskAbort:
        yield gr.update(), gr.update(), gr.update(), _status_html("Magic Mask aborted."), gr.update(visible=True), "", _exit_button_idle(), _abort_button_idle(), None, None
    except Exception as exc:
        yield gr.update(), gr.update(), gr.update(), _status_html(exc, "error"), gr.update(visible=True), "", _exit_button_idle(), _abort_button_idle(), None, None
    finally:
        if acquired:
            release_gpu(state, magic_mask.PROCESS_ID)


@dataclass
class MagicMaskUI:
    trigger: gr.Button | None = None
    panel: gr.Group | None = None
    keywords: gr.Textbox | None = None
    negative_mask: gr.Checkbox | None = None
    status: gr.HTML | None = None
    progress_html: gr.HTML | None = None
    cancel_btn: gr.Button | None = None
    abort_btn: gr.Button | None = None
    generate_btn: gr.Button | None = None
    abort_token: gr.State | None = None
    pending_image_mask_guide: gr.State | None = None
    pending_image_mask: gr.State | None = None

    @staticmethod
    def hidden_trigger():
        return gr.Button(MAGIC_WAND_LABEL, size="sm", min_width=1, visible=False, elem_classes=["wangp-magic-mask-trigger", "wangp-magic-mask-trigger--hidden"])

    @staticmethod
    def button_updates(image_mode, video_prompt_type):
        return magic_mask_button_updates(image_mode, video_prompt_type)

    @staticmethod
    def patch_image_editor():
        global _ORIGINAL_IMAGE_EDITOR
        if _ORIGINAL_IMAGE_EDITOR is not None:
            return True
        original = gr.ImageEditor
        original_init = original.__init__
        if getattr(original_init, "__wangp_magic_mask_patch__", False):
            _ORIGINAL_IMAGE_EDITOR = original_init
            return True

        @wraps(original_init)
        def patched_init(self, *args, **kwargs):
            original_init(self, *args, **kwargs)
            self._wangp_magic_mask_patch_enabled = True

        patched_init.__wangp_magic_mask_patch__ = True
        _ORIGINAL_IMAGE_EDITOR = original_init
        original.__init__ = patched_init
        return True

    @staticmethod
    def get_css():
        return r"""
.wangp-magic-mask-anchor {
    position: relative;
    gap: 0 !important;
    padding: 0 !important;
    margin: 0 !important;
}

.wangp-magic-mask-anchor--image-editor {
    position: relative;
}

.wangp-magic-mask-anchor > .form,
.wangp-magic-mask-anchor > .styler {
    gap: 0 !important;
    padding: 0 !important;
    margin: 0 !important;
}

.wangp-magic-mask-anchor:not(:has(> .block:not(.hide):not(.hidden), > button:not(.hide):not(.hidden), > .gr-group:not(.hide):not(.hidden))) {
    display: none !important;
}

.wangp-magic-mask-trigger,
.wangp-magic-mask-trigger button {
    width: 34px !important;
    min-width: 34px !important;
    max-width: 34px !important;
    height: 34px;
    min-height: 34px;
    padding: 0 !important;
    border: 1px solid rgba(17, 84, 118, 0.14);
    border-radius: 12px;
    background: linear-gradient(180deg, rgba(255, 255, 255, 0.99) 0%, rgba(236, 244, 249, 0.99) 100%);
    color: #155574;
    box-shadow: 0 10px 18px rgba(11, 44, 63, 0.08);
    font-weight: 700;
    line-height: 1;
}

.wangp-magic-mask-trigger--overlay {
    position: absolute !important;
    top: 28px;
    right: 8px;
    z-index: 35;
}

.wangp-magic-mask-trigger--editor {
    display: none !important;
}

.wangp-magic-mask-toolbar-button {
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    flex: 0 0 auto !important;
    margin: var(--spacing-xxs, 2px) !important;
    width: 28px !important;
    min-width: 28px !important;
    max-width: 28px !important;
    height: 28px !important;
    min-height: 28px !important;
    padding: 0 !important;
    border: 1px solid transparent !important;
    border-radius: var(--radius-xs, 4px) !important;
    background: transparent !important;
    color: var(--block-label-text-color) !important;
    box-shadow: none !important;
    font-size: 16px !important;
    line-height: 1 !important;
}

.wangp-magic-mask-toolbar-button:hover {
    cursor: pointer !important;
    background: var(--background-fill-secondary) !important;
    color: var(--color-accent) !important;
    transform: none !important;
}

.wangp-magic-mask-toolbar-button.wangp-magic-mask-unavailable {
    cursor: not-allowed !important;
    opacity: 0.52;
    filter: grayscale(1);
}

.wangp-magic-mask-toolbar-button.wangp-magic-mask-unavailable:hover {
    cursor: not-allowed !important;
    color: var(--body-text-color-subdued) !important;
}

.wangp-magic-mask-toolbar-button[hidden] {
    display: none !important;
}

.wangp-magic-mask-synthetic-toolbar {
    position: absolute !important;
    top: var(--block-label-margin, 8px) !important;
    right: var(--block-label-margin, 8px) !important;
    z-index: var(--layer-3, 1000) !important;
}

.wangp-magic-mask-trigger:hover,
.wangp-magic-mask-trigger button:hover {
    transform: translateY(-1px);
    box-shadow: 0 14px 24px rgba(11, 44, 63, 0.12);
}

.wangp-magic-mask-trigger:hover::after {
    content: "Magic Mask";
    position: absolute;
    top: 40px;
    right: 0;
    width: max-content;
    max-width: 160px;
    padding: 5px 7px;
    border-radius: 4px;
    background: rgba(0, 0, 0, 0.82);
    color: #ffffff;
    font-size: calc(12px * var(--wangp-ui-scale));
    font-weight: 600;
    line-height: 1.2;
    pointer-events: none;
}

.wangp-magic-mask-panel.hide {
    display: none !important;
}

.wangp-magic-mask-panel:not(.hide) {
    position: absolute !important;
    inset: 0;
    z-index: 40;
    display: flex !important;
    align-items: stretch;
    justify-content: stretch;
    margin: 0 !important;
    padding: 0 !important;
    border: 0 !important;
    background: #ffffff !important;
    box-shadow: none !important;
    overflow: hidden !important;
    pointer-events: auto;
    box-sizing: border-box;
}

.wangp-magic-mask-anchor--image-editor .wangp-magic-mask-panel:not(.hide) {
    padding-left: 20px !important;
    padding-right: 20px !important;
}

.wangp-magic-mask-panel:not(.hide),
.wangp-magic-mask-panel:not(.hide) *,
.wangp-magic-mask-card,
.wangp-magic-mask-card * {
    box-sizing: border-box;
    scrollbar-width: none;
}

.wangp-magic-mask-panel:not(.hide)::-webkit-scrollbar,
.wangp-magic-mask-panel:not(.hide) *::-webkit-scrollbar {
    width: 0 !important;
    height: 0 !important;
    display: none !important;
}

.wangp-magic-mask-panel > .form,
.wangp-magic-mask-panel > .styler {
    width: 100% !important;
    height: 100% !important;
    display: flex !important;
    align-items: stretch !important;
    justify-content: stretch !important;
    padding: 0 !important;
    border: 0 !important;
    background: transparent !important;
    box-shadow: none !important;
    overflow: hidden !important;
}

.wangp-magic-mask-card {
    display: flex !important;
    flex-direction: column !important;
    width: 100% !important;
    min-width: 0 !important;
    height: 100% !important;
    min-height: 0 !important;
    margin: 0 !important;
    padding: 0 !important;
    gap: 0 !important;
    border: 0 !important;
    border-radius: 0 !important;
    background: #ffffff !important;
    box-shadow: none !important;
    overflow: hidden !important;
}

.wangp-magic-mask-card > .form {
    display: flex !important;
    flex: 1 1 auto !important;
    flex-direction: column !important;
    height: 100% !important;
    min-height: 0 !important;
    padding: 0 !important;
    border: 0 !important;
    background: transparent !important;
    box-shadow: none !important;
    overflow: hidden !important;
}

.wangp-magic-mask-card .block,
.wangp-magic-mask-card .html-container,
.wangp-magic-mask-card .prose {
    max-width: 100% !important;
    overflow: hidden !important;
}

.wangp-magic-mask-titlebar {
    padding: 10px 16px 9px;
    background: linear-gradient(180deg, rgba(16, 86, 121, 0.98) 0%, rgba(10, 59, 84, 0.98) 100%);
    color: #f3fbff;
}

.wangp-magic-mask-heading {
    font-size: calc(0.95rem * var(--wangp-ui-scale));
    font-weight: 800;
    letter-spacing: 0;
    color: #f3fbff !important;
}

.wangp-magic-mask-body {
    flex: 0 0 auto !important;
    padding: 14px 18px 0;
    overflow: hidden !important;
}

.wangp-magic-mask-body .block,
.wangp-magic-mask-body .form,
.wangp-magic-mask-body .wrap {
    margin: 0 !important;
    overflow: hidden !important;
}

.wangp-magic-mask-intro {
    margin: 0 0 12px;
    color: #164f70;
    font-size: calc(0.88rem * var(--wangp-ui-scale));
    line-height: 1.45;
}

.wangp-magic-mask-keyword-row {
    align-items: center;
}

.wangp-magic-mask-keyword-row > .form {
    align-items: center;
}

.wangp-magic-mask-keywords textarea {
    min-height: 38px !important;
    height: 38px !important;
    overflow-y: hidden !important;
    resize: none !important;
}

.wangp-magic-mask-negative {
    flex: 0 0 150px !important;
    min-width: 150px !important;
    padding: 0 !important;
    border: 0 !important;
}

.wangp-magic-mask-negative,
.wangp-magic-mask-negative > .form,
.wangp-magic-mask-negative > .styler,
.wangp-magic-mask-negative .block,
.wangp-magic-mask-negative .wrap,
.wangp-magic-mask-negative .checkbox-wrap,
.wangp-magic-mask-negative-checkbox,
.wangp-magic-mask-negative-checkbox label,
.wangp-magic-mask-negative label {
    background: transparent !important;
    background-color: transparent !important;
    box-shadow: none !important;
    border-color: transparent !important;
    padding: 0 !important;
}

.wangp-magic-mask-negative label {
    white-space: nowrap !important;
}

.wangp-magic-mask-message {
    flex: 0 0 auto !important;
    margin: 12px 18px 0;
    color: #164f70;
    font-size: calc(0.9rem * var(--wangp-ui-scale));
    line-height: 1.5;
    font-weight: 600;
}

.wangp-magic-mask-message.is-error {
    color: #b33434;
}

.wangp-magic-mask-progress {
    flex: 0 0 auto !important;
    margin: 12px 18px 0;
}

.wangp-magic-mask-progress-label {
    margin-bottom: 6px;
    color: #164f70;
    font-size: calc(0.82rem * var(--wangp-ui-scale));
    font-weight: 700;
}

.wangp-magic-mask-progress-track {
    width: 100%;
    height: 8px;
    overflow: hidden;
    border-radius: 999px;
    background: rgba(19, 91, 126, 0.14);
}

.wangp-magic-mask-progress-bar {
    height: 100%;
    border-radius: 999px;
    background: linear-gradient(90deg, #2d89b7 0%, #56b18e 100%);
    transition: width 0.22s ease;
}

.wangp-magic-mask-spacer {
    flex: 1 1 auto !important;
    min-height: 0 !important;
    height: auto !important;
    margin: 0 !important;
    padding: 0 !important;
    border: 0 !important;
    overflow: hidden !important;
}

.wangp-magic-mask-actions {
    flex: 0 0 auto !important;
    justify-content: flex-end;
    align-items: flex-end;
    gap: 10px;
    margin-top: auto !important;
    padding: 18px;
    overflow: hidden !important;
}

.wangp-magic-mask-actions > .form {
    justify-content: flex-end !important;
    align-items: flex-end !important;
    overflow: hidden !important;
}

.wangp-magic-mask-btn,
.wangp-magic-mask-btn button {
    min-width: 92px;
    height: 40px;
    min-height: 40px;
    border-radius: 14px;
    border: 1px solid rgba(17, 84, 118, 0.14);
    background: linear-gradient(180deg, rgba(255, 255, 255, 0.99) 0%, rgba(237, 245, 250, 0.99) 100%);
    color: #155574;
    box-shadow: 0 10px 18px rgba(11, 44, 63, 0.08);
    font-weight: 700;
}

.wangp-magic-mask-btn--primary,
.wangp-magic-mask-btn--primary button {
    color: #f4fbff;
    border-color: rgba(10, 59, 84, 0.12);
    background: linear-gradient(180deg, rgba(16, 86, 121, 0.98) 0%, rgba(10, 59, 84, 0.98) 100%);
}

.wangp-magic-mask-btn--danger,
.wangp-magic-mask-btn--danger button {
    color: #ffffff;
    border-color: rgba(142, 45, 45, 0.16);
    background: linear-gradient(180deg, rgba(188, 67, 67, 0.98) 0%, rgba(132, 41, 41, 0.98) 100%);
}

.wangp-magic-mask-btn:disabled,
.wangp-magic-mask-btn button:disabled {
    cursor: not-allowed !important;
    filter: grayscale(0.9);
    opacity: 0.48;
}

@media (prefers-color-scheme: dark) {
    .wangp-magic-mask-trigger,
    .wangp-magic-mask-trigger button,
    .wangp-magic-mask-btn,
    .wangp-magic-mask-btn button {
        color: #ecf4f9;
        border-color: rgba(103, 132, 151, 0.22);
        background: linear-gradient(180deg, rgba(10, 10, 10, 0.99) 0%, rgba(21, 21, 21, 0.99) 100%);
        box-shadow: 0 10px 18px rgba(0, 0, 0, 0.22);
    }

    .wangp-magic-mask-panel:not(.hide),
    .wangp-magic-mask-card {
        background: #000000 !important;
    }

    .wangp-magic-mask-intro,
    .wangp-magic-mask-message,
    .wangp-magic-mask-progress-label {
        color: #ecf4f9;
    }

    .wangp-magic-mask-message.is-error {
        color: #ff9e9e;
    }

    .wangp-magic-mask-progress-track {
        background: rgba(236, 244, 249, 0.18);
    }
}
"""

    @staticmethod
    def get_javascript():
        return r"""
window.__wangpMagicMaskNS = window.__wangpMagicMaskNS || {};
const WMM = window.__wangpMagicMaskNS;
WMM.init = WMM.init || false;
WMM.observer = WMM.observer || null;
WMM.raf = WMM.raf || null;
WMM.interval = WMM.interval || null;

WMM.isVisible = function (element) {
    if (!element) return false;
    const rect = element.getBoundingClientRect();
    const style = getComputedStyle(element);
    return rect.width > 2 && rect.height > 2 && style.display !== 'none' && style.visibility !== 'hidden';
};

WMM.findImageEditorToolbar = function (editor) {
    const synthetic = editor.querySelector('.wangp-magic-mask-synthetic-toolbar');
    const topControls = editor.querySelector('.icon-button-wrapper.top-panel:not(.wangp-magic-mask-synthetic-toolbar), .icon-button-wrapper:not(.wangp-magic-mask-synthetic-toolbar)');
    if (topControls) {
        if (synthetic) synthetic.remove();
        return topControls;
    }

    const imageToolbars = Array.from(editor.querySelectorAll('.toolbar-wrap'));
    const primaryToolbar = imageToolbars.find((toolbar) => !toolbar.closest('.toolbar-wrap-wrap')) || imageToolbars[0];
    if (primaryToolbar) {
        if (synthetic) synthetic.remove();
        return primaryToolbar;
    }

    const explicit = editor.querySelector('[role="toolbar"], [class*="toolbar"], [class*="Toolbar"], [class*="tools"], [class*="Tools"]');
    if (explicit && !explicit.classList.contains('wangp-magic-mask-synthetic-toolbar')) {
        if (synthetic) synthetic.remove();
        return explicit;
    }

    const editorRect = editor.getBoundingClientRect();
    let best = null;
    let bestScore = -Infinity;
    const candidates = new Map();
    editor.querySelectorAll('button').forEach((button) => {
        let node = button.parentElement;
        for (let depth = 0; node && node !== editor && depth < 5; depth += 1, node = node.parentElement) {
            candidates.set(node, node.querySelectorAll('button').length);
        }
    });
    candidates.forEach((buttonCount, node) => {
        const rect = node.getBoundingClientRect();
        if (rect.width <= 0 || rect.height <= 0 || rect.top > editorRect.top + 100) return;
        const topDistance = Math.abs(rect.top - editorRect.top);
        const rightBias = rect.left > editorRect.left + editorRect.width * 0.45 ? 16 : 0;
        const heightPenalty = Math.max(0, rect.height - 52) * 2;
        const score = buttonCount * 24 + rightBias - topDistance - heightPenalty;
        if (score > bestScore) {
            best = node;
            bestScore = score;
        }
    });
    if (best) {
        if (synthetic) synthetic.remove();
        return best;
    }

    let fallback = synthetic;
    if (!fallback) {
        fallback = document.createElement('div');
        fallback.className = 'icon-button-wrapper top-panel wangp-magic-mask-synthetic-toolbar';
        fallback.setAttribute('role', 'toolbar');
        fallback.setAttribute('aria-label', 'ImageEditor tools');
        if (getComputedStyle(editor).position === 'static') editor.style.position = 'relative';
        editor.appendChild(fallback);
    }
    return fallback;
};

WMM.visibleImageEditors = function () {
    return Array.from(document.querySelectorAll('#img_editor, [data-testid="image-editor"], .imageeditor, .image-editor'))
        .map((candidate) => candidate.closest('.block') || candidate.closest('[id^="component-"]') || candidate)
        .filter((editor, index, editors) => WMM.isVisible(editor) && editors.indexOf(editor) === index);
};

WMM.focusImageEditor = function (editor, clickBrush) {
    if (!editor || !WMM.isVisible(editor)) return false;
    const focusTarget = editor.querySelector('.pixi-target canvas, canvas, .pixi-target, [data-testid="image"]') || editor;
    if (!focusTarget.hasAttribute('tabindex')) focusTarget.setAttribute('tabindex', '-1');
    try {
        focusTarget.focus({ preventScroll: true });
    } catch (_) {
        focusTarget.focus();
    }
    if (clickBrush) {
        const brushButton = editor.querySelector('button[aria-label="Brush"]') || document.querySelector('button[aria-label="Brush"]');
        if (brushButton && !brushButton.disabled) brushButton.click();
    }
    return true;
};

WMM.focusVisibleImageEditor = function (clickBrush) {
    const editor = WMM.visibleImageEditors()[0];
    return WMM.focusImageEditor(editor, clickBrush);
};

WMM.refocusImageEditorAfterMagicMask = function () {
    const openPanel = Array.from(document.querySelectorAll('.wangp-magic-mask-panel:not(.hide)')).some((panel) => WMM.isVisible(panel));
    if (openPanel) return;
    setTimeout(() => WMM.focusVisibleImageEditor(true), 150);
    setTimeout(() => WMM.focusVisibleImageEditor(true), 700);
};

WMM.installImageEditorFocusPatch = function () {
    if (WMM.imageEditorFocusPatchInstalled) return;
    WMM.imageEditorFocusPatchInstalled = true;
    document.addEventListener('pointerdown', (event) => {
        const editor = event.target?.closest?.('#img_editor, [data-testid="image-editor"], .imageeditor, .image-editor')?.closest?.('.block') || event.target?.closest?.('#img_editor, [data-testid="image-editor"], .imageeditor, .image-editor');
        if (!editor || !WMM.isVisible(editor)) return;
        if (event.target?.closest?.('button, input, textarea, select, [role="button"]')) return;
        WMM.focusImageEditor(editor, false);
    }, true);
};

WMM.openMagicMaskPanelInAnchor = function (anchor) {
    const panel = anchor?.querySelector?.('.wangp-magic-mask-panel');
    if (!panel) return false;
    panel.hidden = false;
    panel.classList.remove('hide');
    panel.style.display = '';
    return true;
};

WMM.installOverlayTriggerPatch = function () {
    document.querySelectorAll('.wangp-magic-mask-trigger--overlay').forEach((trigger) => {
        if (trigger.dataset.wangpMagicMaskOverlayBound === '1') return;
        trigger.dataset.wangpMagicMaskOverlayBound = '1';
        trigger.addEventListener('click', () => {
            WMM.openMagicMaskPanelInAnchor(trigger.closest('.wangp-magic-mask-anchor'));
        }, true);
    });
};

WMM.findImageEditorForTrigger = function (trigger) {
    const roots = [
        trigger.closest('.wangp-magic-mask-anchor--image-editor'),
        trigger.parentElement,
        trigger.closest('.column'),
        trigger.closest('[id^="component-"]')?.parentElement,
    ].filter(Boolean);
    for (const root of roots) {
        const editor = Array.from(root.querySelectorAll('#img_editor, [data-testid="image-editor"], .imageeditor, .image-editor'))
            .find((candidate) => {
                const block = candidate.closest('.block') || candidate.closest('[id^="component-"]') || candidate;
                return WMM.isVisible(block);
            });
        if (editor) return editor.closest('.block') || editor.closest('[id^="component-"]') || editor;
    }
    const previousEditors = Array.from(document.querySelectorAll('#img_editor, [data-testid="image-editor"], .imageeditor, .image-editor'))
        .map((candidate) => candidate.closest('.block') || candidate.closest('[id^="component-"]') || candidate)
        .filter((editor) => WMM.isVisible(editor) && editor.compareDocumentPosition(trigger) & Node.DOCUMENT_POSITION_FOLLOWING);
    return previousEditors.pop() || null;
};

WMM.mountImageEditorTriggers = function () {
    WMM.installOverlayTriggerPatch();
    document.querySelectorAll('.wangp-magic-mask-trigger--editor').forEach((trigger) => {
        const anchor = trigger.closest('.wangp-magic-mask-anchor--image-editor') || trigger.parentElement || document.body;
        if (trigger.classList.contains('hidden') || !!trigger.closest('.hidden') || !!trigger.closest('.hide') || !WMM.isVisible(anchor)) return;
        const editor = WMM.findImageEditorForTrigger(trigger);
        if (!editor || !trigger) return;
        const toolbar = WMM.findImageEditorToolbar(editor);
        if (!toolbar) {
            return;
        }
        anchor.querySelectorAll('.wangp-magic-mask-toolbar-button').forEach((button) => {
            if (button.parentElement !== toolbar) button.remove();
        });
        let toolbarButton = toolbar.querySelector('.wangp-magic-mask-toolbar-button');
        if (!toolbarButton) {
            toolbarButton = document.createElement('button');
            toolbarButton.type = 'button';
            toolbarButton.className = 'wangp-magic-mask-toolbar-button';
            toolbarButton.setAttribute('aria-label', 'Magic Mask');
            toolbarButton.setAttribute('title', 'Magic Mask');
            toolbarButton.textContent = '\u{1FA84}';
            toolbar.appendChild(toolbarButton);
        }
        toolbar.classList.add('wangp-magic-mask-toolbar');
        toolbarButton.hidden = trigger.classList.contains('hidden') || !!trigger.closest('.hidden');
        toolbarButton.disabled = trigger.disabled;
        const needsImage = /Upload an image/i.test(editor.innerText || '') && /select the draw tool to start/i.test(editor.innerText || '');
        toolbarButton.classList.toggle('wangp-magic-mask-unavailable', needsImage);
        toolbarButton.title = needsImage ? 'Magic Mask needs a control image' : 'Magic Mask';
        toolbarButton.onclick = (event) => {
            event.preventDefault();
            event.stopPropagation();
            if (needsImage) return;
            trigger.click();
        };
    });
};

WMM.scheduleMount = function () {
    if (WMM.raf) cancelAnimationFrame(WMM.raf);
    WMM.raf = requestAnimationFrame(WMM.mountImageEditorTriggers);
};

if (!WMM.init) {
    WMM.init = true;
    WMM.observer = new MutationObserver(WMM.scheduleMount);
    const root = document.querySelector('gradio-app') || document.body;
    if (root) WMM.observer.observe(root, { childList: true, subtree: true });
    window.addEventListener('resize', WMM.scheduleMount);
    window.addEventListener('load', WMM.scheduleMount);
    WMM.interval = window.setInterval(WMM.scheduleMount, 500);
}
WMM.installImageEditorFocusPatch();
WMM.scheduleMount();
"""

    @staticmethod
    def focus_image_editor_javascript():
        return "() => { window.__wangpMagicMaskNS?.refocusImageEditorAfterMagicMask?.(); }"

    def render(self, visible=False, trigger_mode="overlay"):
        self.abort_token = gr.State(str(uuid.uuid4()))
        self.pending_image_mask_guide = gr.State(None)
        self.pending_image_mask = gr.State(None)
        self.trigger = gr.Button(MAGIC_WAND_LABEL, size="sm", min_width=1, visible=visible, elem_classes=["wangp-magic-mask-trigger", f"wangp-magic-mask-trigger--{trigger_mode}"])
        with gr.Group(visible=False, elem_classes=["wangp-magic-mask-panel"]) as self.panel:
            with gr.Column(elem_classes=["wangp-magic-mask-card"]):
                gr.HTML("<div class='wangp-magic-mask-titlebar'><div class='wangp-magic-mask-heading'>Magic Mask</div></div>")
                with gr.Column(elem_classes=["wangp-magic-mask-body"]):
                    gr.HTML("<div class='wangp-magic-mask-intro'>Enter the list of Object or Persons to track and that will be used to build the Mask. Each object / person should be separated by a \",\". For example: \"blue car, woman to the right\"</div>")
                    with gr.Row(elem_classes=["wangp-magic-mask-keyword-row"]):
                        self.keywords = gr.Textbox(show_label=False, placeholder="person, car, sky", lines=1, scale=4, elem_classes=["wangp-magic-mask-keywords"])
                        with gr.Group(elem_classes=["wangp-magic-mask-negative"]):
                            self.negative_mask = gr.Checkbox(label="Negative Mask", value=False, container=False, min_width=1, elem_classes=["wangp-magic-mask-negative-checkbox"])
                self.status = gr.HTML("")
                self.progress_html = gr.HTML("")
                gr.HTML("", elem_classes=["wangp-magic-mask-spacer"], padding=False)
                with gr.Row(elem_classes=["wangp-magic-mask-actions"]):
                    self.cancel_btn = gr.Button("Exit", size="sm", elem_classes=["wangp-magic-mask-btn"])
                    self.abort_btn = gr.Button("Abort", size="sm", visible=False, elem_classes=["wangp-magic-mask-btn", "wangp-magic-mask-btn--danger"])
                    self.generate_btn = gr.Button("Generate", size="sm", elem_classes=["wangp-magic-mask-btn", "wangp-magic-mask-btn--primary"])
        return self

    def mount(
        self,
        *,
        state,
        image_mode,
        video_guide,
        image_mask_guide,
        image_guide,
        image_mask,
        video_mask,
        download_assets: Callable[[dict[str, Any]], Any],
        acquire_gpu: Callable[[Any, str, str], Any],
        release_gpu: Callable[[Any, str], Any],
        get_model_settings: Callable[[Any], dict],
    ):
        self.trigger.click(fn=_open_panel, inputs=[], outputs=[self.panel, self.status, self.progress_html, self.cancel_btn, self.abort_btn, self.pending_image_mask_guide, self.pending_image_mask], show_progress="hidden")
        self.cancel_btn.click(
            fn=_close_panel,
            inputs=[],
            outputs=[self.panel, self.status, self.progress_html, self.abort_btn, self.pending_image_mask_guide, self.pending_image_mask],
            show_progress="hidden",
        )
        self.abort_btn.click(fn=_abort_magic_mask, inputs=[self.abort_token], outputs=[self.status, self.abort_btn], show_progress="hidden")

        def generate(state_value, keywords_text, negative_mask_value, image_mode_value, video_guide_value, image_mask_guide_value, image_guide_value, abort_token_value):
            yield from _generate_magic_mask(
                state_value,
                keywords_text,
                negative_mask_value,
                image_mode_value,
                video_guide_value,
                image_mask_guide_value,
                image_guide_value,
                abort_token_value,
                download_assets=download_assets,
                acquire_gpu=acquire_gpu,
                release_gpu=release_gpu,
                get_model_settings=get_model_settings,
            )

        generate_event = self.generate_btn.click(
            fn=generate,
            inputs=[state, self.keywords, self.negative_mask, image_mode, video_guide, image_mask_guide, image_guide, self.abort_token],
            outputs=[image_mask_guide, image_mask, video_mask, self.status, self.panel, self.progress_html, self.cancel_btn, self.abort_btn, self.pending_image_mask_guide, self.pending_image_mask],
            show_progress="hidden",
        )
        generate_event.then(fn=None, inputs=[], outputs=[], js=MagicMaskUI.focus_image_editor_javascript())