Agent Audio Visualizer Grid

'use client';

import React, {
  type CSSProperties,
  Children,
  type ComponentProps,
  type ReactNode,
  cloneElement,
  isValidElement,
  memo,
  useMemo,
} from 'react';
import { type VariantProps, cva } from 'class-variance-authority';
import type { LocalAudioTrack, RemoteAudioTrack } from 'livekit-client';
import {
  type AgentState,
  type TrackReferenceOrPlaceholder,
  useMultibandTrackVolume,
} from '@livekit/components-react';
import {
  type Coordinate,
  useAgentAudioVisualizerGridAnimator,
} from '@/hooks/agents-ui/use-agent-audio-visualizer-grid';
import { cn } from '@/lib/utils';

/**
 * Resizes an array of per-band volume values to exactly `count` entries.
 * Excess values are trimmed from the end; if there are too few, the last
 * value is duplicated to fill the remainder. An empty array is padded with 0s.
 */
export function normalizeVolumeBands(bands: number[], count: number): number[] {
  if (bands.length === count) return bands;
  if (bands.length > count) return bands.slice(0, count);
  const lastValue = bands[bands.length - 1] ?? 0;
  return [...bands, ...new Array(count - bands.length).fill(lastValue)];
}

function cloneSingleChild(
  children: ReactNode | ReactNode[],
  props?: Record<string, unknown>,
  key?: unknown,
) {
  return Children.map(children, (child) => {
    // Checking isValidElement is the safe way and avoids a typescript error too.
    if (isValidElement(child) && Children.only(children)) {
      const childProps = child.props as Record<string, unknown>;
      if (childProps.className) {
        // make sure we retain classnames of both passed props and child
        props ??= {};
        props.className = cn(childProps.className as string, props.className as string);
        props.style = {
          ...(childProps.style as CSSProperties),
          ...(props.style as CSSProperties),
        };
      }
      return cloneElement(child, { ...props, key: key ? String(key) : undefined });
    }
    return child;
  });
}

export const AgentAudioVisualizerGridCellVariants = cva(
  [
    'w-1 h-1 rounded-full bg-current/10 place-self-center transition-all ease-out',
    'data-[lk-highlighted=true]:bg-current',
  ],
  {
    variants: {
      size: {
        icon: ['w-[2px] h-[2px]'],
        sm: ['w-[4px] h-[4px]'],
        md: ['w-[8px] h-[8px]'],
        lg: ['w-[12px] h-[12px]'],
        xl: ['w-[16px] h-[16px]'],
      },
    },
    defaultVariants: {
      size: 'md',
    },
  },
);

export const AgentAudioVisualizerGridVariants = cva('grid', {
  variants: {
    size: {
      icon: ['gap-[2px]'],
      sm: ['gap-[4px]'],
      md: ['gap-[8px]'],
      lg: ['gap-[12px]'],
      xl: ['gap-[16px]'],
    },
  },
  defaultVariants: {
    size: 'md',
  },
});

/**
 * Configuration options for the grid visualizer.
 */
export interface GridOptions {
  /**
   * The radius for the animation spread effect.
   */
  radius?: number;
  /**
   * The interval in milliseconds between animation frames.
   * @defaultValue 100
   */
  interval?: number;
  /**
   * The number of rows in the grid.
   * @defaultValue 5
   */
  rowCount?: number;
  /**
   * The number of columns in the grid.
   * @defaultValue 5
   */
  columnCount?: number;
  /**
   * Additional CSS class names to apply to the container.
   */
  className?: string;
}

const sizeDefaults = {
  icon: 3,
  sm: 5,
  md: 5,
  lg: 5,
  xl: 5,
};

function useGrid(
  size: VariantProps<typeof AgentAudioVisualizerGridVariants>['size'] = 'md',
  columnCount = sizeDefaults[size as keyof typeof sizeDefaults],
  rowCount = sizeDefaults[size as keyof typeof sizeDefaults],
) {
  return useMemo(() => {
    const _columnCount = columnCount;
    const _rowCount = rowCount ?? columnCount;
    const items = new Array(_columnCount * _rowCount).fill(0).map((_, idx) => idx);

    return { columnCount: _columnCount, rowCount: _rowCount, items };
  }, [columnCount, rowCount]);
}

