"""
  Shared tokenization helpers for query understanding.
  """
  
  from __future__ import annotations
  
  from dataclasses import dataclass
  import re

  from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple

  _HAN_PATTERN = re.compile(r"[\u4e00-\u9fff]")
  _TOKEN_PATTERN = re.compile(r"[\u4e00-\u9fff]+|[^\W_]+(?:[-'][^\W_]+)*", re.UNICODE)

  
  
  def normalize_query_text(text: Optional[str]) -> str:
      if text is None:
          return ""
      return " ".join(str(text).strip().casefold().split())
  
  
  def simple_tokenize_query(text: str) -> List[str]:
      """
      Lightweight tokenizer for coarse query matching.
  
      - Consecutive CJK characters form one token
      - Latin / digit runs (with internal hyphens) form tokens
      """
      if not text:
          return []
      return _TOKEN_PATTERN.findall(text)
  
  

  def contains_han_text(text: Optional[str]) -> bool:
      return bool(text and _HAN_PATTERN.search(str(text)))
  
  

  def extract_token_strings(tokenizer_result: Any) -> List[str]:
      """Normalize tokenizer output into a flat token string list."""
      if not tokenizer_result:
          return []
      if isinstance(tokenizer_result, str):
          token = tokenizer_result.strip()
          return [token] if token else []
  
      tokens: List[str] = []
      for item in tokenizer_result:
          token: Optional[str] = None
          if isinstance(item, str):
              token = item
          elif isinstance(item, (list, tuple)) and item:
              token = str(item[0])
          elif item is not None:
              token = str(item)
  
          if token is None:
              continue
          token = token.strip()
          if token:
              tokens.append(token)
      return tokens
  
  
  def _dedupe_preserve_order(values: Iterable[str]) -> List[str]:
      result: List[str] = []
      seen = set()
      for value in values:
          normalized = normalize_query_text(value)
          if not normalized or normalized in seen:
              continue
          seen.add(normalized)
          result.append(normalized)
      return result
  
  
  def _build_phrase_candidates(tokens: Sequence[str], max_ngram: int) -> List[str]:
      if not tokens:
          return []
  
      phrases: List[str] = []
      upper = max(1, int(max_ngram))
      for size in range(1, upper + 1):
          if size > len(tokens):
              break
          for start in range(0, len(tokens) - size + 1):
              phrase = " ".join(tokens[start:start + size]).strip()
              if phrase:
                  phrases.append(phrase)
      return phrases
  
  

  def _build_coarse_tokens(
      text: str,
      *,
      language_hint: Optional[str],
      tokenizer_tokens: Sequence[str],
  ) -> List[str]:
      normalized_language = normalize_query_text(language_hint)
      if normalized_language == "zh" or (contains_han_text(text) and tokenizer_tokens):
          # Chinese coarse tokenization should follow the model tokenizer rather than a
          # regex that collapses the whole sentence into one CJK span.
          return list(_dedupe_preserve_order(tokenizer_tokens))
      return _dedupe_preserve_order(simple_tokenize_query(text))

  @dataclass(frozen=True)
  class TokenizedText:
      text: str
      normalized_text: str
      fine_tokens: Tuple[str, ...]
      coarse_tokens: Tuple[str, ...]
      candidates: Tuple[str, ...]
  
  

  class QueryTextAnalysisCache:
      """Per-parse cache for tokenizer output and derived token bundles."""
  
      def __init__(self, *, tokenizer: Optional[Callable[[str], Any]] = None) -> None:
          self.tokenizer = tokenizer
          self._tokenizer_results: Dict[str, Any] = {}
          self._tokenized_texts: Dict[Tuple[str, int], TokenizedText] = {}
          self._language_hints: Dict[str, str] = {}
  
      @staticmethod
      def _normalize_input(text: Optional[str]) -> str:
          return str(text or "").strip()
  
      def set_language_hint(self, text: Optional[str], language: Optional[str]) -> None:
          normalized_input = self._normalize_input(text)
          normalized_language = normalize_query_text(language)
          if normalized_input and normalized_language:
              self._language_hints[normalized_input] = normalized_language
  
      def get_language_hint(self, text: Optional[str]) -> Optional[str]:
          normalized_input = self._normalize_input(text)
          if not normalized_input:
              return None
          return self._language_hints.get(normalized_input)
  
      def _should_use_model_tokenizer(self, text: str) -> bool:
          if self.tokenizer is None:
              return False
          language_hint = self.get_language_hint(text)
          has_han = contains_han_text(text)
          if language_hint == "zh":
              return has_han
          return has_han
  
      def get_tokenizer_result(self, text: Optional[str]) -> Any:
          normalized_input = self._normalize_input(text)
          if not normalized_input:
              return []
          if not self._should_use_model_tokenizer(normalized_input):
              return simple_tokenize_query(normalized_input)
          if normalized_input not in self._tokenizer_results:
              self._tokenizer_results[normalized_input] = self.tokenizer(normalized_input)
          return self._tokenizer_results[normalized_input]
  
      def get_tokenized_text(self, text: Optional[str], *, max_ngram: int = 3) -> TokenizedText:
          normalized_input = self._normalize_input(text)
          cache_key = (normalized_input, max(1, int(max_ngram)))
          cached = self._tokenized_texts.get(cache_key)
          if cached is not None:
              return cached
  
          normalized_text = normalize_query_text(normalized_input)
          fine_raw = extract_token_strings(self.get_tokenizer_result(normalized_input))
          fine_tokens = _dedupe_preserve_order(fine_raw)

          coarse_tokens = _build_coarse_tokens(
              normalized_input,
              language_hint=self.get_language_hint(normalized_input),
              tokenizer_tokens=fine_tokens,
          )

  
          bundle = TokenizedText(
              text=normalized_input,
              normalized_text=normalized_text,
              fine_tokens=tuple(fine_tokens),
              coarse_tokens=tuple(coarse_tokens),
              candidates=tuple(
                  _dedupe_preserve_order(
                      list(fine_tokens)
                      + list(coarse_tokens)
                      + _build_phrase_candidates(fine_tokens, max_ngram=max_ngram)
                      + _build_phrase_candidates(coarse_tokens, max_ngram=max_ngram)
                      + ([normalized_text] if normalized_text else [])
                  )
              ),
          )
          self._tokenized_texts[cache_key] = bundle
          return bundle
  
  

  def tokenize_text(
      text: str,
      *,
      tokenizer: Optional[Callable[[str], Any]] = None,
      max_ngram: int = 3,
  ) -> TokenizedText:

      return QueryTextAnalysisCache(tokenizer=tokenizer).get_tokenized_text(
          text,
          max_ngram=max_ngram,

      )