import { parseProviderUsage } from "./provider-usage";

export type TokenLimitCategory = "NONE" | "CONTEXT_LIMIT" | "OUTPUT_LIMIT" | "MAX_OUTPUT_TOKENS" | "UNKNOWN";

export interface TokenLimitClassification {
  category: TokenLimitCategory;
  detected: boolean;
  retryable: boolean;
  reason: string | null;
  actual: number | null;
  limit: number | null;
}

export interface TokenLimitOptions {
  contextTokens?: number;
  outputTokens?: number;
  maxOutputTokens?: number;
}

function textOf(input: unknown): string {
  try {
    const serialized = JSON.stringify(input);
    return (serialized ?? String(input)).toLowerCase();
  } catch { return String(input).toLowerCase(); }
}

export const classifyLimit = classifyTokenLimit;

function numberAt(input: unknown, names: Set<string>, seen = new Set<object>()): number | null {
  if (typeof input !== "object" || input === null) return null;
  if (seen.has(input)) return null;
  seen.add(input);
  if (Array.isArray(input)) {
    for (const value of input) { const result = numberAt(value, names, seen); if (result !== null) return result; }
    return null;
  }
  for (const [key, value] of Object.entries(input)) {
    const normalized = key.replace(/[A-Z]/gu, (letter) => `_${letter.toLowerCase()}`).toLowerCase();
    if (names.has(normalized) && typeof value === "number" && Number.isFinite(value)) return value;
    const result = numberAt(value, names, seen);
    if (result !== null) return result;
  }
  return null;
}

function result(category: TokenLimitCategory, reason: string | null, actual: number | null, limit: number | null): TokenLimitClassification {
  return { category, detected: category !== "NONE", retryable: category === "CONTEXT_LIMIT" || category === "OUTPUT_LIMIT", reason, actual, limit };
}

/** Recognize provider errors, finish reasons and measured usage consistently. */
export function classifyTokenLimit(input: unknown, options: TokenLimitOptions = {}): TokenLimitClassification {
  const text = textOf(input);
  const usage = parseProviderUsage(input);
  const contextActual = numberAt(input, new Set(["context_tokens", "contexttokens", "prompt_tokens", "prompttokens", "input_tokens", "inputtokens"]));
  const outputActual = usage.outputTokens ?? numberAt(input, new Set(["output_tokens", "outputtokens", "completion_tokens", "completiontokens"]));
  const contextMarker = /(context[_ -]?length|context[_ -]?window|maximum[_ -]?context|prompt[_ -]?too[_ -]?long|input[_ -]?too[_ -]?long|context[_ -]?limit|context[_ -]?exceed)/u.test(text);
  const outputMarker = /(output[_ -]?limit|completion[_ -]?limit|output[_ -]?too[_ -]?long|output[_ -]?exceed|completion[_ -]?tokens?.{0,20}(limit|exceed|too))/u.test(text);
  const maxMarker = /(max(?:imum)?[_ -]?output[_ -]?tokens|max[_ -]?completion[_ -]?tokens|max[_ -]?tokens|finish[_ -]?reason.{0,12}(length|max_tokens)|stop[_ -]?reason.{0,12}max_tokens|incomplete.{0,30}max_output_tokens)/u.test(text);

  if (contextMarker || (options.contextTokens !== undefined && contextActual !== null && contextActual > options.contextTokens)) {
    return result("CONTEXT_LIMIT", "context token limit exceeded", contextActual, options.contextTokens ?? null);
  }
  if (outputMarker || (options.outputTokens !== undefined && outputActual !== null && outputActual > options.outputTokens)) {
    return result("OUTPUT_LIMIT", "output token limit exceeded", outputActual, options.outputTokens ?? null);
  }
  if (maxMarker || (options.maxOutputTokens !== undefined && outputActual !== null && outputActual >= options.maxOutputTokens && /(?:length|max[_ -]?tokens)/u.test(text))) {
    return result("MAX_OUTPUT_TOKENS", "generation stopped at max_output_tokens", outputActual, options.maxOutputTokens ?? null);
  }
  if (text.includes("token") && /(limit|exceed|truncat|too long)/u.test(text)) return result("UNKNOWN", "provider reported an unclassified token limit", null, null);
  return result("NONE", null, null, null);
}