interface GridCellProps {
  index: number;
  state: AgentState;
  interval: number;
  rowCount: number;
  columnCount: number;
  volumeBands: number[];
  highlightedCoordinate: Coordinate;
  children?: ReactNode;
}

const GridCell = memo(function GridCell({
  index,
  state,
  interval,
  rowCount,
  columnCount,
  volumeBands,
  highlightedCoordinate,
  children,
}: GridCellProps) {
  if (state === 'speaking') {
    const y = Math.floor(index / columnCount);
    const rowMidPoint = Math.floor(rowCount / 2);
    const volumeChunks = 1 / (rowMidPoint + 1);
    const distanceToMid = Math.abs(rowMidPoint - y);
    const threshold = distanceToMid * volumeChunks;
    const isHighlighted = (volumeBands[index % columnCount] ?? 0) >= threshold;

    return cloneSingleChild(children, {
      'data-lk-index': index,
      'data-lk-highlighted': isHighlighted,
    });
  }

  const isHighlighted =
    highlightedCoordinate.x === index % columnCount &&
    highlightedCoordinate.y === Math.floor(index / columnCount);

  const transitionDurationInSeconds = interval / (isHighlighted ? 1000 : 100);

  return cloneSingleChild(children, {
    'data-lk-index': index,
    'data-lk-highlighted': isHighlighted,
    style: {
      transitionDuration: `${transitionDurationInSeconds}s`,
    },
  });
});

/**
 * Props for the AgentAudioVisualizerGrid component.
 */
export type AgentAudioVisualizerGridProps = GridOptions & {
  /**
   * The size of the visualizer.
   * @defaultValue 'md'
   */
  size?: 'icon' | 'sm' | 'md' | 'lg' | 'xl';
  /**
   * The current state of the agent. Determines the animation pattern.
   * @defaultValue 'connecting'
   */
  state?: AgentState;
  /**
   * The color of the grid cells in hexidecimal format.
   */
  color?: `#${string}`;
  /**
   * The audio track to visualize. Can be a local/remote audio track or a track reference.
   */
  audioTrack?: LocalAudioTrack | RemoteAudioTrack | TrackReferenceOrPlaceholder;
  /**
   * Volume values (0-1) to use instead of the values computed from the audioTrack.
   * The volumeBands.length should match columnCount.
   */
  volumeBands?: number[];
  /**
   * Additional CSS class names to apply to the container.
   */
  className?: string;
  /**
   * Custom element to render as grid cells. Each child receives data-lk-index
   * and data-lk-highlighted props.
   */
  children?: ReactNode;
} & VariantProps<typeof AgentAudioVisualizerGridVariants>;

/**
 * A grid-style audio visualizer that responds to agent state and audio levels.
 * Displays an animated grid of cells that react to the current agent state
 * and audio volume when speaking.
 *
 * @extends ComponentProps<'div'>
 *
 * @example
 * ```tsx
 * <AgentAudioVisualizerGrid
 *   size="md"
 *   state="speaking"
 *   rowCount={5}
 *   columnCount={5}
 *   audioTrack={agentAudioTrack}
 * />
 * ```
 */
export function AgentAudioVisualizerGrid({
  size = 'md',
  state = 'connecting',
  radius,
  color,
  rowCount: _rowCount = 5,
  columnCount: _columnCount = 5,
  interval = 100,
  className,
  children,
  audioTrack,
  volumeBands,
  style,
  ...props
}: AgentAudioVisualizerGridProps & ComponentProps<'div'>) {
  const { columnCount, rowCount, items } = useGrid(size, _columnCount, _rowCount);
  const highlightedCoordinate = useAgentAudioVisualizerGridAnimator(
    state,
    rowCount,
    columnCount,
    interval,
    radius,
  );
  const multibandVolume = useMultibandTrackVolume(audioTrack, {
    bands: columnCount,
    loPass: 100,
    hiPass: 200,
  });
  const resolvedVolumeBands = volumeBands
    ? normalizeVolumeBands(volumeBands, columnCount)
    : multibandVolume;

  if (children && Array.isArray(children)) {
    throw new Error('AgentAudioVisualizerGrid children must be a single element.');
  }

  return (
    <div
      data-lk-state={state}
      className={cn(AgentAudioVisualizerGridVariants({ size }), className)}
      style={
        { ...style, gridTemplateColumns: `repeat(${columnCount}, 1fr)`, color } as CSSProperties
      }
      {...props}
    >
      {items.map((idx) => (
        <GridCell
          key={idx}
          index={idx}
          state={state}
          interval={interval}
          rowCount={rowCount}
          columnCount={columnCount}
          volumeBands={resolvedVolumeBands}
          highlightedCoordinate={highlightedCoordinate}
        >
          {children ?? <div className={AgentAudioVisualizerGridCellVariants({ size })} />}
        </GridCell>
      ))}
    </div>
  );
}

