#!/usr/bin/env python3
"""Fresh offline image-to-vector measurements for the local benchmark page.

Runs two tracers that are actually installed on this machine:

  1. the project engine (python/any2svg_engine.py, which drives vtracer)
  2. potrace (brew), on a 1-bit thresholded copy of the same input

Inputs: the two checked-in public PNGs plus one synthetic flat-colour image
generated with Pillow. Fully offline - no network, no upload.

Intermediate SVG/PBM artifacts go to a temp run directory that is deleted at
the end, so this folder only gains scripts and raw JSON (write whitelist).
Raw numbers land in raw-measurements.json next to this script.

Usage: venv/bin/python measure.py
"""
from __future__ import annotations

import hashlib
import importlib.metadata
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import tempfile
import time

from PIL import Image

MEAS_DIR = os.path.dirname(os.path.abspath(__file__))
REPO = "/Users/szj/Downloads/tmp/any2svg"
ENGINE = os.path.join(REPO, "python", "any2svg_engine.py")
PUBLIC = os.path.join(REPO, "public")
OUT_JSON = os.path.join(MEAS_DIR, "raw-measurements.json")

INPUTS = [
    os.path.join(PUBLIC, "apple-touch-icon.png"),
    os.path.join(PUBLIC, "og-image.png"),
]

D_ATTR_RE = re.compile('d="([^"]*)"')
M_CMD_RE = re.compile("[Mm]")
PATH_EL_RE = re.compile("<path[\\s>]")
IMAGE_EL_RE = re.compile(r"<image\b.*?(?:/>|</image>)", re.S)
FILL_RE = re.compile('(?:fill|stroke)="([^"]+)"')
# Inkscape writes colour as style="fill:#rrggbb;..." instead of fill="#rrggbb".
FILL_STYLE_RE = re.compile(r"(?:^|;)\s*(?:fill|stroke)\s*:\s*([^;\"}]+)")
COLOUR_EXCLUDE = {"none", "inherit", "opacity", "currentcolor"}


def sha256(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 16), b""):
            h.update(chunk)
    return h.hexdigest()


def detect_mode_like_engine(img: Image.Image) -> dict:
    """Replicates any2svg_engine.detect_mode() so the number is comparable."""
    small = img.convert("RGB")
    small.thumbnail((256, 256), Image.LANCZOS)
    px = small.load()
    seen = set()
    for y in range(small.height):
        for x in range(small.width):
            r, g, b = px[x, y]
            seen.add((r >> 4, g >> 4, b >> 4))
            if len(seen) > 48:
                return {
                    "quantised_colours_256_thumb": len(seen),
                    "early_exit_over_48": True,
                    "engine_mode": "spline",
                }
    return {
        "quantised_colours_256_thumb": len(seen),
        "early_exit_over_48": False,
        "engine_mode": "polygon",
    }


def describe_input(path: str) -> dict:
    img = Image.open(path)
    rgb = img.convert("RGB")
    colours = rgb.getcolors(maxcolors=1 << 24)
    info = {
        "path": path,
        "bytes": os.path.getsize(path),
        "sha256": sha256(path),
        "width": img.width,
        "height": img.height,
        "mode": img.mode,
        "distinct_rgb_colours": len(colours) if colours else None,
    }
    info.update(detect_mode_like_engine(img))
    return info


def svg_stats(svg_path: str) -> dict:
    text = open(svg_path, encoding="utf-8").read()
    d_blobs = D_ATTR_RE.findall(text)
    fills = set(FILL_RE.findall(text))
    for style in re.findall(r'style="([^"]*)"', text):
        fills.update(FILL_STYLE_RE.findall(style))
    colours = {c.strip() for c in fills
               if c.strip().lower() not in COLOUR_EXCLUDE}
    return {
        "bytes": os.path.getsize(svg_path),
        "sha256": sha256(svg_path),
        "m_commands": sum(len(M_CMD_RE.findall(d)) for d in d_blobs),
        "path_elements": len(PATH_EL_RE.findall(text)),
        "image_elements": len(IMAGE_EL_RE.findall(text)),
        "bytes_without_embedded_image": len(
            IMAGE_EL_RE.sub("", text).encode("utf-8")),
        "distinct_fill_or_stroke_values": len(fills),
        "distinct_colors": len(colours),
        "fill_values": sorted(colours)[:20],
        "contains_image_element": "<image" in text,
    }


