JSON Viewer

'use client';

import * as React from 'react';
import { ChevronRight, ChevronsDownUp, ChevronsUpDown } from 'lucide-react';

import { cn } from '@/lib/utils';
import { CopyButton } from '@/registry/hirael/bases/radix/components/copy-button';
import { Button } from '@/registry/hirael/bases/radix/ui/button';

type JsonKind =
  | 'object'
  | 'array'
  | 'string'
  | 'number'
  | 'bigint'
  | 'boolean'
  | 'null'
  | 'undefined'
  | 'date'
  | 'function'
  | 'symbol';

const kindOf = (value: unknown): JsonKind => {
  if (value === null) return 'null';
  if (Array.isArray(value)) return 'array';
  if (value instanceof Date) return 'date';
  const type = typeof value;
  if (type === 'object') return 'object';
  if (type === 'string') return 'string';
  if (type === 'number') return 'number';
  if (type === 'bigint') return 'bigint';
  if (type === 'boolean') return 'boolean';
  if (type === 'function') return 'function';
  if (type === 'symbol') return 'symbol';

  return 'undefined';
};

const isExpandable = (kind: JsonKind) => {
  return kind === 'object' || kind === 'array';
};

const entriesOf = (value: unknown, kind: JsonKind): [string, unknown][] => {
  if (kind === 'array') {
    return (value as unknown[]).map((v, i) => [String(i), v]);
  }
  if (kind === 'object') {
    return Object.entries(value as Record<string, unknown>);
  }

  return [];
};

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

// Keys that aren't identifiers get bracket-quoted so "a.b" and nested a -> b don't share a path.
const childPath = (path: string, key: string, parentKind: JsonKind) => {
  if (parentKind === 'array') return `${path}[${key}]`;

  return IDENTIFIER.test(key) ? `${path}.${key}` : `${path}[${JSON.stringify(key)}]`;
};

const CIRCULAR = '[Circular]';

const collectPaths = (value: unknown, path = '$', ancestors = new WeakSet<object>()): string[] => {
  const kind = kindOf(value);
  if (!isExpandable(kind)) return [];
  const node = value as object;
  if (ancestors.has(node)) return [];
  ancestors.add(node);
  const out = [path];
  for (const [key, child] of entriesOf(value, kind)) {
    out.push(...collectPaths(child, childPath(path, key, kind), ancestors));
  }
  ancestors.delete(node);

  return out;
};

const stringify = (value: unknown, indent: number) => {
  const ancestors: unknown[] = [];

  return JSON.stringify(
    value,
    function replacer(this: unknown, _key: string, v: unknown) {
      if (typeof v === 'bigint') return v.toString();
      if (typeof v !== 'object' || v === null) return v;
      while (ancestors.length > 0 && ancestors[ancestors.length - 1] !== this) ancestors.pop();
      if (ancestors.includes(v)) return CIRCULAR;
      ancestors.push(v);

      return v;
    },
    indent,
  );
};

type FocusUpdate = string | ((current: string) => string);

interface FocusStore {
  get: () => string;
  set: (next: FocusUpdate) => void;
  subscribe: (listener: () => void) => () => void;
}

const createFocusStore = (initial: string): FocusStore => {
  let current = initial;
  const listeners = new Set<() => void>();

  return {
    get: () => current,
    set: (next) => {
      const resolved = typeof next === 'function' ? next(current) : next;
      if (resolved === current) return;
      current = resolved;
      listeners.forEach((listener) => listener());
    },
    subscribe: (listener) => {
      listeners.add(listener);

      return () => listeners.delete(listener);
    },
  };
};

interface JsonViewerCtx {
  value: unknown;
  maxStringLength: number;
  isExpanded: (path: string, depth: number) => boolean;
  toggle: (path: string, depth: number) => void;
  expandAll: () => void;
  collapseAll: () => void;
  setFocusedPath: (next: FocusUpdate) => void;
  treeRef: React.RefObject<HTMLDivElement | null>;
}

const JsonViewerContext = React.createContext<JsonViewerCtx | null>(null);

const JsonViewerFocusContext = React.createContext<FocusStore | null>(null);

