Message

"use client";

import * as React from "react";
import { Slot } from "radix-ui";
import { Streamdown } from "streamdown";
import { code } from "@streamdown/code";

import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar";
import { MarkdownCodeBlock } from "@/components/pandacoderz-ui/markdown-code-block";
import {
  Tooltip,
  TooltipContent,
  TooltipProvider,
  TooltipTrigger,
} from "@/components/ui/tooltip";
import { Kbd } from "@/components/ui/kbd";
import { cn } from "@/lib/utils";

/**
 * Only the Shiki code plugin ships by default. Add `@streamdown/math`,
 * `@streamdown/mermaid`, or `@streamdown/cjk` here if you need them.
 */
const streamdownPlugins = { code } as const;

const messageMarkdownProseClasses = [
  "prose max-w-none text-foreground font-normal text-sm leading-6.5",
  // headings
  "prose-headings:font-[500] prose-headings:leading-5.5 prose-h2:tracking-[-0.45px] prose-headings:mb-4 prose-headings:mt-6 prose-h1:text-xl prose-h2:text-lg prose-h3:text-base prose-h3:leading-4.5 prose-h3:tracking-[-0.4px] prose-h4:text-sm prose-h5:text-xs prose-h6:text-xs",
  // heading links
  "prose-headings:[&_a]:no-underline prose-headings:[&_a]:shadow-none prose-headings:[&_a]:text-inherit",
  // body text
  "prose-p:mb-1 prose-p:mt-4",
  // links
  "[&_[data-streamdown=link]]:text-brand [&_[data-streamdown=link]]:font-normal [&_[data-streamdown=link]]:underline [&_[data-streamdown=link]]:underline-offset-2",
  // strong
  "[&_[data-streamdown=strong]]:text-foreground [&_[data-streamdown=strong]]:font-[550]",
  // lists
  "prose-li:my-[-0.5px] prose-li:marker:text-muted-foreground/50 prose-ul:my-0 prose-ol:my-0 prose-ol:pl-3",
] as const;

type MessageFrom = "user" | "assistant";

type MessageContextValue = {
  from: MessageFrom;
};

const MessageContext = React.createContext<MessageContextValue | null>(null);

function useMessageContext() {
  return React.useContext(MessageContext);
}

type MessageProps = React.HTMLAttributes<HTMLDivElement> & {
  from: MessageFrom;
};

const Message = React.forwardRef<HTMLDivElement, MessageProps>(function Message(
  {
    className,
    from,
    children,
    "aria-label": ariaLabelProp,
    "aria-labelledby": ariaLabelledBy,
    ...props
  },
  ref,
) {
  const ariaLabel =
    ariaLabelProp ??
    (ariaLabelledBy == null
      ? from === "user"
        ? "User message"
        : "Assistant message"
      : undefined);

  return (
    <MessageContext.Provider value={{ from }}>
      <div
        ref={ref}
        data-slot="message"
        data-from={from}
        role="article"
        aria-label={ariaLabel}
        aria-labelledby={ariaLabelledBy}
        className={cn(
          "group/message flex w-full max-w-[90%] items-start gap-2",
          from === "user" ? "ms-auto" : "me-auto",
          className,
        )}
        {...props}
      >
        {children}
      </div>
    </MessageContext.Provider>
  );
});

type MessageStackProps = React.HTMLAttributes<HTMLDivElement>;

function MessageStack({ className, ...props }: MessageStackProps) {
  const ctx = useMessageContext();
  const from = ctx?.from ?? "assistant";

  return (
    <div
      data-slot="message-stack"
      className={cn(
        "flex w-full flex-col gap-2",
        from === "user" ? "items-end" : "items-start",
        className,
      )}
      {...props}
    />
  );
}

type MessageContentProps = React.HTMLAttributes<HTMLDivElement>;

