JSON Tree

"use client";

import * as React from "react";
import { Braces, ChevronRight, Search } from "lucide-react";
import { cn } from "@/lib/utils";
import { CopyButton } from "@/registry/paragon/ui/copy-button";
import {
  Tooltip,
  TooltipContent,
  TooltipProvider,
  TooltipTrigger,
} from "@/registry/paragon/ui/tooltip";

const IDENTIFIER_RE = /^[A-Za-z_$][\w$]*$/;

/** `user.plan`, `items[3]`, `headers["x-request-id"]`. */
function joinPath(parent: string, key: string | number, isIndex: boolean): string {
  if (isIndex) return `${parent}[${key}]`;
  if (IDENTIFIER_RE.test(String(key))) return parent ? `${parent}.${key}` : String(key);
  return `${parent}["${key}"]`;
}

function isContainer(value: unknown): value is object {
  return typeof value === "object" && value !== null;
}

function entriesOf(value: object): [string | number, unknown][] {
  return Array.isArray(value)
    ? value.map((v, i) => [i, v] as [number, unknown])
    : Object.entries(value);
}

const VALUE_CLASSES = {
  string: "text-[oklch(0.52_0.1_152)] dark:text-[oklch(0.78_0.1_152)]",
  number: "text-[oklch(0.5_0.13_252)] dark:text-[oklch(0.75_0.11_252)]",
  boolean: "text-[oklch(0.51_0.15_300)] dark:text-[oklch(0.77_0.12_300)]",
  nullish: "text-muted-foreground italic",
} as const;

function formatPrimitive(value: unknown): { text: string; className: string } {
  if (typeof value === "string")
    return { text: JSON.stringify(value), className: VALUE_CLASSES.string };
  if (typeof value === "number")
    return { text: String(value), className: VALUE_CLASSES.number };
  if (typeof value === "boolean")
    return { text: String(value), className: VALUE_CLASSES.boolean };
  return { text: String(value), className: VALUE_CLASSES.nullish };
}

/** Splits `text` on case-insensitive `query` hits and wraps them in <mark>. */
function highlightKey(text: string, query: string): React.ReactNode {
  if (!query) return text;
  const lower = text.toLowerCase();
  const parts: React.ReactNode[] = [];
  let i = 0;
  let at = lower.indexOf(query);
  while (at !== -1) {
    if (at > i) parts.push(text.slice(i, at));
    parts.push(
      <mark
        key={at}
        className="-mx-px rounded-[2px] bg-warning/30 px-px text-inherit"
      >
        {text.slice(at, at + query.length)}
      </mark>,
    );
    i = at + query.length;
    at = lower.indexOf(query, i);
  }
  if (i < text.length) parts.push(text.slice(i));
  return parts;
}

export interface JsonTreeProps
  extends Omit<React.ComponentProps<"div">, "children"> {
  /** Any JSON-serializable value. */
  data: unknown;
  /** Renders a root row with this name; paths then start with it. */
  rootLabel?: string;
  /** Containers at depth < this start expanded. Default 1 (top level). */
  defaultExpandedDepth?: number;
  /** Key search field that auto-expands and highlights matches. */
  searchable?: boolean;
  /** Copy-path button on row hover. */
  showCopyPath?: boolean;
  /** Vertical guide rails per level. */
  showGuides?: boolean;
  defaultQuery?: string;
  /** Disables expand/collapse and icon-swap motion. */
  static?: boolean;
}

/**
 * A collapsible JSON explorer. Values are type-tinted, branches open through
 * the sanctioned grid-rows transition, and the key search auto-expands every
 * ancestor of a match while highlighting the hit — clearing the query
 * collapses the tree back to how the user left it. Rows follow the ARIA tree
 * pattern with a roving tabindex, and hovering a row surfaces copy-path.
 */
