Balloon Button

"use client";

import { useEffect, useRef, type CSSProperties, type ButtonHTMLAttributes, type PointerEvent } from "react";
import * as THREE from "three";
import { cn } from "@/lib/utils";
import { fragment, vertex } from "./shader";

// Room around the button for the shadow and the swell on press.
const PAD = 24;
const STIFFNESS = 520;
// A press always lasts at least this long (ms). A quick click or tap often delivers pointerdown and
// pointerup in the same frame, which would reset the spring before it moved and show no dent.
const MIN_PRESS = 140;
const msSince = (t: number) => performance.now() - t;

export type BalloonButtonProps = Omit<ButtonHTMLAttributes<HTMLButtonElement>, "color"> & {
  label?: string;
  /** Balloon colour, any CSS colour three.js understands. */
  color?: string;
  /** Corner radius in px. */
  radius?: number;
  /** How far the surface swells, relative to half the button height. */
  puff?: number;
  /** Distance between folds along the edge, px. */
  creaseSpacing?: number;
  /** How deep the folds pinch in. */
  creaseDepth?: number;
  /** How much the dent springs back past rest on release, 0 to 0.9. */
  bounce?: number;
};

type Gl = {
  renderer: THREE.WebGLRenderer;
  scene: THREE.Scene;
  camera: THREE.Camera;
  material: THREE.ShaderMaterial;
};