const useIsFocusedPath = (path: string) => {
  const store = React.useContext(JsonViewerFocusContext);
  const subscribe = React.useCallback((listener: () => void) => store?.subscribe(listener) ?? (() => {}), [store]);

  return React.useSyncExternalStore(
    subscribe,
    () => (store ? store.get() === path : path === '$'),
    () => path === '$',
  );
};

interface Ancestry {
  value: object;
  parent: Ancestry | null;
}

const JsonViewerAncestryContext = React.createContext<Ancestry | null>(null);

const isCircular = (ancestry: Ancestry | null, value: unknown) => {
  for (let link = ancestry; link; link = link.parent) {
    if (link.value === value) return true;
  }

  return false;
};

const useJsonViewer = () => {
  const ctx = React.useContext(JsonViewerContext);
  if (!ctx) {
    throw new Error('JsonViewer compound parts must be used inside <JsonViewer>');
  }

  return ctx;
};

interface JsonViewerNodeCtx {
  expandable: boolean;
  expanded: boolean;
}

const JsonViewerNodeContext = React.createContext<JsonViewerNodeCtx | null>(null);

const useJsonViewerNode = () => {
  const ctx = React.useContext(JsonViewerNodeContext);
  if (!ctx) {
    throw new Error('JsonViewer node parts must be used inside <JsonViewerNode>');
  }

  return ctx;
};

export interface JsonViewerProps extends Omit<React.ComponentProps<'div'>, 'children'> {
  value: unknown;
  /**
   * `true` expands everything, `false` collapses the root, a number expands
   * nodes shallower than that depth (`1` opens only the root).
   */
  defaultExpanded?: number | boolean;
  /** Controlled set of expanded paths, e.g. `["$", "$.user", "$.items[0]"]`. */
  expanded?: string[];
  onExpandedChange?: (paths: string[]) => void;
  /** Strings longer than this are truncated with an inline expand control. */
  maxStringLength?: number;
  children?: React.ReactNode;
}

interface UncontrolledState {
  depth: number;
  overrides: Record<string, boolean>;
}

const JsonViewer = ({
  value,
  defaultExpanded = 2,
  expanded: expandedProp,
  onExpandedChange,
  maxStringLength = 80,
  className,
  children,
  ...props
}: JsonViewerProps) => {
  const [internal, setInternal] = React.useState<UncontrolledState>(() => ({
    depth: defaultExpanded === true ? Infinity : defaultExpanded === false ? 0 : defaultExpanded,
    overrides: {},
  }));
  const [focusStore] = React.useState(() => createFocusStore('$'));
  const setFocusedPath = focusStore.set;
  const treeRef = React.useRef<HTMLDivElement | null>(null);

  const controlled = expandedProp !== undefined;
  const expandedSet = React.useMemo(() => (expandedProp ? new Set(expandedProp) : null), [expandedProp]);

  const isExpanded = React.useCallback(
    (path: string, depth: number) => {
      if (expandedSet) return expandedSet.has(path);
      const override = internal.overrides[path];

      return override ?? depth < internal.depth;
    },
    [expandedSet, internal],
  );

  const toggle = React.useCallback(
    (path: string, depth: number) => {
      const next = !isExpanded(path, depth);
      // Keep the roving tab stop on something visible when a parent closes.
      if (!next) {
        setFocusedPath((current) =>
          current !== path && (current.startsWith(`${path}.`) || current.startsWith(`${path}[`)) ? path : current,
        );
      }
      if (controlled) {
        const set = new Set(expandedSet);
        if (next) set.add(path);
        else set.delete(path);
        onExpandedChange?.(Array.from(set));

        return;
      }
      setInternal((s) => ({
        ...s,
        overrides: { ...s.overrides, [path]: next },
      }));
    },
    [controlled, expandedSet, isExpanded, onExpandedChange, setFocusedPath],
  );

  const expandAll = React.useCallback(() => {
    if (controlled) {
      onExpandedChange?.(collectPaths(value));

      return;
    }
    setInternal({ depth: Infinity, overrides: {} });
  }, [controlled, onExpandedChange, value]);

  const collapseAll = React.useCallback(() => {
    setFocusedPath('$');
    if (controlled) {
      onExpandedChange?.([]);

      return;
    }
    setInternal({ depth: 0, overrides: {} });
  }, [controlled, onExpandedChange, setFocusedPath]);

  const ctx = React.useMemo<JsonViewerCtx>(
    () => ({
      value,
      maxStringLength,
      isExpanded,
      toggle,
      expandAll,
      collapseAll,
      setFocusedPath,
      treeRef,
    }),
    [value, maxStringLength, isExpanded, toggle, expandAll, collapseAll, setFocusedPath],
  );

  return (
    <JsonViewerContext.Provider value={ctx}>
      <JsonViewerFocusContext.Provider value={focusStore}>
        <div
          data-slot="json-viewer"
          className={cn(
            'overflow-auto rounded-md border border-border bg-card p-3 font-mono text-xs leading-relaxed text-card-foreground',
            className,
          )}
          {...props}
        >
          {children ?? <JsonViewerTree />}
        </div>
      </JsonViewerFocusContext.Provider>
    </JsonViewerContext.Provider>
  );
};

