Tool

"use client";

import {
  createContext,
  useContext,
  type ComponentProps,
  type ComponentType,
  type CSSProperties,
} from "react";
import {
  ArrowPathIcon,
  CheckCircleIcon,
  ChevronDownIcon,
  ClockIcon,
  WrenchScrewdriverIcon,
  XCircleIcon,
} from "@heroicons-animated/react";

import { Badge } from "@/components/ui/badge";
import {
  Collapsible,
  CollapsibleContent,
  CollapsibleTrigger,
} from "@/components/ui/collapsible";
import { cn } from "@/lib/utils";

import {
  CodeBlock,
  CodeBlockContent,
  CodeBlockShiki,
} from "@/components/pandacoderz-ui/code-block";

type ToolStatus = "pending" | "ready" | "running" | "completed" | "error";

type ToolIcon = ComponentType<{ size?: number; className?: string }>;

type ToolMeta = {
  label: string;
  icon: ToolIcon;
  color: { bg: string; fg: string };
  iconClassName?: string;
};

const TOOL_META: Record<ToolStatus, ToolMeta> = {
  pending: {
    label: "Pending",
    icon: WrenchScrewdriverIcon,
    color: { bg: "var(--color-slate-100)", fg: "var(--color-slate-500)" },
  },
  ready: {
    label: "Ready",
    icon: ClockIcon,
    color: { bg: "var(--color-amber-100)", fg: "var(--color-amber-600)" },
  },
  running: {
    label: "Running",
    icon: ArrowPathIcon,
    color: { bg: "var(--color-violet-100)", fg: "var(--color-violet-600)" },
    iconClassName: "animate-spin",
  },
  completed: {
    label: "Completed",
    icon: CheckCircleIcon,
    color: { bg: "var(--color-emerald-100)", fg: "var(--color-emerald-600)" },
  },
  error: {
    label: "Error",
    icon: XCircleIcon,
    color: { bg: "var(--color-red-100)", fg: "var(--color-red-600)" },
  },
};

type ToolContextValue = {
  status: ToolStatus;
  meta: ToolMeta;
};

const ToolContext = createContext<ToolContextValue | null>(null);

function isToolStatus(value: unknown): value is ToolStatus {
  return (
    typeof value === "string" &&
    Object.prototype.hasOwnProperty.call(TOOL_META, value)
  );
}

function useToolContext(component: string): ToolContextValue {
  const context = useContext(ToolContext);
  if (!context) {
    throw new Error(`${component} must be used within <Tool>`);
  }
  return context;
}

function stringifyToolPayload(payload: unknown): string {
  if (typeof payload === "string") return payload;
  if (payload === undefined) return "";

  try {
    return JSON.stringify(payload, null, 2);
  } catch {
    return String(payload);
  }
}

type ToolProps = ComponentProps<typeof Collapsible> & {
  status: ToolStatus;
};

function Tool({ status, className, style, ...props }: ToolProps) {
  const resolvedStatus = isToolStatus(status) ? status : "pending";
  const meta = TOOL_META[resolvedStatus];

  return (
    <ToolContext.Provider value={{ status: resolvedStatus, meta }}>
      <Collapsible
        data-slot="tool"
        data-status={resolvedStatus}
        className={cn(
          "not-prose w-full max-w-100 rounded-xl border bg-card shadow-xs dark:border-accent",
          className,
        )}
        style={
          {
            "--tool-color": meta.color.fg,
            "--tool-bg": meta.color.bg,
            ...style,
          } as CSSProperties
        }
        {...props}
      />
    </ToolContext.Provider>
  );
}

type ToolTriggerProps = Omit<
  ComponentProps<typeof CollapsibleTrigger>,
  "children"
> & {
  name: string;
};