import { useEffect, useState } from 'react';
import { type AgentState } from '@livekit/components-react';

export interface Coordinate {
  x: number;
  y: number;
}

export function generateConnectingSequence(rows: number, columns: number, radius: number) {
  const seq = [];
  const centerY = Math.floor(rows / 2);

  // Calculate the boundaries of the ring based on the ring distance
  const topLeft = {
    x: Math.max(0, centerY - radius),
    y: Math.max(0, centerY - radius),
  };
  const bottomRight = {
    x: columns - 1 - topLeft.x,
    y: Math.min(rows - 1, centerY + radius),
  };

  // Top edge
  for (let x = topLeft.x; x <= bottomRight.x; x++) {
    seq.push({ x, y: topLeft.y });
  }

  // Right edge
  for (let y = topLeft.y + 1; y <= bottomRight.y; y++) {
    seq.push({ x: bottomRight.x, y });
  }

  // Bottom edge
  for (let x = bottomRight.x - 1; x >= topLeft.x; x--) {
    seq.push({ x, y: bottomRight.y });
  }

  // Left edge
  for (let y = bottomRight.y - 1; y > topLeft.y; y--) {
    seq.push({ x: topLeft.x, y });
  }

  return seq;
}

export function generateListeningSequence(rows: number, columns: number) {
  const center = { x: Math.floor(columns / 2), y: Math.floor(rows / 2) };
  const noIndex = { x: -1, y: -1 };

  return [center, noIndex, noIndex, noIndex, noIndex, noIndex, noIndex, noIndex, noIndex];
}

export function generateThinkingSequence(rows: number, columns: number) {
  const seq = [];
  const y = Math.floor(rows / 2);
  for (let x = 0; x < columns; x++) {
    seq.push({ x, y });
  }
  for (let x = columns - 1; x >= 0; x--) {
    seq.push({ x, y });
  }

  return seq;
}

export function useAgentAudioVisualizerGridAnimator(
  state: AgentState,
  rows: number,
  columns: number,
  interval: number,
  radius?: number,
): Coordinate {
  const [index, setIndex] = useState(0);
  const [sequence, setSequence] = useState<Coordinate[]>(() => [
    {
      x: Math.floor(columns / 2),
      y: Math.floor(rows / 2),
    },
  ]);

  useEffect(() => {
    const clampedRadius = radius
      ? Math.min(radius, Math.floor(Math.max(rows, columns) / 2))
      : Math.floor(Math.max(rows, columns) / 2);

    if (state === 'thinking') {
      setSequence(generateThinkingSequence(rows, columns));
    } else if (state === 'connecting' || state === 'initializing') {
      const sequence = [...generateConnectingSequence(rows, columns, clampedRadius)];
      setSequence(sequence);
    } else if (state === 'listening') {
      setSequence(generateListeningSequence(rows, columns));
    } else {
      setSequence([{ x: Math.floor(columns / 2), y: Math.floor(rows / 2) }]);
    }
    setIndex(0);
  }, [state, rows, columns, radius]);

  useEffect(() => {
    if (state === 'speaking') {
      return;
    }

    const indexInterval = setInterval(() => {
      setIndex((prev) => {
        return prev + 1;
      });
    }, interval);

    return () => clearInterval(indexInterval);
  }, [interval, columns, rows, state, sequence.length]);

  return (
    sequence[index % sequence.length] ?? { x: Math.floor(columns / 2), y: Math.floor(rows / 2) }
  );
}