import type { Message, ToolDef } from '../../llm/openai-compat.js';

/**
 * Conservative prompt-size estimation for local preflight guards.
 *
 * A single "chars * N" multiplier is too crude here: ASCII-heavy tool output,
 * JSON, and TypeScript source are commonly 3–5 chars/token, while Japanese can
 * be close to 1 char/token. Treating every character as 1.5 tokens caused
 * local guards to report token counts an order of magnitude above provider
 * `usage.prompt_tokens` for normal code/log-heavy tasks.
 */
export const UNKNOWN_CHARS_TO_TOKENS_FACTOR = 1.2;
const ASCII_CHARS_PER_TOKEN = 3.5;
const CJK_TOKENS_PER_CHAR = 1.2;
const OTHER_TOKENS_PER_CHAR = 1.0;
export const IMAGE_CONTENT_TOKENS = 1024;

export function estimateTokensFromChars(chars: number): number {
  return Math.ceil(chars * UNKNOWN_CHARS_TO_TOKENS_FACTOR);
}

export function estimateTokensFromText(text: string): number {
  let asciiChars = 0;
  let cjkChars = 0;
  let otherChars = 0;

  for (const char of text) {
    const code = char.codePointAt(0) ?? 0;
    if (code <= 0x7f) {
      asciiChars++;
    } else if (
      (code >= 0x3040 && code <= 0x30ff) || // Hiragana + Katakana
      (code >= 0x3400 && code <= 0x9fff) || // CJK ideographs
      (code >= 0xf900 && code <= 0xfaff) || // CJK compatibility ideographs
      (code >= 0xff00 && code <= 0xffef)    // full-width forms
    ) {
      cjkChars++;
    } else {
      otherChars++;
    }
  }

  return Math.ceil(
    (asciiChars / ASCII_CHARS_PER_TOKEN) +
    (cjkChars * CJK_TOKENS_PER_CHAR) +
    (otherChars * OTHER_TOKENS_PER_CHAR),
  );
}

export function estimateMessageTokens(message: Message): number {
  let tokens = estimateTokensFromText(message.role) + 8;
  if (typeof message.content === 'string') {
    tokens += estimateTokensFromText(message.content);
  } else if (Array.isArray(message.content)) {
    for (const part of message.content) {
      if (part.type === 'text') {
        tokens += estimateTokensFromText(part.text);
      } else {
        tokens += IMAGE_CONTENT_TOKENS;
      }
    }
  }
  if (message.tool_call_id) tokens += estimateTokensFromText(message.tool_call_id);
  if (message.name) tokens += estimateTokensFromText(message.name);
  if (message.tool_calls) {
    for (const toolCall of message.tool_calls) {
      tokens += estimateTokensFromText(toolCall.id);
      tokens += estimateTokensFromText(toolCall.function.name);
      tokens += estimateTokensFromText(toolCall.function.arguments);
    }
  }
  return Math.ceil(tokens);
}

export function estimateMessagesTokens(messages: Message[]): number {
  let total = 0;
  for (const message of messages) total += estimateMessageTokens(message);
  return total;
}

/**
 * Token cost of a tool definition list. Caller should cache the result for
 * the lifetime of a movement — `tools` does not change once built.
 */
export function estimateToolsTokens(tools: ToolDef[]): number {
  return estimateTokensFromText(JSON.stringify(tools));
}

export function estimatePromptTokens(messages: Message[], tools: ToolDef[]): number {
  return estimateMessagesTokens(messages) + estimateToolsTokens(tools);
}
