"""
Moondream Segmentation Playground
=================================
Gradio 6.8.0 Space recreating the moondream.ai playground segmentation demo.

Key features matching the official playground:
- Multi-object segmentation: model.point() finds all instances, then
  model.segment() with spatial_refs segments each one individually
- SVG path trace animation: stroke-dashoffset traces the outline before
  the fill overlay fades in
- Toggle switch for "Find all instances" (enabled by default)
- gr.HTML custom component for SVG overlay rendering

Requires: MOONDREAM_API_KEY env var (free key at console.moondream.ai)
"""

import gradio as gr
import moondream as md
from PIL import Image
import base64
import io
import os
import time

# ---------------------------------------------------------------------------
# Moondream client
# ---------------------------------------------------------------------------

API_KEY = os.environ.get("MOONDREAM_API_KEY", "")


def get_model():
    if not API_KEY:
        raise gr.Error(
            "MOONDREAM_API_KEY not set. "
            "Get a free key at https://console.moondream.ai "
            "and add it as a Space secret."
        )
    return md.vl(api_key=API_KEY)


# ---------------------------------------------------------------------------
# Color palette (matches playground style)
# ---------------------------------------------------------------------------

SEGMENT_COLORS = [
    {"fill": "rgba(255, 45, 32, 0.50)", "stroke": "#ff2d20"},
    {"fill": "rgba(30, 110, 255, 0.50)", "stroke": "#1e6eff"},
    {"fill": "rgba(16, 200, 80, 0.50)", "stroke": "#10c850"},
    {"fill": "rgba(255, 160, 10, 0.50)", "stroke": "#ffa00a"},
    {"fill": "rgba(180, 60, 230, 0.50)", "stroke": "#b43ce6"},
    {"fill": "rgba(255, 40, 90, 0.50)", "stroke": "#ff285a"},
    {"fill": "rgba(0, 200, 255, 0.50)", "stroke": "#00c8ff"},
    {"fill": "rgba(255, 210, 0, 0.50)", "stroke": "#ffd200"},
]


# ---------------------------------------------------------------------------
# Image helpers
# ---------------------------------------------------------------------------


def image_to_data_uri(pil_image: Image.Image, max_dim: int = 1024) -> str:
    img = pil_image.copy()
    img.thumbnail((max_dim, max_dim), Image.LANCZOS)
    if img.mode == "RGBA":
        img = img.convert("RGB")
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/jpeg;base64,{b64}"


def _to_pil(image) -> Image.Image:
    if isinstance(image, Image.Image):
        return image
    if isinstance(image, str):
        return Image.open(image)
    return Image.fromarray(image)


# ---------------------------------------------------------------------------
# HTML builder with SVG path trace animation
# ---------------------------------------------------------------------------

# Three-phase animation using pathLength="1" on the SVG path element so
# stroke-dasharray/offset of 1 gives consistent timing regardless of path
# complexity.  Segments animate one-by-one sequentially.
#
# Phase 1 (seg-trace):   Bright thick stroke traces the path outline
# Phase 2 (seg-fill-in): Semi-transparent fill fades in over the area
# Phase 3 (seg-settle):  Stroke thins down to final border weight

SEG_STAGGER = 1.3  # seconds between each segment's animation start

ANIMATION_CSS = """
<style>
@keyframes seg-trace {
    0%   { stroke-dashoffset: 1; }
    100% { stroke-dashoffset: 0; }
}
@keyframes seg-fill-in {
    0%   { fill-opacity: 0; }
    100% { fill-opacity: 1; }
}
@keyframes seg-settle {
    0%   { stroke-width: 0.025; }
    100% { stroke-width: 0.006; }
}
.seg-path {
    stroke-dasharray: 1;
    stroke-dashoffset: 1;
    stroke-width: 0.025;
    stroke-linecap: round;
    stroke-linejoin: round;
    fill-opacity: 0;
    animation:
        seg-trace 1.0s ease-in-out forwards,
        seg-fill-in 0.4s ease-out 0.7s forwards,
        seg-settle 0.3s ease-out 1.0s forwards;
}
</style>
"""


