Tool Image Gen

"use client"

import * as React from "react"
import { ToolDisplayState, ToolImageGen } from "@/components/ui/tool-image-gen"

const FIT_CSS = "html,body{overflow:hidden}[data-fit-small] [data-fit-hide]{display:none!important}"

/**
 * 21st shows a demo in a small frame: about 674 by 421 on a component page, 358 by 224 on a
 * phone and 245 by 139 on a profile card. A demo taller than its frame gets scrollbars, which
 * render white on dark pages. FitFrame fills the frame exactly instead.
 *
 * - "scale": the content shrinks until it fits, centered. For visuals such as orbs.
 * - "crop": the content stays at full size from the top and fades out at the bottom. For text
 *   that streams in, where a changing scale would make the page breathe.
 *
 * `minWidth` lays out a wide demo (a table, a two column panel) at that width and scales it down
 * to the frame, instead of letting it collapse to a narrow layout. The frame centers it, so it
 * shrinks evenly on both sides.
 *
 * Anything marked `data-fit-hide` (state buttons, the chatcn.dev line) drops out in frames
 * shorter than `hideBelow`, so a profile card shows the visual alone.
 */
function FitFrame({
  children,
  mode = "scale",
  hideBelow = 360,
  minWidth,
}: {
  children: React.ReactNode
  mode?: "scale" | "crop"
  hideBelow?: number
  minWidth?: string
}) {
  const content = React.useRef<HTMLDivElement>(null)
  const [cropped, setCropped] = React.useState(false)

  React.useLayoutEffect(() => {
    const el = content.current
    if (!el) return
    const measure = () => {
      const vw = window.innerWidth
      const vh = window.innerHeight
      el.toggleAttribute("data-fit-small", vh < hideBelow)
      // Layout sizes ignore the transform, so measuring never feeds back.
      const w = Math.max(el.offsetWidth, el.scrollWidth)
      const h = Math.max(el.offsetHeight, el.scrollHeight)
      const scale = mode === "scale" ? Math.min(1, vw / w, vh / h) : 1
      el.style.transform = scale < 1 ? `scale(${scale})` : ""
      setCropped(mode === "crop" && h > vh + 1)
    }
    measure()
    const observer = new ResizeObserver(measure)
    observer.observe(el)
    // Content that changes inside a fixed width box (a headline that cycles words) never resizes
    // the box itself, so changes below it measure again on the next frame.
    let frame = 0
    const mutations = new MutationObserver(() => {
      cancelAnimationFrame(frame)
      frame = requestAnimationFrame(measure)
    })
    mutations.observe(el, { subtree: true, childList: true, characterData: true })
    window.addEventListener("resize", measure)
    return () => {
      observer.disconnect()
      mutations.disconnect()
      cancelAnimationFrame(frame)
      window.removeEventListener("resize", measure)
    }
  }, [mode, hideBelow])

  const fade = "linear-gradient(to bottom, #000 calc(100% - 40px), transparent)"
  return (
    <div
      className="flex w-full justify-center overflow-hidden"
      style={{
        height: "100dvh",
        alignItems: mode === "crop" ? "flex-start" : "center",
        maskImage: cropped ? fade : undefined,
        WebkitMaskImage: cropped ? fade : undefined,
      }}
    >
      <style>{FIT_CSS}</style>
      <div
        ref={content}
        className="w-full shrink-0"
        style={{ transformOrigin: mode === "crop" ? "top center" : "center", minWidth }}
      >
        {children}
      </div>
    </div>
  )
}

interface JobOptions {
  /** Queue positions to step through before generating. Empty skips the queue. */
  queue?: number[]
  queueStepMs?: number
  /** Time from 0 to 100 percent. */
  durationMs?: number
  /** Where the job ends. Default "done". */
  outcome?: "done" | "failed" | "blocked"
  /** Percent at which a failed or blocked job stops. */
  stopAt?: number
  /** Start on mount. Default true. */
  autoStart?: boolean
}

type JobStatus = "queued" | "generating" | "done" | "failed" | "blocked"

/** Small seeded PRNG (mulberry32): same seed, same picture. */
function random(seed: number) {
  let a = seed >>> 0
  return () => {
    a = (a + 0x6d2b79f5) >>> 0
    let t = a
    t = Math.imul(t ^ (t >>> 15), t | 1)
    t ^= t + Math.imul(t ^ (t >>> 7), t | 61)
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296
  }
}