function ToolTrigger({ name, className, ...props }: ToolTriggerProps) {
  const { meta } = useToolContext("ToolTrigger");
  const Icon = meta.icon;

  return (
    <CollapsibleTrigger
      data-slot="tool-trigger"
      className={cn(
        "group flex h-10 w-full cursor-pointer items-center justify-between px-3 py-2",
        className,
      )}
      {...props}
    >
      <div className="flex items-center gap-2">
        <Icon
          data-slot="tool-trigger-icon"
          size={16}
          className={cn("flex text-(--tool-color)", meta.iconClassName)}
        />
        <span
          data-slot="tool-trigger-name"
          className="text-sm leading-6 font-[450] text-foreground"
        >
          {name}
        </span>
        <Badge
          data-slot="tool-trigger-badge"
          className="h-6 bg-(--tool-bg)/60 font-[450] text-(--tool-color) dark:bg-(--tool-color)/10 dark:text-(--tool-color)"
        >
          {meta.label}
        </Badge>
      </div>

      <ChevronDownIcon
        data-slot="tool-trigger-chevron"
        size={16}
        className="flex text-muted-foreground transition-transform duration-200 group-data-[state=open]:rotate-180"
      />
    </CollapsibleTrigger>
  );
}

type ToolContentProps = ComponentProps<typeof CollapsibleContent>;

function ToolContent({ className, ...props }: ToolContentProps) {
  return (
    <CollapsibleContent
      data-slot="tool-content"
      className={cn(
        "flex flex-col gap-6 p-3 pt-4",
        "overflow-hidden data-[state=closed]:animate-collapsible-up data-[state=open]:animate-collapsible-down",
        className,
      )}
      {...props}
    />
  );
}

type ToolPartProps = {
  kind: "input" | "output";
  payload: unknown;
  errorText?: string;
};

function ToolPart({ kind, payload, errorText }: ToolPartProps) {
  const { status } = useToolContext("ToolPart");
  const code = stringifyToolPayload(payload);
  const isOutputError = kind === "output" && status === "error";
  const hasPayload = payload !== undefined && payload !== null;
  const shouldShowCodeblock = !isOutputError || hasPayload;
  const title = kind === "input" ? "Input" : isOutputError ? "Error" : "Output";

  return (
    <div data-slot={`tool-${kind}`} className="flex flex-col gap-3">
      <span
        data-slot={`tool-${kind}-title`}
        className={cn(
          "text-xs leading-4 font-[450] text-muted-foreground uppercase",
          isOutputError && "text-destructive",
        )}
      >
        {title}
      </span>
      {isOutputError ? (
        <div
          data-slot="tool-output-error"
          className="rounded-xl border border-destructive/20 bg-destructive/5 px-4 py-3 text-sm leading-6 text-destructive dark:bg-destructive/10"
        >
          {errorText ?? "Tool execution failed"}
        </div>
      ) : null}
      {shouldShowCodeblock ? (
        <CodeBlock className="rounded-lg" keepBackground>
          <CodeBlockContent>
            <CodeBlockShiki language="json">{code}</CodeBlockShiki>
          </CodeBlockContent>
        </CodeBlock>
      ) : null}
    </div>
  );
}

type ToolPayloadProps = {
  payload: unknown;
};

function ToolInput({ payload }: ToolPayloadProps) {
  return <ToolPart kind="input" payload={payload} />;
}

type ToolOutputProps = ToolPayloadProps & {
  showWhen?: ToolStatus[];
  errorText?: string;
};

function ToolOutput({
  payload,
  showWhen = ["completed"],
  errorText,
}: ToolOutputProps) {
  const { status } = useToolContext("ToolOutput");
  if (!showWhen.includes(status)) return null;

  return <ToolPart kind="output" payload={payload} errorText={errorText} />;
}

export type { ToolStatus };
export { Tool, ToolTrigger, ToolContent, ToolInput, ToolOutput };

"use client";

import {
  CheckIcon,
  CodeBracketIcon,
  DocumentDuplicateIcon,
} from "@heroicons-animated/react";
import {
  createContext,
  useContext,
  useEffect,
  useState,
  type ComponentProps,
  type CSSProperties,
  type ReactNode,
} from "react";

import { cn } from "@/lib/utils";
import { highlight, Themes } from "@/lib/shiki/highlighter";
import type { BundledLanguage } from "shiki/bundle/web";

const highlighterPromise = highlight();