def build_empty_html() -> str:
    return (
        '<div style="display:flex; align-items:center; justify-content:center;'
        " min-height:420px; border:2px dashed #444; border-radius:12px;"
        " color:#777; font-family:system-ui,sans-serif; font-size:15px;"
        ' background:#111827; flex-direction:column; gap:8px;">'
        '<svg width="48" height="48" viewBox="0 0 24 24" fill="none"'
        ' stroke="#555" stroke-width="1.5" stroke-linecap="round"'
        ' stroke-linejoin="round">'
        '<rect x="3" y="3" width="18" height="18" rx="2" ry="2"/>'
        '<circle cx="8.5" cy="8.5" r="1.5"/>'
        '<polyline points="21 15 16 10 5 21"/>'
        "</svg>"
        "<span>Upload an image and enter a segmentation prompt</span>"
        "</div>"
    )


def build_image_html(
    data_uri: str,
    segments: list | None = None,
    animate: bool = True,
) -> str:
    """
    Build HTML with base image + SVG overlays.

    Each segment dict:
        path     : SVG path string, coords [0,1] relative to bbox
        bbox     : {"x_min", "y_min", "x_max", "y_max"} relative to image
        prompt   : the text prompt
        color_idx: palette index
    """
    svg_overlays = ""
    legend_items = ""

    css_block = ANIMATION_CSS if (animate and segments) else ""

    if segments:
        for i, seg in enumerate(segments):
            bbox = seg["bbox"]
            color = SEGMENT_COLORS[seg["color_idx"] % len(SEGMENT_COLORS)]

            left = bbox["x_min"] * 100
            top = bbox["y_min"] * 100
            width = (bbox["x_max"] - bbox["x_min"]) * 100
            height = (bbox["y_max"] - bbox["y_min"]) * 100

            # Sequential stagger: each segment fully completes before next starts
            if animate:
                stagger = i * SEG_STAGGER
                anim_class = "seg-path"
                static_fill = ""
                # Three animation-delay values: trace, fill-in, settle
                delay_style = (
                    f"animation-delay: {stagger:.2f}s,"
                    f" {stagger + 0.7:.2f}s,"
                    f" {stagger + 1.0:.2f}s;"
                )
                sw_attr = ""  # CSS handles stroke-width for animated
                pl_attr = 'pathLength="1"'
            else:
                anim_class = ""
                static_fill = 'fill-opacity="1"'
                delay_style = ""
                sw_attr = 'stroke-width="0.006"'
                pl_attr = ""

            svg_overlays += (
                f'<svg xmlns="http://www.w3.org/2000/svg"'
                f' viewBox="0 0 1 1" preserveAspectRatio="none"'
                f" style=\"position:absolute; left:{left:.4f}%; top:{top:.4f}%;"
                f" width:{width:.4f}%; height:{height:.4f}%;"
                f' pointer-events:none; z-index:2;">'
                f'<path class="{anim_class}" {pl_attr} d="{seg["path"]}"'
                f' fill="{color["fill"]}"'
                f' stroke="{color["stroke"]}"'
                f" {sw_attr}"
                f" {static_fill}"
                f' style="{delay_style}"/>'
                f"</svg>\n"
            )

        # Build legend -- group by prompt, show count
        prompt_counts: dict[str, int] = {}
        prompt_color: dict[str, dict] = {}
        for seg in segments:
            p = seg["prompt"]
            prompt_counts[p] = prompt_counts.get(p, 0) + 1
            if p not in prompt_color:
                prompt_color[p] = SEGMENT_COLORS[
                    seg["color_idx"] % len(SEGMENT_COLORS)
                ]

        for prompt, count in prompt_counts.items():
            color = prompt_color[prompt]
            count_label = f" x{count}" if count > 1 else ""
            legend_items += (
                '<span style="display:inline-flex; align-items:center; gap:5px;'
                " padding:3px 10px; border-radius:6px;"
                f" background:{color['stroke']}22;"
                f" border:1px solid {color['stroke']}55;"
                f' font-size:13px; color:{color["stroke"]};">'
                '<span style="width:10px; height:10px; border-radius:50%;'
                f' background:{color["stroke"]}; display:inline-block;"></span>'
                f"{prompt}{count_label}"
                "</span>\n"
            )

    legend_html = ""
    if legend_items:
        legend_html = (
            '<div style="display:flex; flex-wrap:wrap; gap:6px; padding:8px 0;">'
            f"{legend_items}</div>"
        )

    return (
        f"{css_block}"
        '<div style="font-family:system-ui,sans-serif;">'
        '<div style="position:relative; display:inline-block; max-width:100%;'
        " border-radius:8px; overflow:hidden; background:#000;"
        ' line-height:0;">'
        f'<img src="{data_uri}"'
        ' style="max-width:100%; max-height:600px; display:block;'
        ' width:auto; height:auto;"'
        ' draggable="false" />'
        f"{svg_overlays}"
        "</div>"
        f"{legend_html}"
        "</div>"
    )


