pi-statusline/modules/usage.ts

import type { ModuleContext, ModuleSpec, RenderedModule, UsageTextPart, UsageWindows } from "../types";
import type { UsageLimit, UsageReport } from "../usage";

function usageFraction(limit: UsageLimit): number | undefined {
  if (limit.amount.usedFraction !== undefined) return Math.min(Math.max(limit.amount.usedFraction, 0), 1);
  if (limit.amount.used !== undefined && limit.amount.limit !== undefined && limit.amount.limit > 0) {
    return Math.min(Math.max(limit.amount.used / limit.amount.limit, 0), 1);
  }
  return undefined;
}

function duration(valueMs: number | undefined): string | undefined {
  if (valueMs === undefined || !Number.isFinite(valueMs)) return undefined;
  const ms = Math.max(0, valueMs);
  const days = Math.floor(ms / 86_400_000);
  const hours = Math.floor((ms % 86_400_000) / 3_600_000);
  const minutes = Math.round((ms % 3_600_000) / 60_000);
  if (days > 0) return `${days}d${hours}h`;
  if (hours > 0) return `${hours}h${minutes}m`;
  return `${minutes}m`;
}

function compactNumber(value: number | undefined): string | undefined {
  if (value === undefined) return undefined;
  if (value >= 1_000_000) return `${(value / 1_000_000).toFixed(1)}M`;
  if (value >= 1_000) return `${(value / 1_000).toFixed(1)}k`;
  return `${Math.round(value)}`;
}

function lerp(start: number, end: number, amount: number): number {
  return Math.round(start + (end - start) * amount);
}

function fgRgb(text: string, r: number, g: number, b: number): string {
  return `\x1b[38;2;${r};${g};${b}m${text}\x1b[0m`;
}

function usageColor(fraction: number): { r: number; g: number; b: number } {
  const green = { r: 34, g: 197, b: 94 };
  const yellow = { r: 234, g: 179, b: 8 };
  const red = { r: 239, g: 68, b: 68 };

  // Continuous Green → Yellow → Red Gradient
  if (fraction <= 0.5) {
    const amount = fraction / 0.5;
    return {
      r: lerp(green.r, yellow.r, amount),
      g: lerp(green.g, yellow.g, amount),
      b: lerp(green.b, yellow.b, amount),
    };
  }

  const amount = (fraction - 0.5) / 0.5;
  return {
    r: lerp(yellow.r, red.r, amount),
    g: lerp(yellow.g, red.g, amount),
    b: lerp(yellow.b, red.b, amount),
  };
}

function findWindowLimits(report: UsageReport): UsageWindows {
  const windowId = (limit: UsageLimit) => (limit.window?.id ?? limit.scope.windowId ?? "").toLowerCase();
  const current = report.limits.find(limit => ["5h", "primary"].includes(windowId(limit))) ?? report.limits[0];
  const week = report.limits.find(limit => ["7d", "secondary"].includes(windowId(limit))) ?? report.limits.find(limit => limit !== current);
  return { current, week };
}

function lineBar(limit: UsageLimit | undefined, width: number, theme?: any, color = true): string {
  if (width <= 0) return "";
  const fraction = limit ? usageFraction(limit) ?? 0 : 0;

  const filled = Math.round(fraction * width);
  const empty = width - filled;
  const filledText = "━".repeat(filled);
  const emptyText = "─".repeat(empty);
  if (!color) return `${filledText}${emptyText}`;

  const { r, g, b } = usageColor(fraction);
  const coloredFilled = filledText ? fgRgb(filledText, r, g, b) : "";
  const dimEmpty = emptyText && theme?.fg ? theme.fg("dim", emptyText) : emptyText;
  return `${coloredFilled}${dimEmpty}`;
}

function selectedLimit(moduleCtx: ModuleContext, spec: ModuleSpec): UsageLimit | undefined {
  if (!moduleCtx.state.report || moduleCtx.state.error) return undefined;
  const windows = findWindowLimits(moduleCtx.state.report);
  const window = typeof spec === "string" ? "both" : spec.window ?? "both";
  if (window === "week") return windows.week;
  return windows.current ?? windows.week;
}

function textPart(limit: UsageLimit, part: UsageTextPart): string | undefined {
  const fraction = usageFraction(limit);
  switch (part) {
    case "percent":
      return fraction === undefined ? undefined : `${Math.round(fraction * 100)}%`;
    case "time":
      return duration(limit.window?.resetsAt !== undefined ? limit.window.resetsAt - Date.now() : undefined);
    case "used":
      return compactNumber(limit.amount.used);
    case "limit":
      return compactNumber(limit.amount.limit);
    case "remaining":
      return compactNumber(limit.amount.remaining);
  }
}

function fallbackTextPart(part: UsageTextPart): string | undefined {
  switch (part) {
    case "percent":
      return "0%";
    case "time":
      return "∞";
    case "used":
      return "0";
    case "remaining":
      return "∞";
    case "limit":
      return undefined;
  }
}

function textUsage(moduleCtx: ModuleContext, spec: ModuleSpec): string {
  // Usage Text Parts
  const parts = typeof spec === "string"
    ? moduleCtx.config.modules.usage?.parts ?? ["percent", "time"]
    : spec.parts ?? moduleCtx.config.modules.usage?.parts ?? ["percent", "time"];
  const separator = typeof spec === "string"
    ? moduleCtx.config.modules.usage?.separator ?? " | "
    : spec.separator ?? moduleCtx.config.modules.usage?.separator ?? " | ";

  const limit = !moduleCtx.state.report || moduleCtx.state.error ? undefined : selectedLimit(moduleCtx, spec);
  return parts
    .map(part => limit ? textPart(limit, part) : fallbackTextPart(part))
    .filter((part): part is string => Boolean(part))
    .join(separator);
}

export function usageModule(moduleCtx: ModuleContext, spec: ModuleSpec): RenderedModule {
  const style = typeof spec === "string" ? "line" : spec.style ?? "line";
  const grow = typeof spec === "string" ? 1 : spec.grow ?? 1;

  if (style === "text") return { text: textUsage(moduleCtx, spec) };

  return {
    grow,
    render(width: number, theme?: any): string {
      const limit = !moduleCtx.state.report || moduleCtx.state.error ? undefined : selectedLimit(moduleCtx, spec);
      return lineBar(limit, width, theme, moduleCtx.config.modules.usage?.colorBars !== false);
    },
  };
}