type DivProps = ComponentProps<"div">;
type CodeBlockProps = DivProps & {
  keepBackground?: boolean;
};
type CodeBlockCopyContextValue = {
  content: string;
  setContent: (value: string) => void;
};
const CodeBlockCopyContext = createContext<CodeBlockCopyContextValue | null>(
  null,
);

interface CodeBlockShikiProps extends DivProps {
  code?: string;
  language?: string;
  lineNumbers?: boolean;
  children?: ReactNode;
}

type ShikiToken = {
  content: string;
  htmlStyle?: Record<string, string>;
};

const buildRawTokenRows = (input: string): ShikiToken[][] =>
  input.split(/\r?\n/).map((line) => [{ content: line || " " }]);

const EMPTY_TOKEN_ROW: ShikiToken[] = [{ content: " " }];

const resolveCodeToHighlight = (children: ReactNode, code?: string): string =>
  typeof children === "string"
    ? children
    : Array.isArray(children) &&
        children.length === 1 &&
        typeof children[0] === "string"
      ? children[0]
      : (code ?? "");

const CodeBlock = ({
  children,
  className,
  keepBackground = false,
  ...props
}: CodeBlockProps) => {
  const [copyContent, setCopyContent] = useState("");

  return (
    <CodeBlockCopyContext.Provider
      value={{ content: copyContent, setContent: setCopyContent }}
    >
      <div
        data-slot="code-block"
        className={cn(
          "not-prose",
          "my-0 flex w-full flex-col overflow-hidden rounded-xl",
          keepBackground
            ? "border-none bg-secondary dark:bg-background"
            : "border bg-card dark:border-accent",
          "text-[13px] font-[450]",
          className,
        )}
        {...props}
      >
        {children}
      </div>
    </CodeBlockCopyContext.Provider>
  );
};

const CodeBlockHeader = ({ children, className, ...props }: DivProps) => {
  return (
    <div
      data-slot="code-block-header"
      className={cn(
        "not-prose flex h-9.5 items-center justify-between gap-2 px-4 text-muted-foreground",
        className,
      )}
      {...props}
    >
      {children}
    </div>
  );
};

interface CodeBlockIconProps extends DivProps {
  language?: string;
}

const CodeBlockIcon = ({ className }: CodeBlockIconProps) => {
  return <CodeBracketIcon size={16} className={cn("flex", className)} />;
};

const CodeBlockGroup = ({ children, className, ...props }: DivProps) => {
  return (
    <div
      data-slot="code-block-group"
      className={cn("flex items-center gap-2 text-muted-foreground", className)}
      {...props}
    >
      {children}
    </div>
  );
};

const CodeBlockContent = ({ className, children, ...props }: DivProps) => {
  return (
    <div
      data-slot="code-block-content"
      className={cn(
        "no-scrollbar max-h-96 overflow-auto overscroll-x-none",
        "rounded-xl px-4 text-sm leading-6",
        "font-mono whitespace-pre",
        className,
      )}
      {...props}
    >
      {children}
    </div>
  );
};

const CodeBlockShiki = ({
  code,
  language = "tsx",
  lineNumbers = false,
  className,
  children,
  ...props
}: CodeBlockShikiProps) => {
  const setCopyContent = useContext(CodeBlockCopyContext)?.setContent;
  const codeToHighlight = resolveCodeToHighlight(children, code);
  const [tokenRows, setTokenRows] = useState<ShikiToken[][]>(() =>
    buildRawTokenRows(codeToHighlight),
  );

  useEffect(() => {
    setCopyContent?.(codeToHighlight);
  }, [codeToHighlight, setCopyContent]);

  useEffect(() => {
    let cancelled = false;

    async function clientHighlight() {
      const rawRows = buildRawTokenRows(codeToHighlight);
      if (!cancelled) {
        setTokenRows(rawRows);
      }

      if (!codeToHighlight) {
        return;
      }

      try {
        const highlighter = await highlighterPromise;
        if (!highlighter.getLoadedLanguages().includes(language)) {
          await highlighter.loadLanguage(language as BundledLanguage);
        }
        const result = await highlighter.codeToTokens(codeToHighlight, {
          lang: language as BundledLanguage,
          themes: {
            light: Themes.light,
            dark: Themes.dark,
          },
        });
        if (!cancelled) {
          setTokenRows((result.tokens ?? rawRows) as ShikiToken[][]);
        }
      } catch {
        if (!cancelled) {
          setTokenRows(rawRows);
        }
      }
    }

    void clientHighlight();

    return () => {
      cancelled = true;
    };
  }, [codeToHighlight, language]);

  return (
    <div
      data-slot="code-block-shiki"
      className={cn(
        "no-scrollbar w-full overflow-auto overscroll-x-none py-0",
        className,
      )}
      {...props}
    >
      <pre
        className={cn(
          "shiki",
          lineNumbers ? "shiki-line-numbers" : "no-line-numbers",
        )}
      >
        <code>
          {tokenRows.map((row, rowIndex) => (
            <span key={`row-${rowIndex}`} className="line">
              {(row.length ? row : EMPTY_TOKEN_ROW).map((token, tokenIndex) => (
                <span
                  key={`token-${rowIndex}-${tokenIndex}`}
                  style={token.htmlStyle as CSSProperties | undefined}
                >
                  {token.content || " "}
                </span>
              ))}
              {rowIndex < tokenRows.length - 1 && "\n"}
            </span>
          ))}
        </code>
      </pre>
    </div>
  );
};