function MessageContent({ className, ...props }: MessageContentProps) {
  const ctx = useMessageContext();
  const from = ctx?.from ?? "assistant";

  return (
    <div
      data-slot="message-content"
      className={cn(
        "rounded-2xl text-sm leading-6.5 text-foreground",
        from === "user"
          ? "w-fit bg-brand-soft/70 px-4 py-2 dark:bg-brand-soft/60"
          : "mb-1 w-full bg-transparent px-2",
        className,
      )}
      {...props}
    />
  );
}

type MessageMarkdownProps = React.ComponentProps<typeof Streamdown>;

function MessageMarkdown({
  className,
  components,
  ...props
}: MessageMarkdownProps) {
  const mergedComponents = React.useMemo(() => {
    const defaultComponents = {
      code: MarkdownCodeBlock,
      inlineCode: ({
        children,
        className,
        ...props
      }: React.HTMLAttributes<HTMLElement>) => (
        <code
          className={cn(
            "rounded-md border-none bg-muted px-1.5 py-0.5 font-mono text-xs font-[450]",
            className,
          )}
          data-slot="message-markdown-inline-code"
          {...props}
        >
          {children}
        </code>
      ),
      table: (props: React.HTMLAttributes<HTMLTableElement>) => (
        <div
          data-slot="message-markdown-table-wrap"
          className={[
            "my-6 prose-no-margin overflow-hidden rounded-2xl border border-border bg-muted dark:border-accent dark:bg-background",
            "[&_tbody_tr:first-child_td:first-child]:rounded-ss-xl",
            "[&_tbody_tr:first-child_td:last-child]:rounded-se-xl",
            "[&_tbody_tr:last-child_td:first-child]:rounded-es-xl",
            "[&_tbody_tr:last-child_td:last-child]:rounded-ee-xl",
          ].join(" ")}
        >
          <table
            data-slot="message-markdown-table"
            className="w-full border-separate border-spacing-0 border-none bg-muted text-sm dark:bg-background"
            {...props}
          />
        </div>
      ),
      th: (props: React.ThHTMLAttributes<HTMLTableCellElement>) => (
        <th
          data-slot="message-markdown-th"
          className="border-none px-5 py-2 text-start text-[13px] font-normal! text-muted-foreground! dark:bg-background"
          {...props}
        />
      ),
      td: (props: React.TdHTMLAttributes<HTMLTableCellElement>) => (
        <td
          data-slot="message-markdown-td"
          className="border-0 border-accent bg-card px-5 py-3 text-[13px] text-foreground dark:bg-card [tr:not(:first-child)_&]:border-t"
          {...props}
        />
      ),
    };

    return {
      ...(defaultComponents as object),
      ...((components ?? {}) as object),
    };
  }, [components]);

  return (
    <Streamdown
      data-slot="message-markdown"
      className={cn(
        ...messageMarkdownProseClasses,
        "[&>*:first-child]:mt-0 [&>*:last-child]:mb-0",
        className,
      )}
      components={mergedComponents as MessageMarkdownProps["components"]}
      shikiTheme={["github-light-high-contrast", "github-dark-high-contrast"]}
      plugins={streamdownPlugins}
      {...props}
    />
  );
}

type MessageActionsProps = React.HTMLAttributes<HTMLDivElement>;

function MessageActions({ className, ...props }: MessageActionsProps) {
  const ctx = useMessageContext();
  const from = ctx?.from ?? "assistant";

  return (
    <div
      data-slot="message-actions"
      className={cn(
        "flex w-full",
        from === "user" ? "justify-end" : "justify-start",
        className,
      )}
      {...props}
    />
  );
}

type MessageActionGroupProps = React.HTMLAttributes<HTMLDivElement>;

function MessageActionGroup({ className, ...props }: MessageActionGroupProps) {
  return (
    <div
      data-slot="message-action-group"
      className={cn("flex items-center gap-1", className)}
      {...props}
    />
  );
}