const JsonViewerTree = ({ className, ...props }: Omit<React.ComponentProps<'div'>, 'children'>) => {
  const { value, treeRef } = useJsonViewer();

  return (
    <div ref={treeRef} role="tree" data-slot="json-viewer-tree" className={cn('min-w-max', className)} {...props}>
      <JsonViewerNode value={value} path="$" depth={0} isLast />
    </div>
  );
};

export interface JsonViewerNodeProps extends Omit<React.ComponentProps<'div'>, 'children'> {
  value: unknown;
  /** Key or array index. Omit for the root. */
  name?: string;
  path?: string;
  depth?: number;
  isLast?: boolean;
}

const JsonViewerNodeImpl = ({
  value,
  name,
  path = '$',
  depth = 0,
  isLast = true,
  className,
  ...props
}: JsonViewerNodeProps) => {
  const { isExpanded, toggle, setFocusedPath, treeRef } = useJsonViewer();
  const focused = useIsFocusedPath(path);
  const ancestry = React.useContext(JsonViewerAncestryContext);
  const kind = kindOf(value);
  const circular = (kind === 'object' || kind === 'array') && isCircular(ancestry, value);
  // Compared inline rather than via isExpandable() so the compiler sees a primitive memo dependency.
  const expandable = !circular && (kind === 'object' || kind === 'array');
  const entries = expandable ? entriesOf(value, kind) : [];
  const expanded = expandable && isExpanded(path, depth);
  const open = kind === 'array' ? '[' : '{';
  const close = kind === 'array' ? ']' : '}';

  const nodeCtx = React.useMemo<JsonViewerNodeCtx>(() => ({ expandable, expanded }), [expandable, expanded]);
  const childAncestry = React.useMemo<Ancestry | null>(
    () => (expandable ? { value: value as object, parent: ancestry } : ancestry),
    [expandable, value, ancestry],
  );

  const handleKeyDown = (e: React.KeyboardEvent<HTMLDivElement>) => {
    const row = e.currentTarget;
    const tree = treeRef.current;
    if (!tree) return;
    const rtl = getComputedStyle(row).direction === 'rtl';
    const expandKey = rtl ? 'ArrowLeft' : 'ArrowRight';
    const collapseKey = rtl ? 'ArrowRight' : 'ArrowLeft';
    const rows = Array.from(tree.querySelectorAll<HTMLElement>('[data-slot="json-viewer-row"]'));
    const index = rows.indexOf(row);
    const focusRow = (target: HTMLElement | undefined) => {
      if (!target) return;
      e.preventDefault();
      target.focus();
    };

    switch (e.key) {
      case 'ArrowDown':
        focusRow(rows[index + 1]);
        break;
      case 'ArrowUp':
        focusRow(rows[index - 1]);
        break;
      case 'Home':
        focusRow(rows[0]);
        break;
      case 'End':
        focusRow(rows[rows.length - 1]);
        break;
      case expandKey:
        if (!expandable) break;
        e.preventDefault();
        if (expanded) focusRow(rows[index + 1]);
        else toggle(path, depth);
        break;
      case collapseKey: {
        e.preventDefault();
        if (expandable && expanded) {
          toggle(path, depth);
          break;
        }
        const parent = row
          .closest('[data-slot="json-viewer-node"]')
          ?.parentElement?.closest<HTMLElement>('[data-slot="json-viewer-node"]');
        focusRow(parent?.querySelector<HTMLElement>('[data-slot="json-viewer-row"]') ?? undefined);
        break;
      }
      case 'Enter':
      case ' ': {
        e.preventDefault();
        if (expandable) {
          toggle(path, depth);
          break;
        }
        row.querySelector<HTMLButtonElement>('[data-slot="json-viewer-more"]')?.click();
        break;
      }
    }
  };

  return (
    <JsonViewerNodeContext.Provider value={nodeCtx}>
      <div
        data-slot="json-viewer-node"
        data-path={path}
        data-kind={kind}
        data-state={expandable ? (expanded ? 'open' : 'closed') : undefined}
        className={cn(className)}
        {...props}
      >
        <div
          role="treeitem"
          data-slot="json-viewer-row"
          tabIndex={focused ? 0 : -1}
          aria-level={depth + 1}
          aria-expanded={expandable ? expanded : undefined}
          aria-selected={focused}
          onFocus={() => setFocusedPath(path)}
          onKeyDown={handleKeyDown}
          onClick={(e) => {
            if (!expandable) return;
            if ((e.target as HTMLElement).closest('button')) return;
            toggle(path, depth);
          }}
          className={cn(
            'flex items-start gap-1 rounded-sm px-1 outline-none',
            'hover:bg-accent/50 focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset',
            expandable && 'cursor-pointer',
          )}
        >
          <JsonViewerToggle />
          {name !== undefined && (
            <>
              <JsonViewerKey>{name}</JsonViewerKey>
              <span className="text-muted-foreground">:</span>
            </>
          )}
          {expandable ? (
            expanded ? (
              <span className="text-muted-foreground">{open}</span>
            ) : (
              <span className="text-muted-foreground">
                {open}
                <span data-slot="json-viewer-count" className="px-0.5">
                  {entries.length}
                </span>
                {close}
                {!isLast && ','}
              </span>
            )
          ) : circular ? (
            <>
              <span data-slot="json-viewer-value" data-kind="circular" className="text-muted-foreground italic">
                {CIRCULAR}
              </span>
              {!isLast ? <span className="text-muted-foreground">,</span> : null}
            </>
          ) : (
            <>
              <JsonViewerValue value={value} />
              {!isLast ? <span className="text-muted-foreground">,</span> : null}
            </>
          )}
        </div>
        {expandable && expanded ? (
          <>
            <div role="group" data-slot="json-viewer-children" className="ms-3 border-s border-border ps-2">
              <JsonViewerAncestryContext.Provider value={childAncestry}>
                {entries.map(([key, child], i) => (
                  <JsonViewerNode
                    key={key}
                    name={key}
                    value={child}
                    path={childPath(path, key, kind)}
                    depth={depth + 1}
                    isLast={i === entries.length - 1}
                  />
                ))}
              </JsonViewerAncestryContext.Provider>
            </div>
            <div data-slot="json-viewer-close" className="ps-6 text-muted-foreground">
              {close}
              {!isLast && ','}
            </div>
          </>
        ) : null}
      </div>
    </JsonViewerNodeContext.Provider>
  );
};

