Ai Context

'use client';

import {
  type ComponentProps,
  type ReactNode,
  createContext,
  useContext,
} from 'react';

import Box from '@mui/material/Box';
import Button from '@mui/material/Button';
import LinearProgress from '@mui/material/LinearProgress';
import Typography from '@mui/material/Typography';
import type { LanguageModelUsage } from 'ai';
import { type ModelId, estimateCost } from 'tokenlens';

const PERCENT_MAX = 100;
const ICON_RADIUS = 10;
const ICON_VIEWBOX = 24;
const ICON_CENTER = 12;
const ICON_STROKE_WIDTH = 2;

type ContextSchema = {
  usedTokens: number;
  maxTokens: number;
  usage?: LanguageModelUsage;
  modelId?: ModelId;
};

const ContextContext = createContext<ContextSchema | null>(null);

const useContextValue = () => {
  const context = useContext(ContextContext);

  if (!context) {
    throw new Error('Context components must be used within Context');
  }

  return context;
};

export type ContextProps = {
  children: ReactNode;
} & ContextSchema;

export const Context = ({
  usedTokens,
  maxTokens,
  usage,
  modelId,
  children,
}: ContextProps) => (
  <ContextContext.Provider
    value={{
      usedTokens,
      maxTokens,
      usage,
      modelId,
    }}
  >
    {children}
  </ContextContext.Provider>
);

const ContextIcon = () => {
  const { usedTokens, maxTokens } = useContextValue();
  const circumference = 2 * Math.PI * ICON_RADIUS;
  const usedPercent = usedTokens / maxTokens;
  const dashOffset = circumference * (1 - usedPercent);

  return (
    <svg
      aria-label="Model context usage"
      height="20"
      role="img"
      style={{ color: 'currentcolor' }}
      viewBox={`0 0 ${ICON_VIEWBOX} ${ICON_VIEWBOX}`}
      width="20"
    >
      <circle
        cx={ICON_CENTER}
        cy={ICON_CENTER}
        fill="none"
        opacity="0.25"
        r={ICON_RADIUS}
        stroke="currentColor"
        strokeWidth={ICON_STROKE_WIDTH}
      />
      <circle
        cx={ICON_CENTER}
        cy={ICON_CENTER}
        fill="none"
        opacity="0.7"
        r={ICON_RADIUS}
        stroke="currentColor"
        strokeDasharray={`${circumference} ${circumference}`}
        strokeDashoffset={dashOffset}
        strokeLinecap="round"
        strokeWidth={ICON_STROKE_WIDTH}
        style={{ transformOrigin: 'center', transform: 'rotate(-90deg)' }}
      />
    </svg>
  );
};

export type ContextTriggerProps = ComponentProps<typeof Button> & {
  onOpen?: (event: React.MouseEvent<HTMLElement>) => void;
};

export const ContextTrigger = ({
  children,
  onOpen,
  sx,
  ...props
}: ContextTriggerProps) => {
  const { usedTokens, maxTokens } = useContextValue();
  const usedPercent = usedTokens / maxTokens;
  const renderedPercent = new Intl.NumberFormat('en-US', {
    style: 'percent',
    maximumFractionDigits: 1,
  }).format(usedPercent);

  return (
    children ?? (
      <Button
        type="button"
        variant="text"
        onClick={onOpen}
        sx={{
          minWidth: 'auto',
          p: 1,
          color: 'text.secondary',
          '&:hover': {
            bgcolor: 'action.hover',
          },
          ...sx,
        }}
        {...props}
      >
        <Typography
          component="span"
          sx={{ fontWeight: 500, mr: 1, color: 'text.secondary' }}
        >
          {renderedPercent}
        </Typography>
        <ContextIcon />
      </Button>
    )
  );
};

export type ContextContentProps = ComponentProps<typeof Box>;

export const ContextContent = ({
  sx,
  children,
  ...props
}: ContextContentProps) => (
  <Box
    sx={{
      minWidth: 240,
      '& > *:not(:last-child)': {
        borderBottom: 1,
        borderColor: 'divider',
      },
      ...sx,
    }}
    {...props}
  >
    {children}
  </Box>
);