type MessageActionProps = React.HTMLAttributes<HTMLDivElement> & {
  asChild?: boolean;
  tooltip?:
    | string
    | {
        content?: string;
        side?: "top" | "right" | "bottom" | "left";
        shortcut?: string;
      };
};

function MessageAction({
  asChild = false,
  tooltip,
  ...props
}: MessageActionProps) {
  const Comp = asChild ? Slot.Root : "div";
  const { content, side, shortcut } =
    typeof tooltip === "string" ? { content: tooltip } : (tooltip ?? {});

  if (!content) {
    return <Comp data-slot="message-action" {...props} />;
  }

  return (
    <TooltipProvider delayDuration={200}>
      <Tooltip>
        <TooltipTrigger asChild>
          <Comp data-slot="message-action" {...props} />
        </TooltipTrigger>
        <TooltipContent className="rounded-full" side={side}>
          {content}
          {shortcut ? <Kbd className="rounded-md!">{shortcut}</Kbd> : null}
        </TooltipContent>
      </Tooltip>
    </TooltipProvider>
  );
}

export type MessageAvatarProps = {
  src?: string;
  alt?: string;
  fallback?: React.ReactNode;
  delayMs?: React.ComponentProps<typeof AvatarFallback>["delayMs"];
  size?: React.ComponentProps<typeof Avatar>["size"];
  className?: string;
};

function MessageAvatar({
  src,
  alt = "",
  fallback,
  delayMs,
  size,
  className,
}: MessageAvatarProps) {
  return (
    <Avatar
      data-slot="message-avatar"
      size={size}
      className={cn("size-7 shrink-0", className)}
    >
      {src ? (
        <AvatarImage
          data-slot="message-avatar-image"
          src={src}
          alt={alt}
          className="my-0!"
        />
      ) : null}
      <AvatarFallback
        data-slot="message-avatar-fallback"
        delayMs={delayMs}
        className="my-0! shrink-0"
      >
        {fallback}
      </AvatarFallback>
    </Avatar>
  );
}

export {
  Message,
  MessageStack,
  MessageContent,
  MessageMarkdown,
  MessageActions,
  MessageActionGroup,
  MessageAction,
  MessageAvatar,
};

"use client";

/**
 * Streamdown `components.code` renderer for fenced code blocks.
 * Installed together with Message: same registry item as `message.tsx`.
 *
 * Chrome: optional title row (language label + copy), bordered card, and a
 * horizontally scrolling viewport. Highlighting is done client-side with
 * the `@streamdown/code` Shiki plugin and swaps in once tokens are ready.
 */

import { CheckIcon, DocumentDuplicateIcon } from "@heroicons-animated/react";
import { code as codeHighlighter } from "@streamdown/code";
import {
  type ComponentProps,
  type CSSProperties,
  type DetailedHTMLProps,
  type HTMLAttributes,
  type MouseEventHandler,
  type ReactNode,
  isValidElement,
  memo,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useState,
} from "react";
import type {
  BundledLanguage,
  CodeHighlighterPlugin,
  ExtraProps,
} from "streamdown";
import { StreamdownContext, useIsCodeFenceIncomplete } from "streamdown";

import { cn } from "@/lib/utils";

type MarkdownCodeElementProps = DetailedHTMLProps<
  HTMLAttributes<HTMLElement>,
  HTMLElement
> &
  ExtraProps;

export type MarkdownCodeBlockProps = MarkdownCodeElementProps & {
  /** Language label + copy in a header row. When false, copy floats top-right. */
  showTitleRow?: boolean;
};

type HighlightResult = NonNullable<
  ReturnType<CodeHighlighterPlugin["highlight"]>
>;

const LANGUAGE_REGEX = /language-([^\s]+)/;

function extractCodeString(children: ReactNode): string {
  if (
    isValidElement(children) &&
    children.props &&
    typeof children.props === "object" &&
    "children" in children.props &&
    typeof (children.props as { children?: unknown }).children === "string"
  ) {
    return (children.props as { children: string }).children;
  }
  if (typeof children === "string") return children;
  return "";
}

