///|
/// Heuristic token estimation and prompt utilities.
///
/// These do not implement a real BPE tokenizer; they provide fast, dependency
/// -free approximations suitable for budgeting, truncation, and cost display.
/// For exact counts, use a model-specific tokenizer offline.

///|
/// A family of estimation heuristics with different characters-per-token
/// ratios, since different model families tokenize at different densities.
pub(all) enum TokenizerHint {
  /// ~4 chars/token — a good default for GPT-style English text.
  GptLike
  /// ~3.5 chars/token — Claude tends to produce slightly more tokens.
  ClaudeLike
  /// A custom characters-per-token ratio (must be > 0).
  CharsPerToken(Double)
} derive(Eq, Debug)

///|
/// The characters-per-token ratio for a hint.
fn TokenizerHint::ratio(self : TokenizerHint) -> Double {
  match self {
    GptLike => 4.0
    ClaudeLike => 3.5
    CharsPerToken(r) => if r > 0.0 { r } else { 4.0 }
  }
}

///|
/// Estimate the token count of `text` under a given hint.
pub fn estimate_tokens_with(text : String, hint : TokenizerHint) -> Int {
  let chars = text.length()
  if chars == 0 {
    return 0
  }
  let ratio = hint.ratio()
  let est = (chars.to_double() / ratio).ceil().to_int()
  if est < 1 {
    1
  } else {
    est
  }
}

///|
/// A word-boundary based estimate: counts whitespace-separated words plus a
/// surcharge for punctuation-heavy text. Sometimes closer than the pure
/// char-ratio for natural language.
pub fn estimate_tokens_by_words(text : String) -> Int {
  let mut words = 0
  let mut in_word = false
  for c in text {
    let is_space = c == ' ' || c == '\n' || c == '\t' || c == '\r'
    if is_space {
      in_word = false
    } else if !in_word {
      words = words + 1
      in_word = true
    }
  }
  // Empirically, tokens ≈ words / 0.75 for English.
  (words.to_double() / 0.75).ceil().to_int()
}

///|
/// Estimate the total prompt tokens for a list of messages under a hint,
/// including a per-message overhead (role + formatting).
pub fn estimate_messages_tokens(
  messages : Array[Message],
  hint? : TokenizerHint = GptLike,
) -> Int {
  let mut total = 0
  for m in messages {
    total = total + estimate_tokens_with(m.content.to_text(), hint) + 4
  }
  // A few tokens of reply priming, per OpenAI's counting guidance.
  total + 3
}

///|
/// Approximate a US-dollar cost given token counts and per-1K-token prices.
pub fn estimate_cost(
  prompt_tokens : Int,
  completion_tokens : Int,
  prompt_price_per_1k~ : Double,
  completion_price_per_1k~ : Double,
) -> Double {
  prompt_tokens.to_double() / 1000.0 * prompt_price_per_1k +
  completion_tokens.to_double() / 1000.0 * completion_price_per_1k
}

///|
/// Truncate `text` to at most `max_tokens` (estimated), returning the kept
/// prefix. Useful for clamping a single oversized message.
pub fn truncate_to_tokens(
  text : String,
  max_tokens : Int,
  hint? : TokenizerHint = GptLike,
) -> String {
  if estimate_tokens_with(text, hint) <= max_tokens {
    return text
  }
  let max_chars = (max_tokens.to_double() * hint.ratio()).to_int()
  if max_chars >= text.length() {
    text
  } else if max_chars <= 0 {
    ""
  } else {
    text[:max_chars].to_owned()
  }
}

///|
/// Render a list of messages into a single plain-text transcript, one line
/// per message prefixed by the role — handy for logging or debugging.
pub fn render_transcript(messages : Array[Message]) -> String {
  let out = StringBuilder::new()
  for i, m in messages {
    if i > 0 {
      out.write_string("\n")
    }
    out.write_string(m.role.to_string())
    out.write_string(": ")
    out.write_string(m.content.to_text())
  }
  out.to_string()
}

///|
/// Merge consecutive messages from the same role into a single message,
/// joining their text with a blank line. Some APIs reject two user messages
/// in a row; this collapses them.
pub fn merge_consecutive(messages : Array[Message]) -> Array[Message] {
  let out : Array[Message] = []
  for m in messages {
    match out.last() {
      Some(prev) if prev.role == m.role && prev.tool_calls is None =>
        // Replace the previous entry with a merged one.
        out[out.length() - 1] = {
          role: prev.role,
          content: Str(prev.content.to_text() + "\n\n" + m.content.to_text()),
          tool_calls: None,
          tool_call_id: None,
          name: None,
        }
      _ => out.push(m)
    }
  }
  out
}