/**
 * A scripted generation job for demos: walks the queue, then reports progress in uneven steps
 * every 240 ms, the way real providers do, and lands on the outcome. `run()` starts it again.
 */
function useGenerationJob({
  queue = [2, 1],
  queueStepMs = 900,
  durationMs = 5600,
  outcome = "done",
  stopAt = 46,
  autoStart = true,
}: JobOptions = {}) {
  const [state, setState] = React.useState<{
    status: JobStatus
    progress?: number
    queuePosition?: number
  }>({ status: "queued", queuePosition: queue[0] })
  const [runId, setRunId] = React.useState(autoStart ? 1 : 0)
  const options = React.useRef({ queue, queueStepMs, durationMs, outcome, stopAt })
  options.current = { queue, queueStepMs, durationMs, outcome, stopAt }

  React.useEffect(() => {
    if (runId === 0) return
    const { queue, queueStepMs, durationMs, outcome, stopAt } = options.current
    const timers: number[] = []
    const rand = random(runId * 31)
    const at = (ms: number, fn: () => void) => timers.push(window.setTimeout(fn, ms))
    setState({ status: "queued", queuePosition: queue[0] })
    queue.forEach((position, i) => {
      at(i * queueStepMs, () => setState({ status: "queued", queuePosition: position }))
    })
    const start = queue.length * queueStepMs
    at(start, () => setState({ status: "generating", progress: 0 }))
    const steps = Math.max(4, Math.round(durationMs / 240))
    let p = 0
    let t = start
    for (let i = 1; i <= steps; i++) {
      t += 240
      p = Math.min(100, p + (100 / steps) * (0.4 + rand() * 1.2))
      const value = i === steps ? 100 : Math.min(99, Math.round(p))
      if (outcome !== "done" && value >= stopAt) {
        at(t, () => setState({ status: outcome, progress: value }))
        break
      }
      at(t, () => setState({ status: "generating", progress: value }))
      if (i === steps) at(t + 280, () => setState({ status: "done", progress: 100 }))
    }
    return () => {
      for (const id of timers) window.clearTimeout(id)
    }
  }, [runId])

  const run = React.useCallback(() => setRunId((id) => id + 1), [])
  return { ...state, run }
}

type Scene = "dunes" | "peaks" | "sphere" | "ripples" | "stairs" | "moon"

/** Prompts that match each scene, written the way people prompt image models. */
const PROMPTS: Record<Scene, string> = {
  dunes:
    "Dune ridges at dusk, long raking shadows, medium format film, fine grain, black and white",
  peaks: "Mountain ranges layered in morning fog, telephoto compression, quiet, monochrome",
  sphere: "A matte sphere resting on a seamless studio sweep, single soft key light, monochrome",
  ripples: "Concentric ripples on black water, shot from above at night, high contrast",
  stairs: "Concrete stairs in hard noon light, deep diagonal shadows, black and white",
  moon: "A low moon over a calm sea, its reflection breaking into lines, long exposure",
}

const PROMPT = PROMPTS.peaks

interface SceneOptions {
  /** Seed for the shapes. Different seeds give different variations of one scene. */
  seed?: number
  /** Width over height. Default 1. */
  ratio?: number
  /** Long edge in px. Default 1200. */
  size?: number
}

const r1 = (n: number) => Math.round(n * 10) / 10

/** A smooth ridge line across the frame, closed down to the bottom edge. */
function ridge(
  w: number,
  h: number,
  base: number,
  amp: number,
  rand: () => number,
  jagged = false,
) {
  const steps = jagged ? 9 : 5
  const pts: [number, number][] = []
  for (let i = 0; i <= steps; i++) {
    const x = (i / steps) * w
    const y = base + (rand() - 0.5) * amp * (jagged && i % 2 ? 1.6 : 1)
    pts.push([x, y])
  }
  let d = `M0 ${h} L0 ${r1(pts[0]?.[1] ?? base)}`
  for (let i = 1; i < pts.length; i++) {
    const [x0, y0] = pts[i - 1] ?? [0, base]
    const [x1, y1] = pts[i] ?? [w, base]
    if (jagged) d += ` L${r1(x1)} ${r1(y1)}`
    else {
      const cx = (x0 + x1) / 2
      d += ` C${r1(cx)} ${r1(y0)} ${r1(cx)} ${r1(y1)} ${r1(x1)} ${r1(y1)}`
    }
  }
  return `${d} L${w} ${h} Z`
}