function trimTrailingNewlines(str: string): string {
  let end = str.length;
  while (end > 0 && str[end - 1] === "\n") end--;
  return str.slice(0, end);
}

function buildRawHighlightResult(trimmed: string): HighlightResult {
  return {
    bg: "transparent",
    fg: "inherit",
    tokens: trimmed.split("\n").map((line) => [
      {
        content: line,
        color: "inherit",
        bgColor: "transparent",
        htmlStyle: {},
        offset: 0,
      },
    ]),
  } as HighlightResult;
}

const COPIED_RESET_MS = 1500;

function useCopyButton(
  onCopy: () => void | Promise<void>,
): [checked: boolean, onClick: MouseEventHandler] {
  const [checked, setChecked] = useState(false);
  const callbackRef = useRef(onCopy);
  const timeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);

  useEffect(() => {
    callbackRef.current = onCopy;
  }, [onCopy]);

  const onClick = useCallback<MouseEventHandler>(() => {
    if (timeoutRef.current) clearTimeout(timeoutRef.current);
    void Promise.resolve(callbackRef.current()).then(() => {
      setChecked(true);
      timeoutRef.current = setTimeout(() => {
        setChecked(false);
      }, COPIED_RESET_MS);
    });
  }, []);

  useEffect(() => {
    return () => {
      if (timeoutRef.current) clearTimeout(timeoutRef.current);
    };
  }, []);

  return [checked, onClick];
}

function CopyButton({ text, className }: { text: string; className?: string }) {
  const [checked, onClick] = useCopyButton(() => {
    void navigator.clipboard.writeText(text);
  });

  return (
    <button
      type="button"
      data-checked={checked || undefined}
      className={cn(
        "relative flex size-7 cursor-pointer items-center justify-center rounded-lg text-muted-foreground hover:text-foreground",
        className,
      )}
      aria-label={checked ? "Copied" : "Copy code"}
      onClick={onClick}
    >
      {checked ? (
        <CheckIcon size={16} className="flex text-green-600 dark:text-green-400" />
      ) : (
        <DocumentDuplicateIcon size={16} className="flex" />
      )}
    </button>
  );
}

function TokenSpan({
  token,
}: {
  token: HighlightResult["tokens"][number][number];
}) {
  const tokenStyle: Record<string, string> = {};
  let hasBg = Boolean(token.bgColor);
  if (token.color) tokenStyle["--sdm-c"] = token.color;
  if (token.bgColor) tokenStyle["--sdm-tbg"] = token.bgColor;
  if (token.htmlStyle) {
    for (const [key, value] of Object.entries(token.htmlStyle)) {
      if (value == null) continue;
      if (key === "color") {
        tokenStyle["--sdm-c"] = String(value);
      } else if (key === "background-color") {
        tokenStyle["--sdm-tbg"] = String(value);
        hasBg = true;
      } else {
        tokenStyle[key] = String(value);
      }
    }
  }
  return (
    <span
      className={cn(
        "text-(--sdm-c,inherit)",
        "dark:text-(--shiki-dark,var(--sdm-c,inherit))",
        hasBg && "bg-(--sdm-tbg)",
      )}
      style={tokenStyle as CSSProperties}
    >
      {token.content}
    </span>
  );
}

const Pre = memo(
  function Pre({
    result,
    language,
    className,
    ...rest
  }: Omit<ComponentProps<"pre">, "children"> & {
    result: HighlightResult;
    language: string;
  }) {
    return (
      <pre
        className={cn("w-max min-w-full bg-transparent *:flex *:flex-col", className)}
        data-language={language}
        data-slot="markdown-code-block-body"
        {...rest}
      >
        <code>
          {result.tokens.map((row, rowIndex) => (
            <span key={rowIndex} className="block">
              {row.length === 0 || (row.length === 1 && row[0]?.content === "")
                ? "\n"
                : row.map((token, tokenIndex) => (
                    <TokenSpan key={tokenIndex} token={token} />
                  ))}
            </span>
          ))}
        </code>
      </pre>
    );
  },
  (prev, next) =>
    prev.result === next.result &&
    prev.language === next.language &&
    prev.className === next.className,
);
Pre.displayName = "MarkdownCodeBlockPre";

