"""
Vector Studio — PNG/JPG to SVG
Clean line work for generated images.

Hugging Face Space (Gradio SDK). Run locally: python app.py
"""

import os
import re
import base64
import tempfile
import uuid

import numpy as np
from PIL import Image, ImageOps, ImageFilter
import gradio as gr
import vtracer

MAX_SIDE = 2600          # Cap after upscaling, keeps free CPU Spaces from timing out
OUT_DIR = os.path.join(tempfile.gettempdir(), "svg_out")
os.makedirs(OUT_DIR, exist_ok=True)


# --------------------------------------------------------------------------
# Image preparation
# --------------------------------------------------------------------------
def flatten(img: Image.Image) -> Image.Image:
    """Composite alpha onto white so transparency isn't traced as black."""
    img = ImageOps.exif_transpose(img)
    if img.mode in ("RGBA", "LA", "P"):
        img = img.convert("RGBA")
        bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
        img = Image.alpha_composite(bg, img)
    return img.convert("RGB")


def otsu_threshold(arr: np.ndarray) -> int:
    hist = np.bincount(arr.ravel(), minlength=256).astype(np.float64)
    total = hist.sum()
    omega = np.cumsum(hist) / total
    mu = np.cumsum(hist * np.arange(256)) / total
    mu_t = mu[-1]
    denom = omega * (1.0 - omega)
    denom[denom == 0] = 1e-12
    sigma_b = (mu_t * omega - mu) ** 2 / denom
    return int(np.argmax(sigma_b))