export function BalloonButton({
  label = "Balloon Button",
  color = "#000000",
  radius = 18,
  puff = 0.7,
  creaseSpacing = 14,
  creaseDepth = 0.4,
  bounce = 0.2,
  className,
  onPointerDown,
  onPointerUp,
  onPointerLeave,
  onPointerCancel,
  onKeyDown,
  onKeyUp,
  ...rest
}: BalloonButtonProps) {
  const buttonRef = useRef<HTMLButtonElement>(null);
  const labelRef = useRef<HTMLSpanElement>(null);
  const glRef = useRef<Gl | null>(null);

  // Animation state lives in a ref so the loop never re-renders React.
  const sim = useRef({
    x: 0,
    v: 0,
    target: 0,
    waveT: 99,
    waveAmp: 0,
    raf: 0,
    last: 0,
    damping: 0,
    reduced: false,
    pressedAt: 0,
    releaseTimer: 0,
  });
  const renderRef = useRef<() => void>(() => {});
  const drawRef = useRef<() => void>(() => {});

  useEffect(() => {
    const button = buttonRef.current;
    if (!button) return;

    // A fresh canvas (and so a fresh WebGL context) for every setup. When the effect runs again on
    // the same element (StrictMode, or Next.js showing a kept-alive page again after navigating
    // back), a new renderer on the old canvas would inherit the old context's GL state, its
    // uniforms would go to the wrong program, and the press would stop showing.
    const canvas = document.createElement("canvas");
    canvas.setAttribute("aria-hidden", "true");
    canvas.className = "pointer-events-none absolute";
    Object.assign(canvas.style, { inset: `-${PAD}px`, width: `calc(100% + ${PAD * 2}px)`, height: `calc(100% + ${PAD * 2}px)` });

    let renderer: THREE.WebGLRenderer;
    try {
      renderer = new THREE.WebGLRenderer({ canvas, alpha: true, antialias: false, premultipliedAlpha: true });
    } catch {
      // No WebGL: fall back to a flat fill.
      button.dataset.fallback = "";
      return;
    }
    button.prepend(canvas);
    delete button.dataset.fallback;
    renderer.setClearColor(0x000000, 0);
    // Full device resolution (browser zoom raises it too); the canvas is small.
    const pixelRatio = () => Math.min(window.devicePixelRatio || 1, 8);
    renderer.setPixelRatio(pixelRatio());

    const material = new THREE.ShaderMaterial({
      vertexShader: vertex,
      fragmentShader: fragment,
      uniforms: {
        uCanvas: { value: new THREE.Vector2() },
        uDpr: { value: renderer.getPixelRatio() },
        uHalf: { value: new THREE.Vector2() },
        uRadius: { value: 28 },
        uColor: { value: new THREE.Color() },
        uPuff: { value: 0.9 },
        uCreaseSpacing: { value: 11 },
        uCreaseDepth: { value: 0.5 },
        uPress: { value: 0 },
        uPoint: { value: new THREE.Vector2() },
        uWaveT: { value: 99 },
        uWaveAmp: { value: 0 },
        uLabel: { value: null },
        uInk: { value: new THREE.Color() },
      },
    });

    // The label is painted into a texture and printed onto the surface in the
    // shader, so it bends with the folds and dents. The HTML text stays for
    // layout and screen readers but is invisible.
    const labelCanvas = document.createElement("canvas");
    const labelTexture = new THREE.CanvasTexture(labelCanvas);
    labelTexture.minFilter = THREE.LinearFilter;
    labelTexture.generateMipmaps = false;
    material.uniforms.uLabel.value = labelTexture;
    const scene = new THREE.Scene();
    const geometry = new THREE.PlaneGeometry(2, 2);
    scene.add(new THREE.Mesh(geometry, material));
    const camera = new THREE.Camera();
    glRef.current = { renderer, scene, camera, material };

    const render = () => renderer.render(scene, camera);
    renderRef.current = render;

    const drawLabel = () => {
      const span = labelRef.current;
      if (!span) return;
      const w = button.offsetWidth + PAD * 2;
      const h = button.offsetHeight + PAD * 2;
      // Supersampled 2x past device resolution so the warped label stays crisp.
      const device = renderer.getPixelRatio();
      const dpr = device * Math.min(2, 4096 / (w * device));
      const cw = Math.round(w * dpr);
      const ch = Math.round(h * dpr);
      // three.js allocates texture storage once; a new size needs a fresh texture.
      if (labelCanvas.width !== cw || labelCanvas.height !== ch) labelTexture.dispose();
      labelCanvas.width = cw;
      labelCanvas.height = ch;
      const ctx = labelCanvas.getContext("2d");
      if (!ctx) return;
      const cs = getComputedStyle(span);
      ctx.scale(dpr, dpr);
      ctx.font = `${cs.fontStyle} ${cs.fontWeight} ${cs.fontSize} ${cs.fontFamily}`;
      if ("letterSpacing" in ctx) ctx.letterSpacing = cs.letterSpacing === "normal" ? "0px" : cs.letterSpacing;
      ctx.textAlign = "center";
      ctx.textBaseline = "alphabetic";
      ctx.fillStyle = "#fff";
      // Centre the font's line box where the HTML text sits.
      const m = ctx.measureText(span.textContent ?? "");
      const sr = span.getBoundingClientRect();
      const br = button.getBoundingClientRect();
      const cx = PAD + sr.left - br.left + sr.width / 2;
      const cy = PAD + sr.top - br.top + sr.height / 2;
      ctx.fillText(span.textContent ?? "", cx, cy + (m.fontBoundingBoxAscent - m.fontBoundingBoxDescent) / 2);
      material.uniforms.uInk.value.set(cs.color);
      labelTexture.needsUpdate = true;
      render();
    };
    button.dataset.gl = "";
    drawRef.current = drawLabel;
    document.fonts?.ready.then(drawLabel);

    const resize = () => {
      const w = button.offsetWidth;
      const h = button.offsetHeight;
      renderer.setSize(w + PAD * 2, h + PAD * 2, false);
      material.uniforms.uCanvas.value.set(w + PAD * 2, h + PAD * 2);
      material.uniforms.uHalf.value.set(w / 2, h / 2);
      drawLabel();
    };
    const ro = new ResizeObserver(resize);
    ro.observe(button);
    resize();

    // Re-render sharp when the page moves to another screen or is zoomed.
    let dprQuery: MediaQueryList;
    const onDpr = () => {
      renderer.setPixelRatio(pixelRatio());
      material.uniforms.uDpr.value = renderer.getPixelRatio();
      resize();
      watchDpr();
    };
    const watchDpr = () => {
      dprQuery?.removeEventListener("change", onDpr);
      dprQuery = window.matchMedia(`(resolution: ${window.devicePixelRatio}dppx)`);
      dprQuery.addEventListener("change", onDpr);
    };
    watchDpr();

    const onLost = (e: Event) => {
      e.preventDefault();
      delete button.dataset.gl;
      button.dataset.fallback = "";
    };
    canvas.addEventListener("webglcontextlost", onLost);

    const s = sim.current;
    const mq = window.matchMedia("(prefers-reduced-motion: reduce)");
    s.reduced = mq.matches;

    return () => {
      cancelAnimationFrame(s.raf);
      s.raf = 0;
      clearTimeout(s.releaseTimer);
      ro.disconnect();
      dprQuery.removeEventListener("change", onDpr);
      canvas.removeEventListener("webglcontextlost", onLost);
      geometry.dispose();
      labelTexture.dispose();
      material.dispose();
      renderer.dispose();
      // free the context now (browsers allow only a handful) and drop its canvas
      renderer.forceContextLoss();
      canvas.remove();
      glRef.current = null;
      delete button.dataset.gl;
    };
  }, []);

  // Props feed straight into uniforms.
  useEffect(() => {
    const gl = glRef.current;
    if (!gl) return;
    const u = gl.material.uniforms;
    u.uRadius.value = radius;
    u.uColor.value.set(color);
    u.uPuff.value = puff;
    u.uCreaseSpacing.value = Math.max(creaseSpacing, 2);
    u.uCreaseDepth.value = creaseDepth;
    renderRef.current();
  }, [color, radius, puff, creaseSpacing, creaseDepth]);

  // The label ink comes from the text colour, so a className change repaints it.
  useEffect(() => {
    drawRef.current();
  }, [label, className]);

  useEffect(() => {
    // Damping ratio from bounce: 0 settles without overshoot, 0.9 rings.
    const zeta = 1 - Math.min(Math.max(bounce, 0), 0.9);
    sim.current.damping = 2 * Math.sqrt(STIFFNESS) * zeta;
  }, [bounce]);

  const tick = (now: number) => {
    const s = sim.current;
    const gl = glRef.current;
    if (!gl) {
      // torn down mid-loop: clear the handle, or kick() would think a loop is still running
      s.raf = 0;
      return;
    }
    // the first frame's timestamp can be a little before kick() ran; never step backwards
    const dt = Math.min(Math.max((now - s.last) / 1000, 0), 1 / 30);
    s.last = now;

    // Spring toward the target in small steps for stability.
    const damping = s.reduced ? 2 * Math.sqrt(STIFFNESS) : s.damping;
    for (let i = 0; i < 4; i++) {
      const h = dt / 4;
      s.v += (-STIFFNESS * (s.x - s.target) - damping * s.v) * h;
      s.x += s.v * h;
    }
    s.waveT += dt;

    const u = gl.material.uniforms;
    u.uPress.value = s.x;
    u.uWaveT.value = s.waveT;
    u.uWaveAmp.value = s.reduced ? 0 : s.waveAmp;
    gl.renderer.render(gl.scene, gl.camera);

    const settled = Math.abs(s.x - s.target) < 0.0005 && Math.abs(s.v) < 0.0005 && s.waveT > 2;
    if (settled) {
      s.x = s.target;
      s.raf = 0;
      return;
    }
    s.raf = requestAnimationFrame(tick);
  };

  const kick = () => {
    const s = sim.current;
    if (s.raf) return;
    s.last = performance.now();
    s.raf = requestAnimationFrame(tick);
  };

  const press = (x: number, y: number) => {
    const s = sim.current;
    clearTimeout(s.releaseTimer);
    s.pressedAt = performance.now();
    glRef.current?.material.uniforms.uPoint.value.set(x, y);
    s.target = 1;
    s.waveT = 0;
    s.waveAmp = 1;
    kick();
  };

  const release = () => {
    const s = sim.current;
    if (s.target === 0) return;
    const wait = MIN_PRESS - msSince(s.pressedAt);
    clearTimeout(s.releaseTimer);
    if (wait > 0) {
      // let the dent form before springing back
      s.releaseTimer = window.setTimeout(release, wait);
      return;
    }
    s.target = 0;
    kick();
  };

  const handlePointerDown = (e: PointerEvent<HTMLButtonElement>) => {
    onPointerDown?.(e);
    const rect = e.currentTarget.getBoundingClientRect();
    press(e.clientX - rect.left - rect.width / 2, rect.height / 2 - (e.clientY - rect.top));
  };

  return (
    <button
      ref={buttonRef}
      type="button"
      {...rest}
      onPointerDown={handlePointerDown}
      onPointerUp={(e) => {
        onPointerUp?.(e);
        release();
      }}
      onPointerLeave={(e) => {
        onPointerLeave?.(e);
        release();
      }}
      onPointerCancel={(e) => {
        onPointerCancel?.(e);
        release();
      }}
      onKeyDown={(e) => {
        onKeyDown?.(e);
        if ((e.key === "Enter" || e.key === " ") && !e.repeat) press(0, 0);
      }}
      onKeyUp={(e) => {
        onKeyUp?.(e);
        if (e.key === "Enter" || e.key === " ") release();
      }}
      className={cn(
        "relative inline-flex h-14 cursor-pointer touch-manipulation items-center justify-center px-9 text-base font-semibold tracking-tight text-white select-none data-fallback:bg-(--balloon) focus-visible:outline-2 focus-visible:outline-offset-4 focus-visible:outline-zinc-400",
        className,
      )}
      style={{ borderRadius: radius, WebkitTapHighlightColor: "transparent", "--balloon": color } as CSSProperties}
    >
      <span ref={labelRef} className="relative in-data-gl:opacity-0">
        {label}
      </span>
    </button>
  );
}