def run_engine(src: str, rundir: str, name: str) -> dict:
    outdir = os.path.join(rundir, "engine", name)
    os.makedirs(outdir, exist_ok=True)
    cmd = [sys.executable, ENGINE, "--input", src, "--outdir", outdir,
           "--preset", "balanced", "--mode", "auto"]
    t0 = time.perf_counter()
    proc = subprocess.run(cmd, capture_output=True, text=True, timeout=900)
    wall = time.perf_counter() - t0
    rec = {
        "tool": "any2svg engine (python/any2svg_engine.py, vtracer)",
        "command": " ".join(cmd),
        "wall_seconds": round(wall, 3),
        "exit_code": proc.returncode,
        "stderr_tail": proc.stderr.strip()[-500:],
    }
    try:
        manifest = json.loads(proc.stdout.strip().splitlines()[-1])
    except Exception:
        rec["error"] = "engine stdout was not a JSON manifest"
        rec["stdout_tail"] = proc.stdout.strip()[-500:]
        return rec
    if "error" in manifest:
        rec["error"] = manifest["error"]
        return rec
    svg = manifest["svg"]
    rec["engine_mode"] = manifest.get("mode")
    rec["engine_canvas"] = manifest.get("canvas")
    rec["output"] = svg_stats(svg)
    rec["output_path_ephemeral"] = svg
    return rec


def run_potrace(src: str, rundir: str, name: str) -> dict:
    potrace = shutil.which("potrace")
    if not potrace:
        return {"tool": "potrace", "error": "potrace not found on PATH"}
    prep_dir = os.path.join(rundir, "potrace")
    os.makedirs(prep_dir, exist_ok=True)
    pbm = os.path.join(prep_dir, name + ".pbm")
    out_svg = os.path.join(prep_dir, name + ".svg")

    t0 = time.perf_counter()
    img = Image.open(src).convert("L")
    bw = img.point(lambda v: 255 if v >= 128 else 0, mode="1")
    bw.save(pbm)
    prep_wall = time.perf_counter() - t0

    cmd = [potrace, "-s", "-o", out_svg, pbm]
    t1 = time.perf_counter()
    proc = subprocess.run(cmd, capture_output=True, text=True, timeout=900)
    trace_wall = time.perf_counter() - t1

    rec = {
        "tool": "potrace (1-bit threshold trace)",
        "command": " ".join(cmd),
        "prep_seconds_pillow_threshold": round(prep_wall, 3),
        "wall_seconds": round(trace_wall, 3),
        "wall_seconds_total_prep_plus_trace": round(prep_wall + trace_wall, 3),
        "exit_code": proc.returncode,
        "stderr_tail": proc.stderr.strip()[-500:],
        "note": ("potrace traces monochrome only, so the input is converted to "
                 "1-bit at threshold 128 first; output is 1 colour layer"),
    }
    if proc.returncode != 0 or not os.path.exists(out_svg):
        rec["error"] = "potrace failed"
        return rec
    rec["output"] = svg_stats(out_svg)
    rec["output_path_ephemeral"] = out_svg
    return rec


def run_inkscape(src: str, rundir: str, name: str) -> dict:
    """Headless Inkscape trace via the object-trace action.

    Argument format (from `inkscape --action-list` / the error it prints when
    arguments are wrong): scans, smooth, stack, remove_background, speckles,
    smooth_corners, optimize. Values are parsed as numbers, not true/false.
    Inkscape keeps the source bitmap in the document, so bytes_without_
    embedded_image is reported alongside the as-exported size.
    """
    inkscape = "/opt/homebrew/bin/inkscape"
    if not os.path.exists(inkscape):
        return {"tool": "inkscape", "error": "inkscape not installed"}
    outdir = os.path.join(rundir, "inkscape")
    os.makedirs(outdir, exist_ok=True)
    out_svg = os.path.join(outdir, name + ".svg")
    actions = (f"select-all;object-trace:8,1,0,0,2,1,1;"
               f"export-filename:{out_svg};export-do")
    cmd = [inkscape, f"--actions={actions}", src]
    t0 = time.perf_counter()
    proc = subprocess.run(cmd, capture_output=True, text=True, timeout=900)
    wall = time.perf_counter() - t0
    rec = {
        "tool": "inkscape Trace Bitmap (object-trace action)",
        "command": " ".join(cmd),
        "wall_seconds": round(wall, 3),
        "exit_code": proc.returncode,
        "stdout_tail": (proc.stdout or "").strip()[-300:],
        "stderr_tail": (proc.stderr or "").strip()[-300:],
        "note": ("8 scans, smooth=1, stack=0, remove_background=0, "
                 "speckles=2, smooth_corners=1, optimize=1; source bitmap is "
                 "left in the document, so as-exported bytes include it"),
    }
    if not os.path.exists(out_svg):
        rec["error"] = "no SVG produced"
        return rec
    rec["output"] = svg_stats(out_svg)
    rec["output_path_ephemeral"] = out_svg
    return rec