function ShikiPre({
  code,
  language,
  raw,
  codePlugin,
}: {
  code: string;
  language: string;
  raw: HighlightResult;
  codePlugin: CodeHighlighterPlugin;
}) {
  const { shikiTheme } = useContext(StreamdownContext);
  const [result, setResult] = useState<HighlightResult>(raw);

  useEffect(() => {
    let cancelled = false;
    const sync = codePlugin.highlight(
      {
        code,
        language: language as BundledLanguage,
        themes: shikiTheme,
      },
      (highlighted) => {
        if (!cancelled) setResult(highlighted);
      },
    );
    if (sync) setResult(sync);
    return () => {
      cancelled = true;
    };
  }, [code, language, shikiTheme, codePlugin, raw]);

  return <Pre language={language} result={result} />;
}

function FencedView({
  code,
  language,
  className,
  isIncomplete,
  showTitleRow = true,
}: {
  code: string;
  language: string;
  className?: string;
  isIncomplete?: boolean;
  showTitleRow?: boolean;
}) {
  const trimmed = useMemo(() => trimTrailingNewlines(code), [code]);
  const raw = useMemo(() => buildRawHighlightResult(trimmed), [trimmed]);
  const title = (language || "code").toLowerCase();
  const supported = codeHighlighter.supportsLanguage(
    language as Parameters<typeof codeHighlighter.supportsLanguage>[0],
  );

  return (
    <figure
      className={cn(
        "not-prose relative my-4 w-full overflow-hidden rounded-xl border border-border text-[13px] font-[450] dark:border-accent",
        showTitleRow ? "bg-muted dark:bg-background" : "bg-card",
        className,
      )}
      data-incomplete={isIncomplete || undefined}
      data-language={language}
      data-slot="markdown-code-block"
      dir="ltr"
    >
      {showTitleRow ? (
        <div className="flex h-9.5 items-center gap-2 px-4 text-muted-foreground">
          <figcaption className="flex-1 truncate text-[13px] lowercase">
            {title}
          </figcaption>
          <div className="-me-2 flex shrink-0 items-center">
            <CopyButton text={trimmed} />
          </div>
        </div>
      ) : (
        <div className="absolute top-2 right-2 z-20">
          <CopyButton text={trimmed} />
        </div>
      )}
      <div
        className={cn(
          "no-scrollbar overflow-auto overscroll-x-none bg-card px-4 py-3.5 font-mono text-sm leading-6",
          showTitleRow ? "rounded-t-xl" : "rounded-xl",
        )}
      >
        {supported ? (
          <ShikiPre
            code={trimmed}
            codePlugin={codeHighlighter}
            language={language}
            raw={raw}
          />
        ) : (
          <Pre language={language} result={raw} />
        )}
      </div>
    </figure>
  );
}

/**
 * Drop-in for Streamdown `components.code`. Inline code is handled by
 * `components.inlineCode` in Message, so this only renders fenced blocks.
 */
export function MarkdownCodeBlock({
  className,
  children,
  showTitleRow,
}: MarkdownCodeBlockProps) {
  const isIncomplete = useIsCodeFenceIncomplete();
  const match = className?.match(LANGUAGE_REGEX);
  const language = match?.[1] ?? "text";
  const code = extractCodeString(children);

  return (
    <FencedView
      code={code}
      isIncomplete={isIncomplete}
      language={language}
      showTitleRow={showTitleRow}
    />
  );
}