export type ProviderUsageStatus = "AVAILABLE" | "PARTIAL" | "MISSING" | "INCONSISTENT";
export type ProviderUsageField = "inputTokens" | "outputTokens" | "totalTokens";

export interface ProviderUsage {
  status: ProviderUsageStatus;
  inputTokens: number | null;
  outputTokens: number | null;
  totalTokens: number | null;
  missing: ProviderUsageField[];
  inconsistent: ProviderUsageField[];
  derivedTotal: boolean;
  sources: Partial<Record<ProviderUsageField, string[]>>;
}

const ALIASES: Record<ProviderUsageField, Set<string>> = {
  inputTokens: new Set(["input_tokens", "inputtokens", "prompt_tokens", "prompttokens", "prompt_token_count", "prompttokencount", "input_token_count", "inputtokencount"]),
  outputTokens: new Set(["output_tokens", "outputtokens", "completion_tokens", "completiontokens", "candidates_token_count", "candidatestokencount", "output_token_count", "outputtokencount"]),
  totalTokens: new Set(["total_tokens", "totaltokens", "total_token_count", "totaltokencount"]),
};

function keyName(key: string): string {
  return key.replace(/[A-Z]/gu, (letter) => `_${letter.toLowerCase()}`).replace(/-/gu, "_").toLowerCase();
}

function tokenCount(value: unknown): number | undefined {
  if (typeof value === "number" && Number.isSafeInteger(value) && value >= 0) return value;
  if (typeof value === "string" && /^\d+$/u.test(value)) {
    const parsed = Number(value);
    if (Number.isSafeInteger(parsed)) return parsed;
  }
  return undefined;
}

interface Observation { value: number; path: string; }

/** Normalize OpenAI, Responses, Anthropic, Gemini and nested gateway usage shapes. */
export function parseProviderUsage(input: unknown): ProviderUsage {
  const observations: Record<ProviderUsageField, Observation[]> = { inputTokens: [], outputTokens: [], totalTokens: [] };
  const seen = new Set<object>();
  const visit = (value: unknown, path: string): void => {
    if (typeof value !== "object" || value === null) return;
    if (seen.has(value)) return;
    seen.add(value);
    if (Array.isArray(value)) {
      value.forEach((item, index) => visit(item, `${path}[${index}]`));
      return;
    }
    for (const [key, child] of Object.entries(value)) {
      const normalized = keyName(key);
      for (const field of Object.keys(ALIASES) as ProviderUsageField[]) {
        if (ALIASES[field].has(normalized)) {
          const count = tokenCount(child);
          if (count !== undefined) observations[field].push({ value: count, path: `${path}.${key}` });
        }
      }
      visit(child, `${path}.${key}`);
    }
  };
  visit(input, "$");

  const values = {} as Record<ProviderUsageField, number | null>;
  const inconsistent: ProviderUsageField[] = [];
  const sources: Partial<Record<ProviderUsageField, string[]>> = {};
  for (const field of Object.keys(observations) as ProviderUsageField[]) {
    const items = observations[field];
    const unique = [...new Set(items.map((item) => item.value))];
    values[field] = unique.length === 0 ? null : unique[0];
    sources[field] = items.map((item) => item.path);
    if (unique.length > 1) inconsistent.push(field);
  }

  let derivedTotal = false;
  if (values.totalTokens === null && values.inputTokens !== null && values.outputTokens !== null && inconsistent.length === 0) {
    values.totalTokens = values.inputTokens + values.outputTokens;
    derivedTotal = true;
  }
  if (values.totalTokens !== null && values.inputTokens !== null && values.outputTokens !== null
    && values.totalTokens !== values.inputTokens + values.outputTokens && !inconsistent.includes("totalTokens")) {
    inconsistent.push("totalTokens");
  }
  const missing = (Object.keys(values) as ProviderUsageField[]).filter((field) => values[field] === null);
  const status: ProviderUsageStatus = inconsistent.length > 0
    ? "INCONSISTENT"
    : missing.length === 3
      ? "MISSING"
      : missing.length > 0
        ? "PARTIAL"
        : "AVAILABLE";
  return { status, ...values, missing, inconsistent, derivedTotal, sources };
}

export const normalizeProviderUsage = parseProviderUsage;
export const parseUsage = parseProviderUsage;