Netcatty/lib/searchMatcher.ts

import { pinyin } from "pinyin-pro";

const SEARCH_SPLIT_REGEX = /[\s\p{Pd}_/\\|.,，。;；:：!！?？()（）[\]{}<>《》、"'`~·]+/u;
const SEARCH_REMOVE_REGEX = /[\s\p{Pd}_/\\|.,，。;；:：!！?？()（）[\]{}<>《》、"'`~·]+/gu;
const SEARCH_QUERY_SEGMENT_SPLIT_REGEX = /\s+/u;
const SEARCH_QUERY_PUNCT_REGEX = /[\p{Pd}_/\\|.,，。;；:：!！?？()（）[\]{}<>《》、"'`~·]/u;
const DASH_SEPARATOR_REGEX = /[\u2010-\u2015\u2212\uFE58\uFE63\uFF0D]/gu;
export const PINYIN_CACHE_MAX_ENTRIES = 16_384;
/**
 * Hard cap on the total characters the pinyin cache retains (cache key +
 * full + initials transliterations per entry, i.e. several copies of the
 * source text). A per-entry cap alone is not enough: callers transliterate
 * long fields (e.g. up to 20,000-character session titles), so a scan over
 * many distinct large fields could retain tens or hundreds of megabytes
 * within the entry count. Sources whose transliterations alone exceed this
 * budget are not cached at all — they are simply re-transliterated on the
 * next lookup instead of pinning memory.
 */
export const PINYIN_CACHE_MAX_CHARS = 1_000_000;
const PINYIN_CACHE = new Map<string, { full: string; initials: string }>();
let pinyinCacheChars = 0;
const IPV4_LIKE_REGEX = /^\d{1,3}(?:\.\d{1,3})+$/;
const HOST_PHASE_SCORE_BONUS = {
  strict: 1_000_000,
  loose: 100_000,
} as const;
const HOST_FIELD_SCORE = {
  label: 500,
  hostname: 380,
  group: 280,
  tag: 220,
} as const;
const HOST_METHOD_SCORE = {
  literal: 50,
  compact: 25,
  pinyinFull: 12,
  pinyinInitials: 8,
} as const;

function normalizeText(input: string): string {
  return input.normalize("NFKC").replace(DASH_SEPARATOR_REGEX, "-").toLowerCase().trim();
}

function compactText(input: string): string {
  return normalizeText(input).replace(SEARCH_REMOVE_REGEX, "");
}

function pinyinCacheEntryChars(
  key: string,
  entry: { full: string; initials: string },
): number {
  return key.length + entry.full.length + entry.initials.length;
}

function getPinyinVariants(sourceText: string): { full: string; initials: string } {
  const cacheKey = normalizeText(sourceText);
  const cached = PINYIN_CACHE.get(cacheKey);
  if (cached) {
    PINYIN_CACHE.delete(cacheKey);
    PINYIN_CACHE.set(cacheKey, cached);
    return cached;
  }

  let full = "";
  let initials = "";

  try {
    full = compactText(
      pinyin(sourceText, {
        toneType: "none",
      }),
    );
    initials = compactText(
      pinyin(sourceText, {
        pattern: "first",
        toneType: "none",
      }),
    );
  } catch {
    // ignore conversion failures
  }

  const next = { full, initials };
  const entryChars = pinyinCacheEntryChars(cacheKey, next);
  // Never cache an entry that by itself exceeds the retained-char budget:
  // admitting it (and evicting everything older to fit) would still let one
  // oversized field pin memory above the budget, and skipping keeps it a
  // simple re-transliteration on the next lookup instead of unbounded decay.
  if (entryChars > PINYIN_CACHE_MAX_CHARS) return next;
  pinyinCacheChars += entryChars;
  PINYIN_CACHE.delete(cacheKey);
  PINYIN_CACHE.set(cacheKey, next);
  while (PINYIN_CACHE.size > PINYIN_CACHE_MAX_ENTRIES
    || pinyinCacheChars > PINYIN_CACHE_MAX_CHARS) {
    const oldestKey = PINYIN_CACHE.keys().next().value as string | undefined;
    if (oldestKey === undefined) break;
    const oldest = PINYIN_CACHE.get(oldestKey);
    if (!oldest) break;
    pinyinCacheChars -= pinyinCacheEntryChars(oldestKey, oldest);
    PINYIN_CACHE.delete(oldestKey);
  }
  return next;
}

export function resetPinyinCacheForTests(): void {
  PINYIN_CACHE.clear();
  pinyinCacheChars = 0;
}

export function getPinyinCacheStatsForTests(): {
  size: number;
  chars: number;
  keys: string[];
} {
  return {
    size: PINYIN_CACHE.size,
    chars: pinyinCacheChars,
    keys: [...PINYIN_CACHE.keys()],
  };
}

export function tokenizeSearchQuery(query: string): string[] {
  const normalized = normalizeText(query);
  if (!normalized) return [];
  return normalized.split(SEARCH_SPLIT_REGEX).filter(Boolean);
}

/**
 * Optional matching tweaks. Callers with large unbounded haystacks (e.g.
 * full conversation histories) can keep literal/compact matching global
 * while restricting the (expensive) pinyin fallback to specific fields via
 * `pinyinFields`: each token may still match literally anywhere in the
 * haystack, but is only transliteration-matched against the listed fields.
 */
export type SearchMatchOptions = {
  allowPinyin?: boolean;
  /** Restrict the pinyin fallback to these fields; defaults to every field. */
  pinyinFields?: string[];
};

/**
 * Suggested query bound for callers that rescan many haystacks per keystroke
 * (e.g. session-history search). The shared matcher never truncates: callers
 * apply this cap at their own input so the bound stays local to them and the
 * global match semantics (full-query meaning) are unchanged for everyone else.
 */
export const MAX_SEARCH_QUERY_LENGTH = 256;

/**
 * Bound a caller-side search query to `MAX_SEARCH_QUERY_LENGTH` UTF-16 code
 * units without splitting a UTF-16 surrogate pair when the cap lands mid-pair
 * (e.g. an astral-plane character such as an emoji; drop the lone high
 * surrogate). Truncation is a caller concern: see `MAX_SEARCH_QUERY_LENGTH`.
 */
export function truncateQueryForMatch(query: string): string {
  if (query.length <= MAX_SEARCH_QUERY_LENGTH) return query;
  const truncated = query.slice(0, MAX_SEARCH_QUERY_LENGTH);
  // Avoid splitting a UTF-16 surrogate pair when the cap lands mid-pair (e.g.
  // an astral-plane character such as an emoji); drop the lone high surrogate.
  const tail = truncated.charCodeAt(truncated.length - 1);
  return tail >= 0xd800 && tail <= 0xdbff ? truncated.slice(0, -1) : truncated;
}

/**
 * Pre-normalized search query. Callers that match the same query against many
 * haystacks (e.g. session-history search scans every session per keystroke)
 * can prepare the query once per scan instead of paying NFKC normalization,
 * tokenization and compaction once per session.
 */
export type PreparedSearchQuery = {
  /** Normalized (NFKC, dash-folded, lowercased, trimmed) query. */
  normalized: string;
  /** Non-empty tokens of `normalized` after separator splitting. */
  tokens: string[];
  /** Separator-removed form of `normalized` (same as `compactText(normalized)`). */
  compact: string;
  /** `normalized` looks like a dotted IPv4 address (contiguous-literal mode). */
  ipv4Like: boolean;
};

const EMPTY_PREPARED_QUERY: PreparedSearchQuery = {
  normalized: "",
  tokens: [],
  compact: "",
  ipv4Like: false,
};

export function prepareSearchQuery(query: string): PreparedSearchQuery {
  // No length cap here: truncating globally would silently change match
  // semantics for every caller (an over-limit query's tail would be dropped,
  // so a field containing only the prefix would report a match). Callers that
  // rescan many haystacks bound their own input via `truncateQueryForMatch`.
  const normalized = normalizeText(query);
  if (!normalized) return EMPTY_PREPARED_QUERY;
  return {
    normalized,
    tokens: normalized.split(SEARCH_SPLIT_REGEX).filter(Boolean),
    compact: compactText(normalized),
    ipv4Like: IPV4_LIKE_REGEX.test(normalized),
  };
}

/**
 * Pre-normalized search haystack. Callers that match the same fields against
 * many queries (e.g. session-history search fires on every keystroke) can
 * prepare the haystack once, cache it, and reuse it per query instead of
 * re-running NFKC normalization and compaction over megabytes each time.
 */
export type PreparedSearchFields = {
  /** Per-field normalized text, in input order. */
  normalizedFields: string[];
  /** Normalized fields joined with a single space. */
  haystack: string;
  /** Separator-removed form of `haystack` (same as `compactText(haystack)`). */
  haystackCompact: string;
};

export function prepareSearchFields(
  fields: Array<string | null | undefined>,
): PreparedSearchFields {
  const normalizedFields = fields
    .filter((field): field is string => typeof field === "string" && field.trim().length > 0)
    .map((field) => normalizeText(field));
  const haystack = normalizedFields.join(" ");
  return { normalizedFields, haystack, haystackCompact: compactText(haystack) };
}

export function matchesPreparedSearchQuery(
  query: string | PreparedSearchQuery,
  prepared: PreparedSearchFields,
  options: SearchMatchOptions = {},
): boolean {
  // Raw strings are still accepted (and prepared inline) for one-off callers;
  // repeated-match callers pass `prepareSearchQuery(query)` instead.
  const { normalized: normalizedQuery, tokens, compact: compactQuery, ipv4Like } =
    typeof query === "string" ? prepareSearchQuery(query) : query;
  if (!normalizedQuery) return true;

  if (prepared.normalizedFields.length === 0) return false;

  // For dotted numeric input (IPv4-like), require contiguous literal match.
  if (ipv4Like) {
    return prepared.normalizedFields.some((field) => field.includes(normalizedQuery));
  }

  if (tokens.length === 0) return false;

  const haystack = prepared.haystack;
  if (haystack.includes(normalizedQuery)) {
    return true;
  }

  if (compactQuery && prepared.haystackCompact.includes(compactQuery)) {
    return true;
  }

  if (tokens.every((token) => haystack.includes(token))) {
    return true;
  }

  const hasLatinToken = tokens.some((token) => /[a-z]/i.test(token));
  if (!hasLatinToken || options.allowPinyin === false) return false;

  const pinyinSource = options.pinyinFields === undefined
    ? haystack
    : options.pinyinFields
      .filter((field): field is string => typeof field === "string" && field.trim().length > 0)
      .map((field) => normalizeText(field))
      .join(" ");
  const { full, initials } = getPinyinVariants(pinyinSource);
  if (!full && !initials) return false;

  return tokens.every((token) => {
    if (haystack.includes(token)) return true;
    const compactToken = compactText(token);
    return (
      (full && full.includes(compactToken)) ||
      (initials && initials.includes(compactToken))
    );
  });
}

const DEFAULT_SEARCH_MATCH_OPTIONS: Required<Pick<SearchMatchOptions, "allowPinyin">> = {
  allowPinyin: true,
};

export function matchesSearchQuery(
  query: string,
  ...rest: Array<string | null | undefined | SearchMatchOptions>
): boolean {
  const args = [...rest];
  const options: Pick<SearchMatchOptions, "allowPinyin"> & { pinyinFields?: string[] } =
    args.length > 0
    && typeof args[args.length - 1] === "object"
    && args[args.length - 1] !== null
    ? { ...DEFAULT_SEARCH_MATCH_OPTIONS, ...(args.pop() as SearchMatchOptions) }
    : DEFAULT_SEARCH_MATCH_OPTIONS;
  const fields = args as Array<string | null | undefined>;
  return matchesPreparedSearchQuery(
    query,
    prepareSearchFields(fields),
    { allowPinyin: options.allowPinyin, pinyinFields: options.pinyinFields },
  );
}

/**
 * Word-prefix matcher: the query must prefix-match a whole token (word) of one
 * of the fields — including that token's pinyin full/initials forms. Strictly
 * weaker than matchesSearchQuery's substring behavior; used where substring
 * noise is undesirable, e.g. snippet completions while typing an argument
 * (issue #3622). Dashes/underscores separate tokens, so "set-email" still
 * matches from the word "email".
 */
export function matchesWordPrefixQuery(
  query: string,
  ...fields: Array<string | null | undefined>
): boolean {
  const normalizedQuery = normalizeText(query);
  if (!normalizedQuery) return true;

  const words = fields
    .filter((field): field is string => typeof field === "string" && field.trim().length > 0)
    .flatMap((field) => normalizeText(field).split(SEARCH_SPLIT_REGEX))
    .filter((word) => word.length > 0);
  if (words.length === 0) return false;

  // The compact/pinyin fallback must not discard the query's leading
  // punctuation: option-style needles like "--ve" or "-d" would otherwise
  // match words ("version", "deploy") that never literally start with them
  // and displace the real Fig option completions while typing flags.
  const compactQuery = /^[a-z0-9]/.test(normalizedQuery)
    ? compactText(normalizedQuery)
    : "";
  return words.some((word) => {
    if (word.startsWith(normalizedQuery)) return true;
    if (!compactQuery) return false;
    // ASCII command words cannot gain a different pinyin prefix. Avoid the
    // transliteration/cache work for every token in large snippet libraries.
    if (/^[\x20-\x7E]+$/.test(word)) return word.startsWith(compactQuery);
    const { full, initials } = getPinyinVariants(word);
    return (
      (full !== "" && full.startsWith(compactQuery))
      || (initials !== "" && initials.startsWith(normalizedQuery))
    );
  });
}

/**
 * Host search should avoid mixing label/group tokens with hostname/IP tokens.
 * Otherwise queries like "山东 6-1" can accidentally match:
 * - "山东" from group/label
 * - "6" / "1" from hostname IP
 * across different fields.
 */
export type HostSearchLike = {
  label?: string | null;
  hostname?: string | null;
  group?: string | null;
  username?: string | null;
  tags?: Array<string | null | undefined> | null;
};

export function matchesHostSearchQuery(
  query: string,
  hostLike: HostSearchLike,
): boolean {
  return getHostSearchMatch(query, hostLike).matched;
}

/**
 * Workspace host pickers (Ctrl+Shift+J create/append) reuse vault host search,
 * including tags, and keep username as an extra field.
 */
export function matchesWorkspaceHostPickerQuery(
  query: string,
  hostLike: HostSearchLike,
): boolean {
  return matchesHostSearchQuery(query, hostLike)
    || matchesSearchQuery(query, hostLike.username);
}

type HostMatchField = "label" | "hostname" | "group" | "tag";
type HostMatchMethod = "literal" | "compact" | "pinyinFull" | "pinyinInitials";
type HostMatchPhase = "strict" | "loose";

export type HostSearchMatchDetail = {
  segment: string;
  field: HostMatchField;
  method: HostMatchMethod;
  phase: HostMatchPhase;
};

export type HostSearchMatchResult = {
  matched: boolean;
  phase: HostMatchPhase | "none";
  score: number;
  details: HostSearchMatchDetail[];
};

type HostSearchFieldSource = {
  field: HostMatchField;
  allowPinyin: boolean;
  text: string;
  compact: string;
};

function scoreHostMatch(detail: HostSearchMatchDetail): number {
  return HOST_FIELD_SCORE[detail.field] + HOST_METHOD_SCORE[detail.method];
}

function selectBetterHostMatch(
  left: HostSearchMatchResult,
  right: HostSearchMatchResult,
): HostSearchMatchResult {
  if (!left.matched) return right;
  if (!right.matched) return left;
  if (left.score !== right.score) {
    return left.score > right.score ? left : right;
  }
  if (left.details.length !== right.details.length) {
    return left.details.length < right.details.length ? left : right;
  }
  return left;
}

function matchSegmentAgainstSource(
  segment: string,
  source: HostSearchFieldSource,
): HostSearchMatchDetail | null {
  if (source.text.includes(segment)) {
    return {
      segment,
      field: source.field,
      method: "literal",
      phase: "strict",
    };
  }

  // Keep punctuation semantic (e.g. 6-, 6-1, 10.6.1.8): no compact fallback.
  if (SEARCH_QUERY_PUNCT_REGEX.test(segment)) return null;

  const compactSegment = compactText(segment);
  if (
    compactSegment
    && source.field !== "hostname"
    && source.compact.includes(compactSegment)
  ) {
    return {
      segment,
      field: source.field,
      method: "compact",
      phase: "loose",
    };
  }

  if (!source.allowPinyin || !/[a-z]/i.test(segment)) return null;
  // Literal and compact checks already cover printable ASCII host fields.
  // Keep their ranking above; skip only redundant transliteration/cache work.
  if (/^[\x20-\x7E]+$/.test(source.text)) return null;

  const { full, initials } = getPinyinVariants(source.text);
  if (full && full.includes(compactSegment)) {
    return {
      segment,
      field: source.field,
      method: "pinyinFull",
      phase: "loose",
    };
  }
  if (initials && initials.includes(compactSegment)) {
    return {
      segment,
      field: source.field,
      method: "pinyinInitials",
      phase: "loose",
    };
  }

  return null;
}

function evaluateHostFieldGroup(
  segments: string[],
  sources: HostSearchFieldSource[],
): HostSearchMatchResult {
  if (sources.length === 0) {
    return { matched: false, phase: "none", score: 0, details: [] };
  }

  const details: HostSearchMatchDetail[] = [];
  let isStrict = true;

  for (const segment of segments) {
    let best: HostSearchMatchDetail | null = null;
    let bestScore = -1;

    for (const source of sources) {
      const detail = matchSegmentAgainstSource(segment, source);
      if (!detail) continue;

      const score = scoreHostMatch(detail);
      if (score > bestScore) {
        best = detail;
        bestScore = score;
      }
    }

    if (!best) {
      return { matched: false, phase: "none", score: 0, details: [] };
    }
    if (best.phase !== "strict") {
      isStrict = false;
    }
    details.push(best);
  }

  const phase: HostMatchPhase = isStrict ? "strict" : "loose";
  const baseScore = details.reduce((sum, detail) => sum + scoreHostMatch(detail), 0);
  return {
    matched: true,
    phase,
    score: baseScore + HOST_PHASE_SCORE_BONUS[phase],
    details,
  };
}

export function getHostSearchMatch(
  query: string,
  hostLike: HostSearchLike,
): HostSearchMatchResult {
  const normalizedQuery = normalizeText(query);
  if (!normalizedQuery) {
    return { matched: true, phase: "strict", score: HOST_PHASE_SCORE_BONUS.strict, details: [] };
  }

  const querySegments = normalizedQuery
    .split(SEARCH_QUERY_SEGMENT_SPLIT_REGEX)
    .map((segment) => segment.trim())
    .filter(Boolean);
  if (querySegments.length === 0) {
    return { matched: true, phase: "strict", score: HOST_PHASE_SCORE_BONUS.strict, details: [] };
  }

  const label = normalizeText(hostLike.label ?? "");
  const group = normalizeText(hostLike.group ?? "");
  const hostname = normalizeText(hostLike.hostname ?? "");
  const tags = (hostLike.tags ?? [])
    .filter((tag): tag is string => typeof tag === "string" && tag.trim().length > 0)
    .map((tag) => normalizeText(tag));

  const humanSources: HostSearchFieldSource[] = [];
  if (label) {
    humanSources.push({
      field: "label",
      allowPinyin: true,
      text: label,
      compact: compactText(label),
    });
  }
  if (group) {
    humanSources.push({
      field: "group",
      allowPinyin: true,
      text: group,
      compact: compactText(group),
    });
  }
  for (const tag of tags) {
    humanSources.push({
      field: "tag",
      allowPinyin: true,
      text: tag,
      compact: compactText(tag),
    });
  }

  const networkSources: HostSearchFieldSource[] = hostname
    ? [{
      field: "hostname",
      allowPinyin: false,
      text: hostname,
      compact: compactText(hostname),
    }]
    : [];

  const humanResult = evaluateHostFieldGroup(querySegments, humanSources);
  const networkResult = evaluateHostFieldGroup(querySegments, networkSources);
  return selectBetterHostMatch(humanResult, networkResult);
}