stack/apps/web/src/components/dashboard/TokenBudget.tsx

import type { ReactElement } from "react";
import { Card, SectionHeader, ProgressBar, type ProgressBarVariant } from "@mosaic/ui";
import type { TokenBudgetEntry } from "@/lib/api/dashboard";

export interface TokenBudgetProps {
  budgets?: TokenBudgetEntry[] | undefined;
}

/* ------------------------------------------------------------------ */
/*  Helpers                                                            */
/* ------------------------------------------------------------------ */

const VARIANT_CYCLE: ProgressBarVariant[] = ["blue", "teal", "purple", "amber"];

function formatTokenCount(n: number): string {
  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
  if (n >= 1_000) return `${(n / 1_000).toFixed(0)}K`;
  return String(n);
}

interface ModelBudgetDisplay {
  id: string;
  label: string;
  usage: string;
  value: number;
  variant: ProgressBarVariant;
}

function mapBudgetToDisplay(entry: TokenBudgetEntry, index: number): ModelBudgetDisplay {
  const percent = entry.limit > 0 ? Math.round((entry.used / entry.limit) * 100) : 0;
  const usage =
    entry.limit > 0
      ? `${formatTokenCount(entry.used)} / ${formatTokenCount(entry.limit)}`
      : "unlimited";

  return {
    id: entry.model,
    label: entry.model,
    usage,
    value: percent,
    variant: VARIANT_CYCLE[index % VARIANT_CYCLE.length] ?? "blue",
  };
}

/* ------------------------------------------------------------------ */
/*  Components                                                         */
/* ------------------------------------------------------------------ */

interface ModelRowProps {
  model: ModelBudgetDisplay;
}

function ModelRow({ model }: ModelRowProps): ReactElement {
  return (
    <div style={{ marginBottom: 14 }}>
      <div
        style={{
          display: "flex",
          alignItems: "center",
          justifyContent: "space-between",
          marginBottom: 5,
        }}
      >
        <span
          style={{
            fontSize: "0.8rem",
            fontWeight: 600,
            color: "var(--text-2)",
            fontFamily: "var(--mono)",
          }}
        >
          {model.label}
        </span>
        <span
          style={{
            fontSize: "0.72rem",
            color: "var(--muted)",
            fontFamily: "var(--mono)",
          }}
        >
          {model.usage}
        </span>
      </div>
      <ProgressBar
        value={model.value}
        variant={model.variant}
        label={`${model.label} token usage`}
      />
    </div>
  );
}

export function TokenBudget({ budgets }: TokenBudgetProps): ReactElement {
  const displayModels = budgets ? budgets.map(mapBudgetToDisplay) : [];

  return (
    <Card>
      <SectionHeader title="Token Budget" subtitle="Usage by model" />
      <div>
        {displayModels.length > 0 ? (
          displayModels.map((model) => <ModelRow key={model.id} model={model} />)
        ) : (
          <div
            style={{
              padding: "24px 0",
              textAlign: "center",
              fontSize: "0.8rem",
              color: "var(--muted)",
            }}
          >
            No budget data
          </div>
        )}
      </div>
    </Card>
  );
}