function grain(w: number, h: number, seed: number, opacity = 0.16) {
  return `<filter id="g" x="0" y="0" width="100%" height="100%"><feTurbulence type="fractalNoise" baseFrequency="0.85" numOctaves="2" seed="${seed}" stitchTiles="stitch"/><feColorMatrix type="saturate" values="0"/></filter><rect width="${w}" height="${h}" filter="url(#g)" opacity="${opacity}"/>`
}

function draw(scene: Scene, w: number, h: number, seed: number) {
  const rand = random(seed * 7919 + scene.length)
  const parts: string[] = []
  const m = Math.min(w, h)

  switch (scene) {
    case "dunes": {
      parts.push(`<rect width="${w}" height="${h}" fill="white"/>`)
      parts.push(`<rect width="${w}" height="${h}" fill="black" opacity="0.06"/>`)
      const sunX = w * (0.25 + rand() * 0.5)
      parts.push(`<circle cx="${r1(sunX)}" cy="${r1(h * 0.3)}" r="${r1(m * 0.07)}" fill="white"/>`)
      const layers = [0.1, 0.2, 0.34, 0.52, 0.72, 0.86]
      layers.forEach((alpha, i) => {
        const base = h * (0.42 + i * 0.1)
        parts.push(
          `<path d="${ridge(w, h, base, h * 0.12, rand)}" fill="black" opacity="${alpha}"/>`,
        )
      })
      break
    }
    case "peaks": {
      parts.push(`<rect width="${w}" height="${h}" fill="white"/>`)
      parts.push(`<rect width="${w}" height="${h}" fill="black" opacity="0.04"/>`)
      const layers = [0.08, 0.15, 0.26, 0.42, 0.66]
      layers.forEach((alpha, i) => {
        const base = h * (0.36 + i * 0.12)
        parts.push(
          `<path d="${ridge(w, h, base, h * 0.28, rand, true)}" fill="black" opacity="${alpha}"/>`,
        )
        // Fog banks between the ranges.
        parts.push(
          `<rect y="${r1(base + h * 0.06)}" width="${w}" height="${r1(h * 0.05)}" fill="white" opacity="0.35" filter="url(#b)"/>`,
        )
      })
      parts.unshift(
        `<filter id="b" x="-10%" y="-200%" width="120%" height="500%"><feGaussianBlur stdDeviation="${r1(m * 0.03)}"/></filter>`,
      )
      break
    }
    case "sphere": {
      const cx = w * (0.4 + rand() * 0.2)
      const r = m * (0.18 + rand() * 0.06)
      const floor = h * 0.68
      const cy = floor - r
      parts.push(
        `<defs><radialGradient id="bg" cx="50%" cy="38%" r="75%"><stop offset="0" stop-color="white"/><stop offset="1" stop-color="black" stop-opacity="0.9"/></radialGradient>`,
        `<radialGradient id="s" cx="34%" cy="30%" r="80%"><stop offset="0" stop-color="white"/><stop offset="0.55" stop-color="black" stop-opacity="0.55"/><stop offset="1" stop-color="black" stop-opacity="0.95"/></radialGradient>`,
        `<filter id="b" x="-50%" y="-200%" width="200%" height="500%"><feGaussianBlur stdDeviation="${r1(r * 0.18)}"/></filter></defs>`,
      )
      parts.push(`<rect width="${w}" height="${h}" fill="black"/>`)
      parts.push(`<rect width="${w}" height="${h}" fill="url(#bg)" opacity="0.92"/>`)
      parts.push(
        `<rect y="${r1(floor)}" width="${w}" height="${r1(h - floor)}" fill="black" opacity="0.18"/>`,
      )
      parts.push(
        `<ellipse cx="${r1(cx + r * 0.25)}" cy="${r1(floor + r * 0.04)}" rx="${r1(r * 1.25)}" ry="${r1(r * 0.16)}" fill="black" opacity="0.7" filter="url(#b)"/>`,
      )
      parts.push(`<circle cx="${r1(cx)}" cy="${r1(cy)}" r="${r1(r)}" fill="url(#s)"/>`)
      break
    }
    case "ripples": {
      parts.push(`<rect width="${w}" height="${h}" fill="black"/>`)
      parts.push(`<rect width="${w}" height="${h}" fill="white" opacity="0.06"/>`)
      const cx = w * (0.35 + rand() * 0.3)
      const cy = h * (0.45 + rand() * 0.15)
      const rings = 9 + Math.floor(rand() * 4)
      for (let i = rings; i > 0; i--) {
        const rx = (i / rings) * w * 0.75
        parts.push(
          `<ellipse cx="${r1(cx)}" cy="${r1(cy)}" rx="${r1(rx)}" ry="${r1(rx * 0.32)}" fill="none" stroke="white" stroke-width="${r1(1 + (rings - i) * 0.35)}" opacity="${r1((1 - i / rings) * 0.75 + 0.18)}"/>`,
        )
      }
      parts.push(
        `<ellipse cx="${r1(cx)}" cy="${r1(cy)}" rx="${r1(m * 0.012)}" ry="${r1(m * 0.005)}" fill="white" opacity="0.9"/>`,
      )
      break
    }
    case "stairs": {
      parts.push(`<rect width="${w}" height="${h}" fill="white"/>`)
      parts.push(`<rect width="${w}" height="${h}" fill="black" opacity="0.08"/>`)
      const n = 8 + Math.floor(rand() * 4)
      const sw = w / n
      const sh = (h * 0.78) / n
      const base = h * 0.98
      // Risers: the solid mass under the zigzag.
      let pts = `0,${h}`
      for (let i = 0; i < n; i++) {
        pts += ` ${r1(i * sw)},${r1(base - i * sh)} ${r1(i * sw)},${r1(base - (i + 1) * sh)} ${r1((i + 1) * sw)},${r1(base - (i + 1) * sh)}`
      }
      pts += ` ${w},${h}`
      parts.push(`<polygon points="${pts}" fill="black" opacity="0.5"/>`)
      // Treads catch the sun as thin bright strips.
      for (let i = 0; i < n; i++) {
        parts.push(
          `<rect x="${r1(i * sw)}" y="${r1(base - (i + 1) * sh)}" width="${r1(sw)}" height="${r1(sh * 0.22)}" fill="white" opacity="0.7"/>`,
        )
      }
      // Hand rail and balusters, parallel to the flight.
      const lift = sh * 2.6
      parts.push(
        `<line x1="0" y1="${r1(base - lift)}" x2="${w}" y2="${r1(base - n * sh - lift)}" stroke="black" stroke-width="${r1(m * 0.006)}" opacity="0.75"/>`,
      )
      for (let i = 0; i < n; i++) {
        const x = i * sw + sw * 0.5
        const y = base - (i + 1) * sh
        parts.push(
          `<line x1="${r1(x)}" y1="${r1(y)}" x2="${r1(x)}" y2="${r1(y - lift + sh * 0.5)}" stroke="black" stroke-width="${r1(m * 0.003)}" opacity="0.6"/>`,
        )
      }
      // A hard shadow wedge from the wall above.
      const s0 = w * (0.45 + rand() * 0.15)
      parts.push(
        `<polygon points="0,0 ${r1(s0)},0 ${r1(s0 - w * 0.5)},${h} 0,${h}" fill="black" opacity="0.28"/>`,
      )
      break
    }
    case "moon": {
      parts.push(`<rect width="${w}" height="${h}" fill="black"/>`)
      parts.push(`<rect width="${w}" height="${h}" fill="white" opacity="0.05"/>`)
      const cx = w * (0.3 + rand() * 0.4)
      const horizon = h * 0.6
      const r = m * 0.11
      parts.push(
        `<filter id="b" x="-100%" y="-100%" width="300%" height="300%"><feGaussianBlur stdDeviation="${r1(r * 0.5)}"/></filter>`,
      )
      parts.push(
        `<circle cx="${r1(cx)}" cy="${r1(h * 0.3)}" r="${r1(r * 1.5)}" fill="white" opacity="0.12" filter="url(#b)"/>`,
      )
      parts.push(
        `<circle cx="${r1(cx)}" cy="${r1(h * 0.3)}" r="${r1(r)}" fill="white" opacity="0.92"/>`,
      )
      parts.push(
        `<rect y="${r1(horizon)}" width="${w}" height="${r1(h - horizon)}" fill="white" opacity="0.04"/>`,
      )
      // The reflection breaks into short horizontal strokes.
      for (let i = 0; i < 26; i++) {
        const y = horizon + 6 + i * ((h - horizon) / 26)
        const len = r * (0.4 + rand() * 1.4) * (1 + i * 0.05)
        parts.push(
          `<rect x="${r1(cx - len / 2 + (rand() - 0.5) * r * 0.6)}" y="${r1(y)}" width="${r1(len)}" height="${r1(1.5 + i * 0.12)}" fill="white" opacity="${r1(0.75 - i * 0.025)}"/>`,
        )
      }
      break
    }
  }

  parts.push(grain(w, h, seed % 97, scene === "ripples" || scene === "moon" ? 0.22 : 0.16))
  return parts.join("")
}