def binarize(gray: Image.Image, method: str, level: int, block: int, offset: int) -> np.ndarray:
    """Return a boolean mask where True means ink."""
    arr = np.asarray(gray, dtype=np.uint8)

    if method.startswith("Adaptive"):
        radius = max(1, int(block) // 2)
        local = np.asarray(gray.filter(ImageFilter.BoxBlur(radius)), dtype=np.int16)
        return arr.astype(np.int16) < (local - int(offset))

    if method.startswith("Otsu"):
        level = otsu_threshold(arr)

    return arr < int(level)


def zhang_suen(mask: np.ndarray, max_iter: int = 60) -> np.ndarray:
    """Skeletonize: thin every stroke down to a 1 px centerline."""
    img = mask.copy()
    for _ in range(max_iter):
        changed = False
        for step in (0, 1):
            P = np.pad(img, 1, constant_values=False)
            P2, P3, P4 = P[:-2, 1:-1], P[:-2, 2:], P[1:-1, 2:]
            P5, P6, P7 = P[2:, 2:], P[2:, 1:-1], P[2:, :-2]
            P8, P9 = P[1:-1, :-2], P[:-2, :-2]

            ring = [P2, P3, P4, P5, P6, P7, P8, P9, P2]
            B = sum(x.astype(np.uint8) for x in ring[:8])
            A = np.zeros(img.shape, np.uint8)
            for i in range(8):
                A += (~ring[i] & ring[i + 1]).astype(np.uint8)

            base = img & (B >= 2) & (B <= 6) & (A == 1)
            if step == 0:
                cond = base & ~(P2 & P4 & P6) & ~(P4 & P6 & P8)
            else:
                cond = base & ~(P2 & P4 & P8) & ~(P2 & P6 & P8)

            if cond.any():
                img &= ~cond
                changed = True
        if not changed:
            break
    return img


def set_stroke_width(mask: np.ndarray, width: int) -> np.ndarray:
    """Grow the skeleton back out to one uniform stroke width."""
    if width <= 1:
        return mask
    ink = Image.fromarray(np.where(mask, 255, 0).astype(np.uint8))
    size = width if width % 2 == 1 else width + 1
    ink = ink.filter(ImageFilter.MaxFilter(size))
    return np.asarray(ink) > 127


def cap_size(img: Image.Image) -> Image.Image:
    w, h = img.size
    if max(w, h) <= MAX_SIDE:
        return img
    s = MAX_SIDE / max(w, h)
    return img.resize((max(1, int(w * s)), max(1, int(h * s))), Image.LANCZOS)


# --------------------------------------------------------------------------
# SVG post-processing
# --------------------------------------------------------------------------
def strip_background(svg: str) -> str:
    """Drop near-white fills, which is what the background becomes in color mode."""
    def is_light(hexcol: str) -> bool:
        r, g, b = (int(hexcol[i:i + 2], 16) for i in (1, 3, 5))
        return r > 243 and g > 243 and b > 243

    return re.sub(
        r'<path[^>]*fill="(#[0-9a-fA-F]{6})"[^>]*/>\s*',
        lambda m: "" if is_light(m.group(1)) else m.group(0),
        svg,
    )


def recolor(svg: str, color: str) -> str:
    return re.sub(r'fill="#000000"', f'fill="{color}"', svg)


def rescale_root(svg: str, out_w: int, out_h: int) -> str:
    """Add a viewBox and set the display size back to the source dimensions."""
    m = re.search(r'<svg([^>]*)width="(\d+)"\s+height="(\d+)"', svg)
    if not m:
        return svg
    vw, vh = m.group(2), m.group(3)
    new_tag = (
        f'<svg{m.group(1)}width="{out_w}" height="{out_h}" '
        f'viewBox="0 0 {vw} {vh}"'
    )
    return svg[: m.start()] + new_tag + svg[m.end():]


def preview_html(svg: str, label: str) -> str:
    if len(svg) > 4_000_000:
        return (
            f"<div class='vs-note'>{label} is large ({len(svg) // 1024} KB). "
            "Download it instead of previewing here, or raise "
            "<em>Filter small specks</em> to cut the path count.</div>"
        )
    b64 = base64.b64encode(svg.encode("utf-8")).decode("ascii")
    return (
        "<div class='vs-stage'>"
        f"<img src='data:image/svg+xml;base64,{b64}' alt='{label}'/>"
        "</div>"
    )


# --------------------------------------------------------------------------
# Main pipeline
# --------------------------------------------------------------------------
def vectorize(
    image,
    mode,
    invert,
    thr_method,
    thr_level,
    block,
    offset,
    denoise,
    upscale,
    uniform,
    stroke_w,
    curve_mode,
    speckle,
    corner,
    length_thr,
    splice,
    precision,
    colors,
    layer_diff,
    hierarchical,
    drop_bg,
    line_color,
    progress=gr.Progress(),
):
    if image is None:
        raise gr.Error("Upload a PNG or JPG first.")

    progress(0.1, desc="Preparing image")
    src = flatten(image)
    orig_w, orig_h = src.size

    work = src
    if upscale > 1:
        work = work.resize(
            (int(orig_w * upscale), int(orig_h * upscale)), Image.LANCZOS
        )
    work = cap_size(work)

    line_mode = mode.startswith("Line")

    if line_mode:
        progress(0.3, desc="Isolating lines")
        gray = ImageOps.autocontrast(work.convert("L"), cutoff=1)
        if invert:
            gray = ImageOps.invert(gray)
        if denoise > 0:
            gray = gray.filter(ImageFilter.MedianFilter(int(denoise) * 2 + 1))

        mask = binarize(gray, thr_method, thr_level, block, offset)

        if uniform:
            progress(0.45, desc="Finding centerlines")
            mask = set_stroke_width(zhang_suen(mask), int(stroke_w))

        prepped = Image.fromarray(np.where(mask, 0, 255).astype(np.uint8)).convert("RGB")
    else:
        progress(0.35, desc="Reducing color areas")
        prepped = work
        if denoise > 0:
            prepped = prepped.filter(ImageFilter.MedianFilter(int(denoise) * 2 + 1))

    stem = uuid.uuid4().hex[:10]
    tmp_png = os.path.join(OUT_DIR, f"{stem}.png")
    out_svg = os.path.join(OUT_DIR, "vector.svg" if line_mode else "vector-color.svg")
    prepped.save(tmp_png)

    progress(0.6, desc="Tracing paths")
    vtracer.convert_image_to_svg_py(
        tmp_png,
        out_svg,
        colormode="binary" if line_mode else "color",
        hierarchical="cutout" if hierarchical == "Cutout" else "stacked",
        mode="polygon" if curve_mode.startswith("Polygon") else "spline",
        filter_speckle=int(speckle),
        color_precision=int(colors),
        layer_difference=int(layer_diff),
        corner_threshold=int(corner),
        length_threshold=float(length_thr),
        max_iterations=10,
        splice_threshold=int(splice),
        path_precision=int(precision),
    )

    progress(0.85, desc="Cleaning up SVG")
    svg = open(out_svg, "r", encoding="utf-8").read()
    if drop_bg:
        svg = strip_background(svg)
    if line_mode:
        svg = recolor(svg, line_color)
    svg = rescale_root(svg, orig_w, orig_h)

    with open(out_svg, "w", encoding="utf-8") as f:
        f.write(svg)
    try:
        os.remove(tmp_png)
    except OSError:
        pass

    n_paths = svg.count("<path")
    stats = (
        f"**{n_paths} path{'' if n_paths == 1 else 's'}** · {len(svg) / 1024:.1f} KB · "
        f"canvas {orig_w}×{orig_h} px · traced at {prepped.size[0]}×{prepped.size[1]} px"
    )
    return preview_html(svg, "Vector result"), out_svg, stats


# --------------------------------------------------------------------------
# Interface
# --------------------------------------------------------------------------
CSS = """
.vs-stage {
  background-color: #f7f5f2;
  background-image:
    linear-gradient(45deg, #e6e2dc 25%, transparent 25%, transparent 75%, #e6e2dc 75%),
    linear-gradient(45deg, #e6e2dc 25%, transparent 25%, transparent 75%, #e6e2dc 75%);
  background-size: 18px 18px;
  background-position: 0 0, 9px 9px;
  border-radius: 10px; padding: 14px; min-height: 240px;
  display: flex; align-items: center; justify-content: center;
}
.vs-stage img { max-width: 100%; max-height: 62vh; }
.vs-note { padding: 18px; font-size: 0.9rem; line-height: 1.5; }
footer { display: none !important; }
"""

# Gradio 6 wants theme/css on launch(), Gradio 5 on the Blocks constructor.
_MAJOR = int(gr.__version__.split(".")[0])
_STYLE = {"theme": gr.themes.Soft(), "css": CSS}
_BLOCKS_KW = {} if _MAJOR >= 6 else _STYLE
_LAUNCH_KW = _STYLE if _MAJOR >= 6 else {}

with gr.Blocks(title="PNG/JPG to SVG", **_BLOCKS_KW) as demo:
    gr.Markdown(
        "## PNG / JPG to SVG\n"
        "Turn bitmaps into clean, scalable paths. "
        "Built for AI-generated drawings, logos and sketches."
    )

    with gr.Row():
        # ---------------- left column: input and controls ----------------
        with gr.Column(scale=4):
            image = gr.Image(label="Image", type="pil", sources=["upload", "clipboard"],
                             image_mode="RGBA", height=300)
            mode = gr.Radio(
                ["Line art (black & white)", "Color (filled shapes)"],
                value="Line art (black & white)",
                label="Mode",
            )
            run = gr.Button("Vectorize", variant="primary")

            with gr.Accordion("Isolate lines", open=True) as line_box:
                thr_method = gr.Radio(
                    ["Otsu (automatic)", "Global", "Adaptive (uneven lighting)"],
                    value="Otsu (automatic)",
                    label="Threshold",
                )
                thr_level = gr.Slider(0, 255, 128, step=1, label="Cutoff (Global only)")
                block = gr.Slider(5, 151, 41, step=2, label="Window size (Adaptive only)")
                offset = gr.Slider(0, 40, 8, step=1, label="Sensitivity (Adaptive only)")
                invert = gr.Checkbox(False, label="Invert (light lines on dark background)")
                denoise = gr.Slider(0, 3, 1, step=1, label="Smooth noise")
                upscale = gr.Slider(1, 4, 2, step=1,
                                    label="Upscale before tracing (smoother curves)")
                uniform = gr.Checkbox(
                    False, label="Uniform stroke width (redraw from centerlines)"
                )
                stroke_w = gr.Slider(1, 9, 3, step=2, label="Stroke width in px")

            with gr.Accordion("Path quality", open=False):
                curve_mode = gr.Radio(
                    ["Spline (smooth curves)", "Polygon (hard edges)"],
                    value="Spline (smooth curves)", label="Curve type",
                )
                speckle = gr.Slider(0, 40, 6, step=1, label="Filter small specks")
                corner = gr.Slider(0, 180, 60, step=1, label="Corner threshold")
                length_thr = gr.Slider(0.5, 10, 4, step=0.5, label="Shortest segment")
                splice = gr.Slider(0, 180, 45, step=1, label="Curve splicing")
                precision = gr.Slider(1, 8, 4, step=1, label="Coordinate precision")

            with gr.Accordion("Color and output", open=False):
                colors = gr.Slider(1, 8, 6, step=1, label="Color depth (color mode)")
                layer_diff = gr.Slider(0, 64, 16, step=1, label="Layer spacing (color mode)")
                hierarchical = gr.Radio(
                    ["Stacked", "Cutout"], value="Stacked", label="Layer structure"
                )
                drop_bg = gr.Checkbox(True, label="Remove white background")
                line_color = gr.ColorPicker("#111111", label="Line color")

        # ---------------- right column: result ----------------
        with gr.Column(scale=6):
            preview = gr.HTML("<div class='vs-stage'><span class='vs-note'>"
                              "Your result appears here.</span></div>",
                              label="Preview")
            stats = gr.Markdown("")
            download = gr.File(label="Download SVG", height=90)

    gr.Markdown(
        "**For clean lines:** upscale 2–3×, use spline curves, set speck filtering to 6–12. "
        "If a drawing comes out ragged, turn on *Uniform stroke width* — it pulls every "
        "line onto a centerline and redraws it at a constant weight."
    )

    def toggle(m):
        return gr.update(open=m.startswith("Line"))

    mode.change(toggle, mode, line_box)

    inputs = [image, mode, invert, thr_method, thr_level, block, offset, denoise,
              upscale, uniform, stroke_w, curve_mode, speckle, corner, length_thr,
              splice, precision, colors, layer_diff, hierarchical, drop_bg, line_color]

    run.click(vectorize, inputs, [preview, download, stats])

if __name__ == "__main__":
    demo.queue(max_size=20).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        ssr_mode=False,
        **_LAUNCH_KW,
    )