def tool_versions() -> dict:
    def pip_ver(name: str) -> str:
        try:
            return importlib.metadata.version(name)
        except Exception:
            return "not installed"

    def cmd_ver(argv) -> str:
        exe = shutil.which(argv[0])
        if not exe:
            return "not installed"
        try:
            out = subprocess.run(argv, capture_output=True, text=True, timeout=30)
            return (out.stdout or out.stderr).strip().splitlines()[0]
        except Exception as exc:
            return f"failed: {exc}"

    return {
        "python": sys.version.split()[0] + " (" + sys.executable + ")",
        "vtracer": pip_ver("vtracer"),
        "pillow": pip_ver("pillow"),
        "potrace": cmd_ver(["potrace", "--version"]),
        "inkscape": cmd_ver(["inkscape", "--version"]),
        "homebrew": cmd_ver(["brew", "--version"]),
        "platform": platform.platform(),
        "machine": platform.machine(),
        "cpu_count": os.cpu_count(),
    }


def make_synthetic(rundir: str) -> str:
    """Deterministic flat-colour artwork: 5 colours, hard edges, no gradients."""
    path = os.path.join(rundir, "synthetic-flat.png")
    img = Image.new("RGB", (800, 600), (255, 255, 255))
    px = img.load()
    bands = [((30, 60, 180), 0, 150), ((240, 200, 40), 150, 300),
             ((200, 40, 60), 300, 450), ((40, 150, 90), 450, 600)]
    for colour, y0, y1 in bands:
        for y in range(y0, y1):
            for x in range(800):
                px[x, y] = colour
    for y in range(600):
        for x in range(640, 800):
            px[x, y] = (30, 30, 30)
    for cy in range(200, 400):
        for cx in range(100, 300):
            if (cx - 200) ** 2 + (cy - 300) ** 2 <= 90 ** 2:
                px[cx, cy] = (255, 255, 255)
    img.save(path)
    return path


def main() -> int:
    rundir = tempfile.mkdtemp(prefix="any2svg-bench-")
    inputs = list(INPUTS)
    inputs.append(make_synthetic(rundir))

    results = {
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "generator": os.path.abspath(__file__),
        "repo": REPO,
        "offline": True,
        "tool_versions": tool_versions(),
        "runs": [],
    }

    for src in inputs:
        name = os.path.splitext(os.path.basename(src))[0]
        entry = {"input": describe_input(src), "measurements": []}
        try:
            entry["measurements"].append(run_engine(src, rundir, name))
        except Exception as exc:
            entry["measurements"].append(
                {"tool": "any2svg engine", "error": f"{type(exc).__name__}: {exc}"})
        try:
            entry["measurements"].append(run_potrace(src, rundir, name))
        except Exception as exc:
            entry["measurements"].append(
                {"tool": "potrace", "error": f"{type(exc).__name__}: {exc}"})
        try:
            entry["measurements"].append(run_inkscape(src, rundir, name))
        except Exception as exc:
            entry["measurements"].append(
                {"tool": "inkscape", "error": f"{type(exc).__name__}: {exc}"})
        results["runs"].append(entry)
        print(f"done: {name}", flush=True)

    with open(OUT_JSON, "w", encoding="utf-8") as fh:
        json.dump(results, fh, indent=2, sort_keys=False)
        fh.write("\n")

    shutil.rmtree(rundir, ignore_errors=True)
    print(f"wrote {OUT_JSON}")
    return 0


if __name__ == "__main__":
    sys.exit(main())