const JsonViewerNode = React.memo(JsonViewerNodeImpl);

const JsonViewerToggle = ({ className, ...props }: Omit<React.ComponentProps<'span'>, 'children'>) => {
  const { expandable, expanded } = useJsonViewerNode();

  return (
    <span
      data-slot="json-viewer-toggle"
      data-state={expandable ? (expanded ? 'open' : 'closed') : undefined}
      aria-hidden
      className={cn(
        'mt-[3px] inline-flex size-4 shrink-0 items-center justify-center text-muted-foreground',
        !expandable && 'invisible',
        className,
      )}
      {...props}
    >
      <ChevronRight
        className={cn(
          'size-3.5 transition-transform motion-reduce:transition-none rtl:rotate-180',
          expanded && 'rotate-90 rtl:rotate-90',
        )}
      />
    </span>
  );
};

const JsonViewerKey = ({ className, ...props }: React.ComponentProps<'span'>) => {
  return <span data-slot="json-viewer-key" className={cn('text-foreground', className)} {...props} />;
};

export interface JsonViewerValueProps extends Omit<React.ComponentProps<'span'>, 'children'> {
  value: unknown;
}

const JsonViewerValue = ({ value, className, ...props }: JsonViewerValueProps) => {
  const { maxStringLength } = useJsonViewer();
  const [showAll, setShowAll] = React.useState(false);
  const kind = kindOf(value);

  let text: string;
  switch (kind) {
    case 'string':
      text = value as string;
      break;
    case 'date':
      // toISOString throws a RangeError on an invalid date.
      text = Number.isNaN((value as Date).getTime()) ? 'Invalid Date' : (value as Date).toISOString();
      break;
    case 'number':
      text = Object.is(value, -0) ? '-0' : String(value);
      break;
    case 'bigint':
      text = `${String(value)}n`;
      break;
    case 'boolean':
      text = String(value);
      break;
    case 'null':
      text = 'null';
      break;
    case 'undefined':
      text = 'undefined';
      break;
    case 'function':
      text = `ƒ ${(value as { name?: string }).name || 'anonymous'}()`;
      break;
    case 'symbol':
      text = String(value);
      break;
    default:
      text = stringify(value, 0) ?? String(value);
  }

  const isText = kind === 'string' || kind === 'date';
  const truncated = isText && !showAll && text.length > maxStringLength;
  const shown = truncated ? text.slice(0, maxStringLength) : text;

  return (
    <span
      data-slot="json-viewer-value"
      data-kind={kind}
      className={cn(
        'break-all whitespace-pre-wrap',
        isText && 'text-success',
        (kind === 'number' || kind === 'bigint') && 'text-info',
        kind === 'boolean' && 'text-warning',
        (kind === 'null' || kind === 'undefined') && 'text-muted-foreground italic',
        (kind === 'function' || kind === 'symbol') && 'text-muted-foreground',
        className,
      )}
      {...props}
    >
      {isText ? `"${shown}` : shown}
      {truncated && (
        <>
          <span className="text-muted-foreground">…</span>
          <button
            type="button"
            data-slot="json-viewer-more"
            tabIndex={-1}
            onClick={(e) => {
              e.stopPropagation();
              setShowAll(true);
            }}
            className="mx-1 rounded-sm bg-accent px-1 text-[10px] text-muted-foreground hover:text-foreground"
          >
            +{text.length - maxStringLength}
          </button>
        </>
      )}
      {isText ? `"` : null}
    </span>
  );
};