# ---------------------------------------------------------------------------
# Core segmentation logic
# ---------------------------------------------------------------------------


def segment_objects(
    pil_img: Image.Image,
    prompt: str,
    find_all: bool,
    existing_segments: list,
) -> list:
    """
    Segment objects in the image.

    If find_all is True:
      1. model.point(image, prompt) to find all matching points
      2. model.segment(image, prompt, spatial_refs=[[x, y]]) for each point

    If find_all is False:
      Single model.segment(image, prompt) call.

    Returns the updated segment list.
    """
    model = get_model()
    segments = list(existing_segments) if existing_segments else []
    base_color_idx = len(segments)

    if find_all:
        # Step 1: detect all instances
        try:
            point_result = model.point(pil_img, prompt)
        except Exception as e:
            raise gr.Error(f"Moondream point detection error: {e}")

        points = point_result.get("points", [])
        if not points:
            raise gr.Error(
                f"No instances of '{prompt}' found in the image. "
                "Try a different description."
            )

        # Step 2: segment each detected instance
        for i, pt in enumerate(points):
            try:
                result = model.segment(
                    pil_img,
                    prompt,
                    spatial_refs=[[pt["x"], pt["y"]]],
                )
            except Exception as e:
                raise gr.Error(
                    f"Moondream segment error for point {i + 1}: {e}"
                )

            if result and "path" in result and "bbox" in result:
                segments.append(
                    {
                        "path": result["path"],
                        "bbox": result["bbox"],
                        "prompt": prompt,
                        "color_idx": base_color_idx,
                    }
                )
    else:
        # Single object segmentation
        try:
            result = model.segment(pil_img, prompt)
        except Exception as e:
            raise gr.Error(f"Moondream segment error: {e}")

        if not result or "path" not in result or "bbox" not in result:
            raise gr.Error(
                f"No segmentation mask returned for '{prompt}'. "
                "Try a different description."
            )

        segments.append(
            {
                "path": result["path"],
                "bbox": result["bbox"],
                "prompt": prompt,
                "color_idx": base_color_idx,
            }
        )

    return segments


# ---------------------------------------------------------------------------
# Gradio event handlers
# ---------------------------------------------------------------------------


def on_image_upload(image):
    if image is None:
        return build_empty_html(), []
    pil_img = _to_pil(image)
    data_uri = image_to_data_uri(pil_img)
    return build_image_html(data_uri, animate=False), []


def on_segment(image, prompt, find_all, segments_state):
    if image is None:
        raise gr.Error("Upload an image first.")
    if not prompt or not prompt.strip():
        raise gr.Error(
            "Enter a segmentation prompt (e.g. 'cat', 'boots', 'person')."
        )

    prompt = prompt.strip()
    pil_img = _to_pil(image)

    segments = segment_objects(
        pil_img, prompt, find_all, segments_state
    )

    data_uri = image_to_data_uri(pil_img)

    # Animate only the newly added segments -- rebuild all but only
    # new ones get the animation class. We do this by building a
    # mixed list: old segments without animation, new ones with.
    old_count = len(segments_state) if segments_state else 0
    html = build_segmentation_html(data_uri, segments, old_count)

    found_count = len(segments) - old_count
    status = f"Found {found_count} instance(s) of '{prompt}'"

    return html, segments, "", status


