Assistant Message

"use client"

import * as React from "react"
import { RotateCcwIcon, ThumbsDownIcon, ThumbsUpIcon } from "lucide-react"
import { AnimatePresence, motion } from "motion/react"

import { cn } from "@/lib/utils"
import { useControllableState } from "@/hooks/use-controllable-state"
import { duration, ease, useMotionSafe } from "@/lib/motion"
import { CopyButton } from "@/components/ui/copy-button"

export type MessageFeedback = "up" | "down" | null

export type AssistantMessageProps = Omit<React.ComponentProps<"article">, "children"> & {
  /**
   * The reply as plain text. Blank lines separate paragraphs, lines starting with "- " are
   * bullets, and \`backticks\` mark code.
   */
  content: string
  /** True while more text is arriving. New words fade in; the actions wait until it is done. */
  streaming?: boolean
  /** The assistant's picture. Pass false to leave it out. */
  avatar?: React.ReactNode | false
  /** Who is speaking, for screen readers. */
  name?: string
  /** Shows a retry button. */
  onRetry?: () => void
  feedback?: MessageFeedback
  defaultFeedback?: MessageFeedback
  /** Shows thumbs up and down. Pressing the chosen one again clears it. */
  onFeedback?: (feedback: MessageFeedback) => void
  /** Hides copy, retry and feedback. */
  hideActions?: boolean
}

type Block = { kind: "p" | "li"; text: string; start: number }

function toBlocks(content: string): Block[][] {
  // Paragraphs, each a run of plain lines or of bullets, with every word numbered in order.
  let start = 0
  return content
    .split(/\n{2,}/)
    .map((para) =>
      para
        .split("\n")
        .filter((line) => line.trim())
        .map((line) => {
          const bullet = line.trimStart().startsWith("- ")
          const text = bullet ? line.trimStart().slice(2) : line
          const block: Block = { kind: bullet ? "li" : "p", text, start }
          start += wordCount(text)
          return block
        })
    )
    .filter((group) => group.length > 0)
}