export type ContextContentHeaderProps = ComponentProps<typeof Box>;

export const ContextContentHeader = ({
  children,
  sx,
  ...props
}: ContextContentHeaderProps) => {
  const { usedTokens, maxTokens } = useContextValue();
  const usedPercent = usedTokens / maxTokens;
  const displayPct = new Intl.NumberFormat('en-US', {
    style: 'percent',
    maximumFractionDigits: 1,
  }).format(usedPercent);
  const used = new Intl.NumberFormat('en-US', {
    notation: 'compact',
  }).format(usedTokens);
  const total = new Intl.NumberFormat('en-US', {
    notation: 'compact',
  }).format(maxTokens);

  return (
    <Box sx={{ width: '100%', p: 1.5, ...sx }} {...props}>
      {children ?? (
        <>
          <Box
            sx={{
              display: 'flex',
              alignItems: 'center',
              justifyContent: 'space-between',
              gap: 1.5,
              mb: 1,
            }}
          >
            <Typography variant="caption">{displayPct}</Typography>
            <Typography
              variant="caption"
              sx={{ fontFamily: 'monospace', color: 'text.secondary' }}
            >
              {used} / {total}
            </Typography>
          </Box>
          <LinearProgress
            variant="determinate"
            value={usedPercent * PERCENT_MAX}
            sx={{ height: 6, borderRadius: 1, bgcolor: 'action.hover' }}
          />
        </>
      )}
    </Box>
  );
};

export type ContextContentBodyProps = ComponentProps<typeof Box>;

export const ContextContentBody = ({
  children,
  sx,
  ...props
}: ContextContentBodyProps) => (
  <Box sx={{ width: '100%', p: 1.5, ...sx }} {...props}>
    {children}
  </Box>
);

export type ContextContentFooterProps = ComponentProps<typeof Box>;

export const ContextContentFooter = ({
  children,
  sx,
  ...props
}: ContextContentFooterProps) => {
  const { modelId, usage } = useContextValue();
  const costUSD = modelId
    ? estimateCost({
        modelId,
        usage: {
          input: usage?.inputTokens ?? 0,
          output: usage?.outputTokens ?? 0,
        },
      }).totalUSD
    : undefined;
  const totalCost = new Intl.NumberFormat('en-US', {
    style: 'currency',
    currency: 'USD',
  }).format(costUSD ?? 0);

  return (
    <Box
      sx={{
        display: 'flex',
        width: '100%',
        alignItems: 'center',
        justifyContent: 'space-between',
        gap: 1.5,
        bgcolor: 'action.hover',
        p: 1.5,
        ...sx,
      }}
      {...props}
    >
      {children ?? (
        <>
          <Typography variant="caption" sx={{ color: 'text.secondary' }}>
            Total cost
          </Typography>
          <Typography variant="caption">{totalCost}</Typography>
        </>
      )}
    </Box>
  );
};

export type ContextInputUsageProps = ComponentProps<typeof Box>;

export const ContextInputUsage = ({
  sx,
  children,
  ...props
}: ContextInputUsageProps) => {
  const { usage, modelId } = useContextValue();
  const inputTokens = usage?.inputTokens ?? 0;

  if (children) {
    return <>{children}</>;
  }

  if (!inputTokens) {
    return null;
  }

  const inputCost = modelId
    ? estimateCost({
        modelId,
        usage: { input: inputTokens, output: 0 },
      }).totalUSD
    : undefined;
  const inputCostText = new Intl.NumberFormat('en-US', {
    style: 'currency',
    currency: 'USD',
  }).format(inputCost ?? 0);

  return (
    <Box
      sx={{
        display: 'flex',
        alignItems: 'center',
        justifyContent: 'space-between',
        ...sx,
      }}
      {...props}
    >
      <Typography variant="caption" sx={{ color: 'text.secondary' }}>
        Input
      </Typography>
      <TokensWithCost costText={inputCostText} tokens={inputTokens} />
    </Box>
  );
};

