"""Utilities for token-budget-aware translation text splitting."""
  
  from __future__ import annotations
  
  from typing import Callable, List, Optional
  
  TokenLengthFn = Callable[[str], int]
  
  _CJK_LANGS = {"zh", "ja", "ko"}
  _STRONG_BOUNDARIES = {"\n", "。", "！", "？", "!", "?", "；", ";", "…"}
  _WEAK_BOUNDARIES = {"，", ",", "、", "：", ":", "(", ")", "（", "）", "[", "]", "【", "】", "/", "|"}
  _CLOSING_CHARS = {'"', "'", "”", "’", ")", "]", "}", "）", "】", "》", "」", "』"}
  _NO_SPACE_BEFORE = tuple('.,!?;:)]}%>"\'')
  _NO_SPACE_AFTER = tuple("([{$#@/<")
  
  
  def is_cjk_language(lang: Optional[str]) -> bool:
      return str(lang or "").strip().lower() in _CJK_LANGS
  
  
  def compute_safe_input_token_limit(
      *,
      max_input_length: int,
      max_new_tokens: int,
      decoding_length_mode: str = "fixed",
      decoding_length_extra: int = 0,
      reserve_input_tokens: int = 8,
      reserve_output_tokens: int = 8,
  ) -> int:
      """Derive a conservative source-token budget for translation splitting.
  
      We keep a small reserve for tokenizer special tokens on the input side. If
      the decode side is much tighter than the encode side, we also cap the
      source budget based on decode settings so we split before the model is
      likely to truncate.
      """
  
      input_limit = max(8, int(max_input_length) - max(0, int(reserve_input_tokens)))
      decode_mode = str(decoding_length_mode or "fixed").strip().lower()
      if int(max_new_tokens) <= 0:
          return input_limit
      if decode_mode == "source":
          output_limit = max(8, int(max_new_tokens) - max(0, int(decoding_length_extra)))
          return max(8, min(input_limit, output_limit))
      if int(max_new_tokens) >= int(max_input_length):
          return input_limit
      output_limit = max(8, int(max_new_tokens) - max(0, int(reserve_output_tokens)))
      return max(8, min(input_limit, output_limit))
  
  
  def split_text_for_translation(
      text: str,
      *,
      max_tokens: int,
      token_length_fn: TokenLengthFn,
  ) -> List[str]:
      """Split long text into a few translation-friendly segments.
  
      The splitter prefers sentence boundaries, then clause boundaries, then
      whitespace, and only falls back to character-based splitting when needed.
      """
  
      if not text:
          return [text]
      if token_length_fn(text) <= max_tokens:
          return [text]
      segments = _split_recursive(text, max_tokens=max_tokens, token_length_fn=token_length_fn, level=0)
      return [segment for segment in segments if segment]
  
  
  def join_translated_segments(
      segments: List[Optional[str]],
      *,
      target_lang: Optional[str],
      original_text: str,
  ) -> Optional[str]:
      parts = [segment.strip() for segment in segments if segment and segment.strip()]
      if not parts:
          return None
      separator = "" if is_cjk_language(target_lang) else " "
      if "\n" in original_text and separator:
          separator = "\n"
  
      merged = parts[0]
      for part in parts[1:]:
          if not separator:
              merged += part
              continue
          if merged.endswith(_NO_SPACE_AFTER) or part.startswith(_NO_SPACE_BEFORE):
              merged += part
              continue
          merged += separator + part
      return merged.strip() or None
  
  
  def _split_recursive(
      text: str,
      *,
      max_tokens: int,
      token_length_fn: TokenLengthFn,
      level: int,
  ) -> List[str]:
      if token_length_fn(text) <= max_tokens:
          return [text]
      if level >= 3:
          return _hard_split(text, max_tokens=max_tokens, token_length_fn=token_length_fn)
  
      pieces = _split_by_level(text, level)
      if len(pieces) <= 1:
          return _split_recursive(text, max_tokens=max_tokens, token_length_fn=token_length_fn, level=level + 1)
  
      merged: List[str] = []
      buffer = ""
      for piece in pieces:
          candidate = buffer + piece if buffer else piece
          if token_length_fn(candidate) <= max_tokens:
              buffer = candidate
              continue
          if buffer:
              merged.extend(
                  _split_recursive(buffer, max_tokens=max_tokens, token_length_fn=token_length_fn, level=level + 1)
              )
              buffer = piece
              continue
          merged.extend(_split_recursive(piece, max_tokens=max_tokens, token_length_fn=token_length_fn, level=level + 1))
      if buffer:
          merged.extend(_split_recursive(buffer, max_tokens=max_tokens, token_length_fn=token_length_fn, level=level + 1))
      return merged
  
  
  def _split_by_level(text: str, level: int) -> List[str]:
      parts: List[str] = []
      start = 0
      index = 0
      while index < len(text):
          boundary_end = _match_boundary(text, index, level)
          if boundary_end is None:
              index += 1
              continue
          if boundary_end > start:
              parts.append(text[start:boundary_end])
              start = boundary_end
          index = boundary_end
      if start < len(text):
          parts.append(text[start:])
      return [part for part in parts if part]
  
  
  def _match_boundary(text: str, index: int, level: int) -> Optional[int]:
      char = text[index]
      if level == 0:
          if char in _STRONG_BOUNDARIES:
              return _consume_boundary_tail(text, index + 1)
          if char == "." and _is_sentence_period(text, index):
              return _consume_boundary_tail(text, index + 1)
          return None
      if level == 1:
          if char in _WEAK_BOUNDARIES:
              return _consume_boundary_tail(text, index + 1)
          return None
      if level == 2 and char.isspace():
          end = index + 1
          while end < len(text) and text[end].isspace():
              end += 1
          return end
      return None
  
  
  def _consume_boundary_tail(text: str, index: int) -> int:
      end = index
      while end < len(text) and text[end] in _CLOSING_CHARS:
          end += 1
      while end < len(text) and text[end].isspace():
          end += 1
      return end
  
  
  def _is_sentence_period(text: str, index: int) -> bool:
      prev_char = text[index - 1] if index > 0 else ""
      next_char = text[index + 1] if index + 1 < len(text) else ""
      if prev_char.isdigit() and next_char.isdigit():
          return False
      if not next_char:
          return True
      return next_char.isspace() or next_char in _CLOSING_CHARS
  
  
  def _hard_split(text: str, *, max_tokens: int, token_length_fn: TokenLengthFn) -> List[str]:
      segments: List[str] = []
      remaining = text
      while remaining:
          if token_length_fn(remaining) <= max_tokens:
              segments.append(remaining)
              break
          cut = _largest_prefix_within_limit(remaining, max_tokens=max_tokens, token_length_fn=token_length_fn)
          refined_cut = _refine_cut(remaining, cut, max_tokens=max_tokens, token_length_fn=token_length_fn)
          if refined_cut <= 0:
              refined_cut = max(1, cut)
          segments.append(remaining[:refined_cut])
          remaining = remaining[refined_cut:]
      return segments
  
  
  def _largest_prefix_within_limit(text: str, *, max_tokens: int, token_length_fn: TokenLengthFn) -> int:
      low = 1
      high = len(text)
      best = 1
      while low <= high:
          mid = (low + high) // 2
          if token_length_fn(text[:mid]) <= max_tokens:
              best = mid
              low = mid + 1
              continue
          high = mid - 1
      return best
  
  
  def _refine_cut(text: str, cut: int, *, max_tokens: int, token_length_fn: TokenLengthFn) -> int:
      best = cut
      lower_bound = max(1, cut - 32)
      for candidate in range(cut, lower_bound - 1, -1):
          if text[candidate - 1].isspace() or text[candidate - 1] in _STRONG_BOUNDARIES or text[candidate - 1] in _WEAK_BOUNDARIES:
              if candidate >= max(1, cut // 2) and token_length_fn(text[:candidate]) <= max_tokens:
                  return candidate
              best = max(best, candidate)
      return best