export interface JsonViewerCopyProps extends Omit<React.ComponentProps<typeof CopyButton>, 'value'> {
  /** Defaults to the whole root value. */
  value?: unknown;
  indent?: number;
}

const JsonViewerCopy = ({ value, indent = 2, ...props }: JsonViewerCopyProps) => {
  const root = useJsonViewer();
  const source = value === undefined ? root.value : value;
  const getText = React.useCallback(() => stringify(source, indent) ?? '', [source, indent]);

  return <CopyButton data-slot="json-viewer-copy" size="sm" value={getText} {...props} />;
};

const JsonViewerExpandAll = ({ className, children, ...props }: React.ComponentProps<typeof Button>) => {
  const { expandAll } = useJsonViewer();

  return (
    <Button
      type="button"
      variant="ghost"
      size="xs"
      data-slot="json-viewer-expand-all"
      onClick={expandAll}
      className={cn('font-sans text-muted-foreground', className)}
      {...props}
    >
      <ChevronsUpDown />
      {children ?? 'Expand all'}
    </Button>
  );
};

const JsonViewerCollapseAll = ({ className, children, ...props }: React.ComponentProps<typeof Button>) => {
  const { collapseAll } = useJsonViewer();

  return (
    <Button
      type="button"
      variant="ghost"
      size="xs"
      data-slot="json-viewer-collapse-all"
      onClick={collapseAll}
      className={cn('font-sans text-muted-foreground', className)}
      {...props}
    >
      <ChevronsDownUp />
      {children ?? 'Collapse all'}
    </Button>
  );
};

export {
  JsonViewer,
  JsonViewerTree,
  JsonViewerNode,
  JsonViewerToggle,
  JsonViewerKey,
  JsonViewerValue,
  JsonViewerCopy,
  JsonViewerExpandAll,
  JsonViewerCollapseAll,
  useJsonViewer,
};