def build_segmentation_html(
    data_uri: str,
    segments: list,
    animate_from_idx: int,
) -> str:
    """
    Build HTML where segments before animate_from_idx are static
    and segments from animate_from_idx onward get the trace animation.
    """
    has_new = animate_from_idx < len(segments)
    css_block = ANIMATION_CSS if has_new else ""

    svg_overlays = ""
    legend_items = ""

    new_anim_counter = 0

    for i, seg in enumerate(segments):
        bbox = seg["bbox"]
        color = SEGMENT_COLORS[seg["color_idx"] % len(SEGMENT_COLORS)]

        left = bbox["x_min"] * 100
        top = bbox["y_min"] * 100
        width = (bbox["x_max"] - bbox["x_min"]) * 100
        height = (bbox["y_max"] - bbox["y_min"]) * 100

        is_new = i >= animate_from_idx
        if is_new:
            stagger = new_anim_counter * SEG_STAGGER
            new_anim_counter += 1
            anim_class = "seg-path"
            static_fill = ""
            delay_style = (
                f"animation-delay: {stagger:.2f}s,"
                f" {stagger + 0.7:.2f}s,"
                f" {stagger + 1.0:.2f}s;"
            )
            sw_attr = ""
            pl_attr = 'pathLength="1"'
        else:
            anim_class = ""
            static_fill = 'fill-opacity="1"'
            delay_style = ""
            sw_attr = 'stroke-width="0.006"'
            pl_attr = ""

        svg_overlays += (
            f'<svg xmlns="http://www.w3.org/2000/svg"'
            f' viewBox="0 0 1 1" preserveAspectRatio="none"'
            f" style=\"position:absolute; left:{left:.4f}%; top:{top:.4f}%;"
            f" width:{width:.4f}%; height:{height:.4f}%;"
            f' pointer-events:none; z-index:2;">'
            f'<path class="{anim_class}" {pl_attr} d="{seg["path"]}"'
            f' fill="{color["fill"]}"'
            f' stroke="{color["stroke"]}"'
            f" {sw_attr}"
            f" {static_fill}"
            f' style="{delay_style}"/>'
            f"</svg>\n"
        )

    # Legend grouped by prompt with counts
    prompt_counts: dict[str, int] = {}
    prompt_color: dict[str, dict] = {}
    for seg in segments:
        p = seg["prompt"]
        prompt_counts[p] = prompt_counts.get(p, 0) + 1
        if p not in prompt_color:
            prompt_color[p] = SEGMENT_COLORS[
                seg["color_idx"] % len(SEGMENT_COLORS)
            ]

    for prompt, count in prompt_counts.items():
        color = prompt_color[prompt]
        count_label = f" x{count}" if count > 1 else ""
        legend_items += (
            '<span style="display:inline-flex; align-items:center; gap:5px;'
            " padding:3px 10px; border-radius:6px;"
            f" background:{color['stroke']}22;"
            f" border:1px solid {color['stroke']}55;"
            f' font-size:13px; color:{color["stroke"]};">'
            '<span style="width:10px; height:10px; border-radius:50%;'
            f' background:{color["stroke"]}; display:inline-block;"></span>'
            f"{prompt}{count_label}"
            "</span>\n"
        )

    legend_html = ""
    if legend_items:
        legend_html = (
            '<div style="display:flex; flex-wrap:wrap; gap:6px; padding:8px 0;">'
            f"{legend_items}</div>"
        )

    return (
        f"{css_block}"
        '<div style="font-family:system-ui,sans-serif;">'
        '<div style="position:relative; display:inline-block; max-width:100%;'
        " border-radius:8px; overflow:hidden; background:#000;"
        ' line-height:0;">'
        f'<img src="{data_uri}"'
        ' style="max-width:100%; max-height:600px; display:block;'
        ' width:auto; height:auto;"'
        ' draggable="false" />'
        f"{svg_overlays}"
        "</div>"
        f"{legend_html}"
        "</div>"
    )


def on_clear_segments(image):
    if image is None:
        return build_empty_html(), [], ""
    pil_img = _to_pil(image)
    data_uri = image_to_data_uri(pil_img)
    return build_image_html(data_uri, animate=False), [], ""


def on_undo_segment(image, segments_state):
    if image is None:
        return build_empty_html(), [], ""
    pil_img = _to_pil(image)
    segments = list(segments_state) if segments_state else []

    if segments:
        # Remove the last prompt group (all segments with the last prompt)
        last_prompt = segments[-1]["prompt"]
        while segments and segments[-1]["prompt"] == last_prompt:
            segments.pop()

    data_uri = image_to_data_uri(pil_img)
    html = build_segmentation_html(data_uri, segments, len(segments))
    return html, segments, ""


# ---------------------------------------------------------------------------
# CSS
# ---------------------------------------------------------------------------

