import type { ChatMessage } from '@genesislcap/foundation-ai';
import DOMPurify from 'dompurify';

/** Per-message token usage when the provider reports it. */
const messageTokenUsage = (m: ChatMessage): number => (m.inputTokens ?? 0) + (m.outputTokens ?? 0);

/**
 * Sum input + output tokens across a message list, recursing into each tool
 * call's `subAgentTrace` so sub-agent LLM usage contributes to the total.
 *
 * Returns 0 when no message carries token metadata.
 *
 * @internal
 */
export function sumTokens(messages: readonly ChatMessage[]): number {
  let total = 0;
  for (const m of messages) {
    total += messageTokenUsage(m);
    for (const tc of m.toolCalls ?? []) {
      if (tc.subAgentTrace) total += sumTokens(tc.subAgentTrace);
    }
  }
  return total;
}

const TOKEN_SCALE_DECIMAL_FACTOR = 10;
const TOKENS_PER_MILLION = 1_000_000;
const TOKENS_PER_THOUSAND = 1_000;
const SUFFIX_ROLLOVER_THRESHOLD = 1000;

const roundScaled = (scaled: number): number =>
  Math.round(scaled * TOKEN_SCALE_DECIMAL_FACTOR) / TOKEN_SCALE_DECIMAL_FACTOR;

const formatScaledTokenMarkup = (rounded: number, suffix: string): string => {
  const value = Number.isInteger(rounded) ? String(rounded) : rounded.toFixed(1);
  return `<span class="value">${value}</span><span class="suffix">${suffix}</span>`;
};

const scaledSuffixLabel = (tokens: number): string | undefined => {
  if (tokens >= TOKENS_PER_MILLION) return 'M';
  if (tokens >= TOKENS_PER_THOUSAND) {
    const kRounded = roundScaled(tokens / TOKENS_PER_THOUSAND);
    if (kRounded >= SUFFIX_ROLLOVER_THRESHOLD) return 'M';
    return 'K';
  }
  return undefined;
};

const scaledNumericValue = (tokens: number): number | undefined => {
  const suffix = scaledSuffixLabel(tokens);
  if (!suffix) return undefined;
  if (suffix === 'M') return roundScaled(tokens / TOKENS_PER_MILLION);
  return roundScaled(tokens / TOKENS_PER_THOUSAND);
};

/**
 * Compact plain-text token count (e.g. 1_400_000 → "1.4M").
 *
 * @internal
 */
export function formatCompactTokenCountPlain(tokens: number): string {
  const suffix = scaledSuffixLabel(tokens);
  if (!suffix) return tokens.toLocaleString();
  const rounded = scaledNumericValue(tokens)!;
  const value = Number.isInteger(rounded) ? String(rounded) : rounded.toFixed(1);
  return `${value}${suffix}`;
}

/**
 * Compact display for large token counts (e.g. 1_400_000 → value "1.4" + suffix "M").
 * Returns HTML markup so the suffix can be styled separately in the settings modal.
 *
 * @internal
 */
export function formatCompactTokenCount(tokens: number): string {
  const suffix = scaledSuffixLabel(tokens);
  if (!suffix) return tokens.toLocaleString();
  return formatScaledTokenMarkup(scaledNumericValue(tokens)!, suffix);
}

/** {@link formatCompactTokenCount} passed through DOMPurify for `:innerHTML` bindings. */
export function formatCompactTokenCountSafe(tokens: number): string {
  return DOMPurify.sanitize(formatCompactTokenCount(tokens));
}