type CodeBlockCopyButtonProps = ComponentProps<"button">;

const CodeBlockCopyButton = ({
  className,
  ...props
}: CodeBlockCopyButtonProps) => {
  const content = useContext(CodeBlockCopyContext)?.content ?? "";
  const [isCopied, setIsCopied] = useState<boolean>(false);

  useEffect(() => {
    if (!isCopied) return;

    const timeout = setTimeout(() => {
      setIsCopied(false);
    }, 2000);
    return () => clearTimeout(timeout);
  }, [isCopied]);

  const handleCopy = async () => {
    if (!content) return;
    try {
      await navigator.clipboard.writeText(content);
      setIsCopied(true);
    } catch (err) {
      console.error("Failed to copy text: ", err);
    }
  };

  return (
    <button
      type="button"
      title="Copy to clipboard"
      aria-label={isCopied ? "Copied" : "Copy code"}
      data-slot="code-block-copy-button"
      className={cn(
        "relative flex size-7 cursor-pointer items-center justify-center rounded-full text-muted-foreground hover:text-foreground",
        className,
      )}
      onClick={handleCopy}
      {...props}
    >
      {isCopied ? (
        <CheckIcon
          size={14}
          className="flex animate-in text-green-600 duration-200 zoom-in-50 dark:text-green-400"
        />
      ) : (
        <DocumentDuplicateIcon
          size={14}
          className="flex animate-in duration-200 zoom-in-50"
        />
      )}
    </button>
  );
};

export {
  CodeBlock,
  CodeBlockHeader,
  CodeBlockIcon,
  CodeBlockGroup,
  CodeBlockContent,
  CodeBlockShiki,
  CodeBlockCopyButton,
};

import { createJavaScriptRegexEngine } from "shiki/engine/javascript";
import {
  createHighlighter,
  type Highlighter,
  type RegexEngine,
} from "shiki/bundle/web";

let jsEngine: RegexEngine | null = null;
let highlighter: Promise<Highlighter> | null = null;

/** High-contrast variants so highlighted code clears WCAG AA on both surfaces. */
const Themes = {
  light: "github-light-high-contrast",
  dark: "github-dark-high-contrast",
} as const;

/**
 * Languages preloaded for tool payloads and code blocks. Others are loaded
 * on demand via `loadLanguage`.
 */
const PRELOADED_LANGS = [
  "json",
  "tsx",
  "ts",
  "js",
  "jsx",
  "bash",
  "python",
  "html",
  "css",
  "md",
];

const getJsEngine = (): RegexEngine => {
  jsEngine ??= createJavaScriptRegexEngine();
  return jsEngine;
};

const highlight = async (): Promise<Highlighter> => {
  highlighter ??= createHighlighter({
    langs: PRELOADED_LANGS,
    themes: [Themes.light, Themes.dark],
    engine: getJsEngine(),
  });
  return highlighter;
};

export { highlight, Themes };