CUSTOM_CSS = """
.main-wrap { max-width: 1040px !important; margin: 0 auto !important; }
.seg-display { min-height: 320px; }
.seg-display > div { background: transparent !important; }
.prompt-row { gap: 8px !important; }
.action-btn { min-width: 90px !important; }
.toggle-row { padding: 4px 0 !important; }
.status-text { min-height: 24px; }
.status-text span { font-size: 13px; color: #999; }
footer { display: none !important; }
"""

# ---------------------------------------------------------------------------
# Build the app
# ---------------------------------------------------------------------------

with gr.Blocks(
    css=CUSTOM_CSS,
    theme=gr.themes.Base(
        primary_hue="indigo",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
    ),
    title="Moondream Segmentation Playground",
) as demo:

    segments_state = gr.State([])

    # Header
    gr.HTML(
        '<div style="text-align:center; padding:16px 0 4px;">'
        '<h1 style="margin:0 0 4px; font-size:1.5em; font-weight:700;">'
        "Moondream Segmentation Playground</h1>"
        '<p style="margin:0; color:#999; font-size:0.9em;">'
        "Open-vocabulary image segmentation &mdash; describe any object "
        "to segment all instances &mdash; powered by "
        '<a href="https://moondream.ai" target="_blank"'
        ' style="color:#818cf8; text-decoration:none;">Moondream</a></p>'
        "</div>"
    )

    with gr.Row(elem_classes="main-wrap"):

        # Left column: controls
        with gr.Column(scale=1, min_width=280):

            image_input = gr.Image(
                label="Upload Image",
                type="pil",
                sources=["upload", "clipboard"],
                height=260,
            )

            prompt_input = gr.Textbox(
                label="Object to Segment",
                placeholder="e.g. boot, cat, person, sky, cracked tiles...",
                lines=1,
                max_lines=1,
                submit_btn="Segment",
            )

            find_all_toggle = gr.Checkbox(
                label="Find all instances",
                value=True,
                info="Detect and segment every matching object in the image",
                elem_classes="toggle-row",
            )

            with gr.Row(elem_classes="prompt-row"):
                clear_btn = gr.Button(
                    "Clear All",
                    variant="secondary",
                    size="sm",
                    elem_classes="action-btn",
                )
                undo_btn = gr.Button(
                    "Undo Last",
                    variant="secondary",
                    size="sm",
                    elem_classes="action-btn",
                )

            status_output = gr.Markdown(
                value="", elem_classes="status-text"
            )

            gr.Markdown(
                "**How to use**\n\n"
                "1. Upload or paste an image\n"
                "2. Type what you want to segment\n"
                "3. Press Enter or click **Segment**\n"
                "4. With *Find all instances* on, every matching object is "
                "detected via `model.point()` then segmented individually\n"
                "5. Add more prompts for different object types\n"
                "6. Each prompt group gets a distinct color\n\n"
                "*Moondream cloud API &mdash; ~5s per object &mdash; "
                "free $5/mo at "
                "[console.moondream.ai](https://console.moondream.ai)*"
            )

        # Right column: visualization
        with gr.Column(scale=2, min_width=400):
            seg_display = gr.HTML(
                value=build_empty_html(),
                label="Segmentation Result",
                elem_classes="seg-display",
            )

    # Prompt suggestions
    gr.Markdown(
        "**Try these prompts** &nbsp; "
        "`boot` &middot; `cat` &middot; `dog` &middot; `person` &middot; "
        "`sky` &middot; `car` &middot; `tree` &middot; `building` &middot; "
        "`water` &middot; `cracked tiles` &middot; "
        "`the person on the left` &middot; `red object`"
    )

    # -----------------------------------------------------------------
    # Event wiring
    # -----------------------------------------------------------------

    image_input.change(
        fn=on_image_upload,
        inputs=[image_input],
        outputs=[seg_display, segments_state],
    )

    prompt_input.submit(
        fn=on_segment,
        inputs=[image_input, prompt_input, find_all_toggle, segments_state],
        outputs=[seg_display, segments_state, prompt_input, status_output],
        show_progress="minimal",
    )

    clear_btn.click(
        fn=on_clear_segments,
        inputs=[image_input],
        outputs=[seg_display, segments_state, status_output],
    )

    undo_btn.click(
        fn=on_undo_segment,
        inputs=[image_input, segments_state],
        outputs=[seg_display, segments_state, status_output],
    )


# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    demo.launch(
        show_error=True,
    )