import { Suspense, useCallback } from "react";
import type { PressEvent } from "react-aria";
import { Pressable } from "react-aria";

import type { TextProps } from "@phoenix/components";
import {
  Loading,
  RichTooltip,
  TooltipArrow,
  TooltipTrigger,
} from "@phoenix/components";

import { SpanTokenCountDetails } from "./SpanTokenCountDetails";
import { TokenCount } from "./TokenCount";

type SpanTokenCountProps = {
  /**
   * The total number of tokens in the prompt and completion
   */
  tokenCountTotal: number;
  /**
   * The id of the node (span, trace, session, etc.)
   */

  nodeId: string;
  /**
   * The size of the icon and text
   */
  size?: TextProps["size"];
  /**
   * The color of the icon and text
   */
  color?: TextProps["color"];
};

/**
 * Displays the number of tokens in the prompt and completion
 */
export function SpanTokenCount(props: SpanTokenCountProps) {
  const handlePress = useCallback((e: PressEvent) => {
    e.continuePropagation(); // allow click to propagate to parent
  }, []);

  return (
    <TooltipTrigger>
      <Pressable onPress={handlePress}>
        <TokenCount
          size={props.size}
          color={props.color}
          role="button"
          tabIndex={0}
        >
          {props.tokenCountTotal}
        </TokenCount>
      </Pressable>
      <RichTooltip>
        <TooltipArrow />
        <Suspense fallback={<Loading />}>
          <SpanTokenCountDetails spanNodeId={props.nodeId} />
        </Suspense>
      </RichTooltip>
    </TooltipTrigger>
  );
}