export function JsonTree({
  data,
  rootLabel,
  defaultExpandedDepth = 1,
  searchable = true,
  showCopyPath = true,
  showGuides = true,
  defaultQuery = "",
  static: isStatic = false,
  className,
  ...props
}: JsonTreeProps) {
  const [expanded, setExpanded] = React.useState<Set<string>>(() => {
    const ids = new Set<string>();
    const walk = (value: unknown, path: string, depth: number) => {
      if (!isContainer(value)) return;
      if (depth < defaultExpandedDepth) ids.add(path);
      for (const [key, child] of entriesOf(value)) {
        walk(child, joinPath(path, key, Array.isArray(value)), depth + 1);
      }
    };
    if (rootLabel !== undefined) {
      walk(data, rootLabel, 0);
    } else if (isContainer(data)) {
      for (const [key, child] of entriesOf(data)) {
        walk(child, joinPath("", key, Array.isArray(data)), 0);
      }
    }
    return ids;
  });
  const [query, setQuery] = React.useState(defaultQuery);
  const [focused, setFocused] = React.useState<string | null>(null);
  const rowRefs = React.useRef(new Map<string, HTMLDivElement>());

  const search = query.trim().toLowerCase();

  // Matches (key contains query) and the ancestors that must open to show them.
  const { matches, ancestors } = React.useMemo(() => {
    const matchSet = new Set<string>();
    const ancestorSet = new Set<string>();
    if (search) {
      const walk = (value: unknown, path: string, chain: string[]) => {
        if (!isContainer(value)) return;
        for (const [key, child] of entriesOf(value)) {
          const childPath = joinPath(path, key, Array.isArray(value));
          if (String(key).toLowerCase().includes(search)) {
            matchSet.add(childPath);
            for (const ancestor of chain) ancestorSet.add(ancestor);
          }
          walk(child, childPath, [...chain, childPath]);
        }
      };
      if (rootLabel !== undefined) {
        walk(data, rootLabel, [rootLabel]);
      } else if (isContainer(data)) {
        for (const [key, child] of entriesOf(data)) {
          const childPath = joinPath("", key, Array.isArray(data));
          if (String(key).toLowerCase().includes(search)) matchSet.add(childPath);
          walk(child, childPath, [childPath]);
        }
      }
    }
    return { matches: matchSet, ancestors: ancestorSet };
  }, [data, search, rootLabel]);

  const isExpanded = React.useCallback(
    (path: string) => expanded.has(path) || (search !== "" && ancestors.has(path)),
    [expanded, search, ancestors],
  );

  // Visible rows in document order — the keyboard navigation space.
  const flat = React.useMemo(() => {
    const list: { path: string; parent: string | null; container: boolean; firstChild: string | null }[] = [];
    const walk = (value: unknown, path: string, parent: string | null) => {
      const container = isContainer(value) && entriesOf(value).length > 0;
      let firstChild: string | null = null;
      if (container && isExpanded(path)) {
        const [firstKey] = entriesOf(value as object)[0];
        firstChild = joinPath(path, firstKey, Array.isArray(value));
      }
      list.push({ path, parent, container, firstChild });
      if (container && isExpanded(path)) {
        for (const [key, child] of entriesOf(value as object)) {
          walk(child, joinPath(path, key, Array.isArray(value)), path);
        }
      }
    };
    if (rootLabel !== undefined) {
      walk(data, rootLabel, null);
    } else if (isContainer(data)) {
      for (const [key, child] of entriesOf(data)) {
        walk(child, joinPath("", key, Array.isArray(data)), null);
      }
    }
    return list;
  }, [data, rootLabel, isExpanded]);

  const focusTarget = focused ?? flat[0]?.path ?? null;

  const toggle = (path: string) => {
    setExpanded((prev) => {
      const next = new Set(prev);
      if (next.has(path)) next.delete(path);
      else next.add(path);
      return next;
    });
  };

  const focusPath = (path: string | null | undefined) => {
    if (!path) return;
    setFocused(path);
    rowRefs.current.get(path)?.focus();
  };

  const handleKeyDown = (event: React.KeyboardEvent, path: string) => {
    const index = flat.findIndex((f) => f.path === path);
    const row = flat[index];
    if (!row) return;
    switch (event.key) {
      case "ArrowDown":
        event.preventDefault();
        focusPath(flat[index + 1]?.path);
        break;
      case "ArrowUp":
        event.preventDefault();
        focusPath(flat[index - 1]?.path);
        break;
      case "ArrowRight":
        event.preventDefault();
        if (!row.container) break;
        if (!isExpanded(path)) toggle(path);
        else focusPath(row.firstChild);
        break;
      case "ArrowLeft":
        event.preventDefault();
        if (row.container && isExpanded(path)) toggle(path);
        else focusPath(row.parent);
        break;
      case "Home":
        event.preventDefault();
        focusPath(flat[0]?.path);
        break;
      case "End":
        event.preventDefault();
        focusPath(flat[flat.length - 1]?.path);
        break;
      case "Enter":
      case " ":
        event.preventDefault();
        if (row.container) toggle(path);
        break;
    }
  };

  const renderRow = (
    key: string | number | null,
    value: unknown,
    path: string,
    depth: number,
  ): React.ReactNode => {
    const container = isContainer(value);
    const count = container ? entriesOf(value).length : 0;
    const expandable = container && count > 0;
    const open = expandable && isExpanded(path);
    const matched = matches.has(path);
    const label = key === null ? (rootLabel ?? "root") : String(key);

    return (
      <li key={path} role="none">
        <div
          role="treeitem"
          aria-level={depth + 1}
          aria-expanded={expandable ? open : undefined}
          aria-selected={focusTarget === path}
          tabIndex={focusTarget === path ? 0 : -1}
          ref={(el) => {
            if (el) rowRefs.current.set(path, el);
            else rowRefs.current.delete(path);
          }}
          onClick={() => {
            setFocused(path);
            if (expandable) toggle(path);
          }}
          onKeyDown={(event) => handleKeyDown(event, path)}
          onFocus={() => setFocused(path)}
          className={cn(
            "group/row flex h-7 cursor-default items-center gap-1 rounded-md pr-1 pl-1.5 font-mono text-[12.5px] select-none",
            "transition-colors duration-(--duration-fast) hover:bg-muted/40",
            matched && "bg-warning/[0.07]",
          )}
        >
          {expandable ? (
            <ChevronRight
              aria-hidden
              className={cn(
                "size-3.5 shrink-0 text-muted-foreground/70",
                open && "rotate-90",
                !isStatic &&
                  "transition-[rotate] duration-(--duration-quick) ease-(--ease-out) motion-reduce:transition-none",
              )}
            />
          ) : (
            <span aria-hidden className="w-3.5 shrink-0" />
          )}
          <span className="shrink-0 font-medium text-foreground">
            {highlightKey(label, search)}
          </span>
          <span aria-hidden className="shrink-0 text-muted-foreground/60">
            :
          </span>
          {container ? (
            <span className="truncate text-[11px] text-muted-foreground">
              {Array.isArray(value)
                ? count === 0
                  ? "[]"
                  : `[${count}]`
                : count === 0
                  ? "{}"
                  : `{${count}}`}
              {expandable && !open && <span aria-hidden> …</span>}
            </span>
          ) : (
            (() => {
              const formatted = formatPrimitive(value);
              return (
                <span
                  className={cn("min-w-0 truncate", formatted.className)}
                  title={formatted.text.length > 32 ? formatted.text : undefined}
                >
                  {formatted.text}
                </span>
              );
            })()
          )}
          {showCopyPath && (
            <Tooltip>
              <TooltipTrigger asChild>
                <CopyButton
                  value={path}
                  static={isStatic}
                  aria-label={`Copy path ${path}`}
                  onClick={(event) => event.stopPropagation()}
                  tabIndex={-1}
                  className={cn(
                    "ml-auto size-5 rounded-[5px] opacity-0 transition-opacity duration-(--duration-fast)",
                    "after:size-6 group-hover/row:opacity-100 group-focus-within/row:opacity-100 [&_svg]:size-3",
                    focusTarget === path && "group-focus-within/row:opacity-100",
                  )}
                />
              </TooltipTrigger>
              <TooltipContent side="left">Copy path</TooltipContent>
            </Tooltip>
          )}
        </div>

        {expandable && (
          <div
            className={cn(
              "grid",
              open ? "grid-rows-[1fr]" : "grid-rows-[0fr]",
              !isStatic &&
                "transition-[grid-template-rows] motion-reduce:transition-none",
              !isStatic &&
                (open
                  ? "duration-(--duration-base) ease-(--ease-out)"
                  : "duration-(--duration-quick) ease-(--ease-exit)"),
            )}
          >
            <div inert={!open} className="min-h-0 overflow-hidden">
              <ul
                role="group"
                className={cn(
                  "ml-[13px] flex flex-col pl-1.5",
                  showGuides && "border-l border-border",
                )}
              >
                {entriesOf(value).map(([childKey, child]) =>
                  renderRow(
                    childKey,
                    child,
                    joinPath(path, childKey, Array.isArray(value)),
                    depth + 1,
                  ),
                )}
              </ul>
            </div>
          </div>
        )}
      </li>
    );
  };

  const topLevel: React.ReactNode =
    rootLabel !== undefined
      ? renderRow(null, data, rootLabel, 0)
      : isContainer(data)
        ? entriesOf(data).map(([key, child]) =>
            renderRow(key, child, joinPath("", key, Array.isArray(data)), 0),
          )
        : renderRow(null, data, "value", 0);

  const empty = isContainer(data) && entriesOf(data).length === 0 && rootLabel === undefined;

  return (
    <TooltipProvider>
      <div
        data-slot="json-tree"
        className={cn("w-full text-sm", className)}
        {...props}
      >
        {searchable && (
          <div className="relative mb-2">
            <Search
              aria-hidden
              className="pointer-events-none absolute top-1/2 left-2.5 size-3.5 -translate-y-1/2 text-muted-foreground"
            />
            <input
              type="text"
              value={query}
              onChange={(event) => setQuery(event.target.value)}
              onKeyDown={(event) => {
                if (event.key === "Escape" && query) {
                  event.stopPropagation();
                  setQuery("");
                }
              }}
              placeholder="Search keys…"
              aria-label="Search keys"
              className={cn(
                "h-8 w-full rounded-lg border border-input bg-transparent pr-20 pl-8 text-[13px] text-foreground",
                "transition-[border-color,box-shadow] duration-150 ease-out",
                "placeholder:text-muted-foreground",
                "outline-none focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/25",
              )}
            />
            {search && (
              <span
                role="status"
                className="pointer-events-none absolute top-1/2 right-2.5 -translate-y-1/2 text-[11px] text-muted-foreground tabular-nums"
              >
                {matches.size} {matches.size === 1 ? "match" : "matches"}
              </span>
            )}
          </div>
        )}

        {empty ? (
          <div className="flex flex-col items-center gap-1 rounded-lg border border-dashed px-4 py-8 text-center">
            <Braces aria-hidden className="size-4 text-muted-foreground/60" />
            <p className="text-xs text-muted-foreground">No data</p>
          </div>
        ) : (
          <ul role="tree" aria-label={rootLabel ?? "JSON"} className="flex flex-col">
            {topLevel}
          </ul>
        )}
      </div>
    </TooltipProvider>
  );
}