/** Splits a line into words, keeping `code` spans whole, and numbers every word across the message. */
function Words({ text, start }: { text: string; start: number }) {
  const parts = text.split(/(`[^`]*`?)/).filter(Boolean)
  let index = start
  return (
    <>
      {parts.map((part, p) => {
        if (part.startsWith("`")) {
          const i = index++
          return (
            <code key={p} data-word={i} className="rounded-md bg-muted px-1.5 py-0.5 font-mono text-[0.88em]">
              {part.replace(/`/g, "")}
            </code>
          )
        }
        return part.split(/(\s+)/).map((word, w) =>
          /^\s+$/.test(word) || !word ? (
            word
          ) : (
            <span key={`${p}-${w}`} data-word={index++}>
              {word}
            </span>
          )
        )
      })}
    </>
  )
}

const wordCount = (text: string) =>
  text
    .split(/(`[^`]*`?)/)
    .filter(Boolean)
    .reduce((n, part) => n + (part.startsWith("`") ? 1 : part.split(/\s+/).filter(Boolean).length), 0)

/**
 * One reply from an assistant. While it streams, each new word fades in, so the text fills in
 * from top to bottom. When it is done, copy, retry and feedback appear beneath it.
 */
function AssistantMessage({
  content,
  streaming = false,
  avatar,
  name = "Assistant",
  onRetry,
  feedback: feedbackProp,
  defaultFeedback = null,
  onFeedback,
  hideActions = false,
  className,
  ...props
}: AssistantMessageProps) {
  const [feedback, setFeedback] = useControllableState<MessageFeedback>({
    value: feedbackProp,
    defaultValue: defaultFeedback,
    onChange: onFeedback,
  })
  const body = React.useRef<HTMLDivElement>(null)
  const shown = React.useRef(0)
  const { reduced, safe } = useMotionSafe()
  const blocks = React.useMemo(() => toBlocks(content), [content])

  // Fade in only the words that are new since the last render.
  React.useLayoutEffect(() => {
    const node = body.current
    if (!node) return
    const words = node.querySelectorAll<HTMLElement>("[data-word]")
    if (!reduced && streaming) {
      let order = 0
      for (const word of words) {
        if (Number(word.dataset.word) < shown.current) continue
        word.animate(
          [
            { opacity: 0, filter: "blur(3px)" },
            { opacity: 1, filter: "blur(0px)" },
          ],
          { duration: 420, delay: order++ * 18, easing: "cubic-bezier(0.22, 1, 0.36, 1)", fill: "backwards" }
        )
      }
    }
    shown.current = words.length
  }, [content, streaming, reduced])

  const mark = avatar === undefined ? <DefaultAvatar /> : avatar

  return (
    <article
      data-slot="assistant-message"
      data-streaming={streaming || undefined}
      aria-busy={streaming || undefined}
      className={cn("flex w-full gap-3", className)}
      {...props}
    >
      {mark !== false && (
        <div aria-hidden="true" className="mt-0.5 size-7 shrink-0 overflow-hidden rounded-full">
          {mark}
        </div>
      )}
      <div className="flex min-w-0 flex-1 flex-col gap-2">
        <h2 className="sr-only">{name}</h2>
        <div ref={body} className="flex flex-col gap-3 text-[15px] leading-7 text-pretty">
          {blocks.map((group, g) => {
            const bullets = group.every((b) => b.kind === "li")
            if (bullets)
              return (
                <ul key={g} className="flex list-disc flex-col gap-1 pl-5 marker:text-muted-foreground">
                  {group.map((b, i) => (
                    <li key={i}>
                      <Words text={b.text} start={b.start} />
                    </li>
                  ))}
                </ul>
              )
            return (
              <p key={g}>
                {group.map((b, i) => (
                  <React.Fragment key={i}>
                    {i > 0 && <br />}
                    <Words text={b.text} start={b.start} />
                  </React.Fragment>
                ))}
              </p>
            )
          })}
        </div>

        <AnimatePresence initial={false}>
          {!streaming && !hideActions && content && (
            <motion.div
              initial={{ opacity: 0, y: reduced ? 0 : 4 }}
              animate={{ opacity: 1, y: 0 }}
              exit={{ opacity: 0 }}
              transition={safe({ duration: duration.base, ease: ease.out })}
              className="-ml-2 flex items-center gap-0.5 text-muted-foreground"
            >
              <CopyButton value={content} label="Copy reply" size="icon-sm" variant="ghost" className="text-muted-foreground hover:text-foreground" />
              {onRetry && (
                <ActionButton label="Retry" onClick={onRetry}>
                  <RotateCcwIcon />
                </ActionButton>
              )}
              {onFeedback !== undefined || feedbackProp !== undefined ? (
                <>
                  <ActionButton
                    label="Good reply"
                    pressed={feedback === "up"}
                    onClick={() => setFeedback(feedback === "up" ? null : "up")}
                  >
                    <ThumbsUpIcon className={cn(feedback === "up" && "fill-current")} />
                  </ActionButton>
                  <ActionButton
                    label="Bad reply"
                    pressed={feedback === "down"}
                    onClick={() => setFeedback(feedback === "down" ? null : "down")}
                  >
                    <ThumbsDownIcon className={cn(feedback === "down" && "fill-current")} />
                  </ActionButton>
                </>
              ) : null}
            </motion.div>
          )}
        </AnimatePresence>
      </div>
    </article>
  )
}

function ActionButton({
  label,
  pressed,
  onClick,
  children,
}: {
  label: string
  pressed?: boolean
  onClick: () => void
  children: React.ReactNode
}) {
  return (
    <button
      type="button"
      aria-label={label}
      aria-pressed={pressed}
      onClick={onClick}
      className={cn(
        "grid size-8 cursor-pointer place-items-center rounded-md outline-none transition-colors hover:bg-accent hover:text-foreground focus-visible:ring-2 focus-visible:ring-ring [&_svg]:size-4",
        pressed && "text-foreground"
      )}
    >
      {children}
    </button>
  )
}

function DefaultAvatar() {
  return (
    <span className="grid size-full place-items-center bg-foreground text-background">
      <svg viewBox="0 0 16 16" className="size-3.5" fill="currentColor">
        <path d="M8 0c.4 3.9 2.1 5.6 6 6-3.9.4-5.6 2.1-6 6-.4-3.9-2.1-5.6-6-6 3.9-.4 5.6-2.1 6-6Z" />
        <path d="M13.5 10c.15 1.6.9 2.35 2.5 2.5-1.6.15-2.35.9-2.5 2.5-.15-1.6-.9-2.35-2.5-2.5 1.6-.15 2.35-.9 2.5-2.5Z" opacity=".7" />
      </svg>
    </span>
  )
}

export { AssistantMessage }