/** One picture of a scene as an SVG data URL, plus its intrinsic size. */
function scene(kind: Scene, { seed = 1, ratio = 1, size = 1200 }: SceneOptions = {}) {
  const w = ratio >= 1 ? size : Math.round(size * ratio)
  const h = ratio >= 1 ? Math.round(size / ratio) : size
  const svg = `<svg xmlns="http://www.w3.org/2000/svg" width="${w}" height="${h}" viewBox="0 0 ${w} ${h}">${draw(kind, w, h, seed)}</svg>`
  return { src: `data:image/svg+xml;charset=utf-8,${encodeURIComponent(svg)}`, width: w, height: h }
}

const images = [1, 2].map((seed) => ({
  id: `peaks-${seed}`,
  ...scene("peaks", { seed, ratio: 16 / 9 }),
}))

/** The model writes the prompt, the tool renders two variations, then the call settles. */
function ToolImageGenDemo() {
  const [typed, setTyped] = React.useState(0)
  const job = useGenerationJob({ queue: [], durationMs: 5200, autoStart: false })
  const { run } = job

  React.useEffect(() => {
    if (typed >= PROMPT.length) {
      run()
      return
    }
    const id = window.setTimeout(() => setTyped((n) => Math.min(PROMPT.length, n + 4)), 40)
    return () => window.clearTimeout(id)
  }, [typed, run])

  const state: ToolDisplayState =
    typed < PROMPT.length
      ? "input-streaming"
      : job.status === "done"
        ? "output-available"
        : "input-available"

  return (
    <div className="w-full max-w-xl">
      <ToolImageGen
        state={state}
        prompt={PROMPT.slice(0, typed)}
        revisedPrompt="Layered mountain ranges at sunrise with fog pooling in the valleys between them, shot on a long telephoto lens, muted black and white tones"
        aspectRatio="16:9"
        count={2}
        progress={state === "input-available" ? job.progress : undefined}
        images={state === "input-streaming" ? [] : images}
        onVary={() => {}}
      />
    </div>
  )
}

/** The line under each 21st demo that points back to the kit. */
function MoreOn({ slug, children }: { slug: string; children: React.ReactNode }) {
  return (
    <p data-fit-hide className="text-[0.75rem]/[1rem] text-(color:--chatcn-fg-faint) text-center text-balance">
      {children}{" "}
      <a
        href={`https://chatcn.dev/docs/components/${slug}`}
        target="_blank"
        rel="noopener"
        className="text-(color:--chatcn-fg-subtle) hover:text-(color:--chatcn-fg) underline underline-offset-2"
      >
        chatcn.dev
      </a>
    </p>
  )
}

/** The kit's own example, framed for the 21st preview. */
export default function ToolImageGenPreview() {
  return (
    <FitFrame minWidth="36rem">
      <div className="flex w-full flex-col items-center justify-center gap-4 p-6">
        <div className="flex h-[24rem] w-full max-w-2xl items-start justify-center">
          <ToolImageGenDemo />
        </div>
        <MoreOn slug="tool-image-gen">An image tool call, from prompt to picture. On</MoreOn>
      </div>
    </FitFrame>
  )
}