export type ContextOutputUsageProps = ComponentProps<typeof Box>;

export const ContextOutputUsage = ({
  sx,
  children,
  ...props
}: ContextOutputUsageProps) => {
  const { usage, modelId } = useContextValue();
  const outputTokens = usage?.outputTokens ?? 0;

  if (children) {
    return <>{children}</>;
  }

  if (!outputTokens) {
    return null;
  }

  const outputCost = modelId
    ? estimateCost({
        modelId,
        usage: { input: 0, output: outputTokens },
      }).totalUSD
    : undefined;
  const outputCostText = new Intl.NumberFormat('en-US', {
    style: 'currency',
    currency: 'USD',
  }).format(outputCost ?? 0);

  return (
    <Box
      sx={{
        display: 'flex',
        alignItems: 'center',
        justifyContent: 'space-between',
        ...sx,
      }}
      {...props}
    >
      <Typography variant="caption" sx={{ color: 'text.secondary' }}>
        Output
      </Typography>
      <TokensWithCost costText={outputCostText} tokens={outputTokens} />
    </Box>
  );
};

export type ContextReasoningUsageProps = ComponentProps<typeof Box>;

export const ContextReasoningUsage = ({
  sx,
  children,
  ...props
}: ContextReasoningUsageProps) => {
  const { usage, modelId } = useContextValue();
  const reasoningTokens = usage?.reasoningTokens ?? 0;

  if (children) {
    return <>{children}</>;
  }

  if (!reasoningTokens) {
    return null;
  }

  const reasoningCost = modelId
    ? estimateCost({
        modelId,
        usage: { reasoningTokens },
      }).totalUSD
    : undefined;
  const reasoningCostText = new Intl.NumberFormat('en-US', {
    style: 'currency',
    currency: 'USD',
  }).format(reasoningCost ?? 0);

  return (
    <Box
      sx={{
        display: 'flex',
        alignItems: 'center',
        justifyContent: 'space-between',
        ...sx,
      }}
      {...props}
    >
      <Typography variant="caption" sx={{ color: 'text.secondary' }}>
        Reasoning
      </Typography>
      <TokensWithCost costText={reasoningCostText} tokens={reasoningTokens} />
    </Box>
  );
};

export type ContextCacheUsageProps = ComponentProps<typeof Box>;

export const ContextCacheUsage = ({
  sx,
  children,
  ...props
}: ContextCacheUsageProps) => {
  const { usage, modelId } = useContextValue();
  const cacheTokens = usage?.cachedInputTokens ?? 0;

  if (children) {
    return <>{children}</>;
  }

  if (!cacheTokens) {
    return null;
  }

  const cacheCost = modelId
    ? estimateCost({
        modelId,
        usage: { cacheReads: cacheTokens, input: 0, output: 0 },
      }).totalUSD
    : undefined;
  const cacheCostText = new Intl.NumberFormat('en-US', {
    style: 'currency',
    currency: 'USD',
  }).format(cacheCost ?? 0);

  return (
    <Box
      sx={{
        display: 'flex',
        alignItems: 'center',
        justifyContent: 'space-between',
        ...sx,
      }}
      {...props}
    >
      <Typography variant="caption" sx={{ color: 'text.secondary' }}>
        Cache
      </Typography>
      <TokensWithCost costText={cacheCostText} tokens={cacheTokens} />
    </Box>
  );
};

const TokensWithCost = ({
  tokens,
  costText,
}: {
  tokens?: number;
  costText?: string;
}) => (
  <Typography component="span" variant="caption">
    {tokens === undefined
      ? '—'
      : new Intl.NumberFormat('en-US', {
          notation: 'compact',
        }).format(tokens)}
    {costText ? (
      <Typography
        component="span"
        variant="caption"
        sx={{ ml: 1, color: 'text.secondary' }}
      >
        • {costText}
      </Typography>
    ) : null}
  </Typography>
);

export * from './ai-context';