imagegen-tools/src/components/canvas/prompt-mention-preview.tsx

"use client";

import type { ReactElement, ReactNode } from "react";
import {
  parsePromptMentionUrl,
  type PromptMentionSource,
} from "@/lib/canvas/prompt-mentions";
import { cn } from "@/lib/utils";

export type PromptSegment =
  | { type: "text"; value: string }
  | { type: "mention"; source: PromptMentionSource };

type PromptMentionPreviewProps = {
  prompt: string;
  mentionSources: PromptMentionSource[];
  className?: string;
  textClassName?: string;
  chipClassName?: string;
  variant?: "link" | "image";
  placeholder?: ReactNode;
};

export function PromptMentionPreview({
  prompt,
  mentionSources,
  className,
  textClassName,
  chipClassName,
  variant = "link",
  placeholder,
}: PromptMentionPreviewProps): ReactElement | null {
  const segments = splitPromptSegments(prompt, mentionSources);

  if (!segments.length) {
    return placeholder ? <>{placeholder}</> : null;
  }

  return (
    <div className={cn("whitespace-pre-wrap break-words leading-6", className)}>
      {segments.map((segment, index) => {
        if (segment.type === "text") {
          return (
            <span
              key={`text-${index}`}
              className={cn("text-zinc-900", textClassName)}
            >
              {segment.value}
            </span>
          );
        }

        if (variant === "image") {
          return (
            <span
              key={`mention-${segment.source.nodeId}-${index}`}
              className={cn(
                "mx-0.5 inline-flex max-w-full align-middle items-center gap-1 rounded-full border border-zinc-200 bg-white px-1.5 py-0.5 text-xs font-medium text-zinc-800 shadow-sm",
                chipClassName,
              )}
            >
              <span className="flex size-4 shrink-0 items-center justify-center overflow-hidden rounded-full bg-zinc-100">
                {segment.source.imageUrl ? (
                  <img
                    alt={segment.source.description}
                    className="h-full w-full object-cover"
                    src={segment.source.imageUrl}
                  />
                ) : null}
              </span>
              <span className="min-w-0">
                <span className="block text-sky-700">{segment.source.label}</span>
              </span>
            </span>
          );
        }

        return (
          <span
            key={`mention-${segment.source.nodeId}-${index}`}
            className={cn(
              "my-0.5 inline-flex h-7 max-w-full align-middle items-center rounded-md border border-sky-200 bg-sky-50 px-2 text-xs font-medium text-sky-800",
              chipClassName,
            )}
          >
            <span className="truncate">{segment.source.label}</span>
          </span>
        );
      })}
    </div>
  );
}

export function splitPromptSegments(
  prompt: string,
  mentionSources: PromptMentionSource[],
): PromptSegment[] {
  if (!prompt) {
    return [];
  }

  const sortedSources = [...mentionSources].sort(
    (left, right) =>
      Math.max(right.token.length, right.alias.length) -
      Math.max(left.token.length, left.alias.length),
  );
  const segments: PromptSegment[] = [];

  let cursor = 0;
  while (cursor < prompt.length) {
    const markdownMatch = parseMarkdownMentionAt(prompt, cursor);
    const matchedSource = markdownMatch
      ? sortedSources.find((source) => source.nodeId === markdownMatch.nodeId)
      : sortedSources.find((source) => prompt.startsWith(source.alias, cursor));

    if (!matchedSource) {
      let nextCursor = cursor + 1;
      while (nextCursor < prompt.length) {
        const hasMentionAhead =
          Boolean(parseMarkdownMentionAt(prompt, nextCursor)) ||
          sortedSources.some((source) => prompt.startsWith(source.alias, nextCursor));
        if (hasMentionAhead) break;
        nextCursor += 1;
      }

      segments.push({
        type: "text",
        value: prompt.slice(cursor, nextCursor),
      });
      cursor = nextCursor;
      continue;
    }

    segments.push({
      type: "mention",
      source: matchedSource,
    });
    cursor += markdownMatch?.length ?? matchedSource.alias.length;
  }

  return segments;
}

function parseMarkdownMentionAt(prompt: string, cursor: number) {
  const match = /^\[([^\]]+)\]\(([^)]+)\)/.exec(prompt.slice(cursor));
  if (!match) return null;

  const nodeId = parsePromptMentionUrl(match[2]);
  if (!nodeId) return null;

  return {
    label: match[1],
    nodeId,
    length: match[0].length,
  };
}