douyin-archive/lib/search.ts

// 子串搜索的查询解析与摘要高亮工具。
// 匹配本身在 Postgres 里用正则匹配 lower() 后的文本完成，这里负责：
// - 把用户输入拆成关键词（空格分隔，全部需要命中；以 - 开头的词表示排除）
// - 在原文里找命中位置，截取摘要片段并标出高亮

export type SnippetPart = { text: string; hit: boolean };

export type ParsedQuery = { terms: string[]; excludes: string[] };

const MAX_TERMS = 10;
const MAX_TERM_LENGTH = 50;

export function parseSearchQuery(raw: string): ParsedQuery {
  const terms: string[] = [];
  const excludes: string[] = [];
  const seen = new Set<string>();
  for (const token of raw.trim().split(/\s+/)) {
    if (!token) continue;
    const isExclude = token.length > 1 && token.startsWith("-");
    const word = (isExclude ? token.slice(1) : token)
      .slice(0, MAX_TERM_LENGTH)
      .toLowerCase();
    const key = `${isExclude ? "-" : "+"}${word}`;
    if (!word || seen.has(key)) continue;
    seen.add(key);
    (isExclude ? excludes : terms).push(word);
    if (terms.length + excludes.length >= MAX_TERMS) break;
  }
  return { terms, excludes };
}

type Range = { start: number; end: number; term: string };

function escapeRegExp(s: string) {
  return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
}

// 纯英文/数字关键词要求「词首」对齐：前面不能紧挨字母数字。
// 这样 "ai" 不会命中 "said"，但 "用AI做"、"AIGC" 仍能命中；中文关键词是纯子串匹配。
const ASCII_WORD = /^[a-z0-9]/;

/** 关键词对应的 Postgres 正则（ARE），用于匹配已 lower() 的文本。 */
export function termToPgPattern(term: string) {
  const escaped = escapeRegExp(term);
  return ASCII_WORD.test(term) ? `(^|[^a-z0-9])${escaped}` : escaped;
}

function termToJsPattern(term: string) {
  const escaped = escapeRegExp(term);
  return ASCII_WORD.test(term) ? `(?<![a-z0-9])${escaped}` : escaped;
}

function findMatches(text: string, terms: string[]): Range[] {
  if (!text || terms.length === 0) return [];
  // 长词优先，避免 "老百姓" 被 "百姓" 抢先匹配成更短的高亮
  const sorted = [...terms].sort((a, b) => b.length - a.length);
  const re = new RegExp(sorted.map(termToJsPattern).join("|"), "giu");
  const ranges: Range[] = [];
  for (const m of text.matchAll(re)) {
    const start = m.index ?? 0;
    ranges.push({
      start,
      end: start + m[0].length,
      term: m[0].toLowerCase(),
    });
  }
  return ranges;
}

// 不要把 UTF-16 代理对（emoji 等）切成两半
function safeStart(text: string, i: number) {
  const c = text.charCodeAt(i);
  return c >= 0xdc00 && c <= 0xdfff ? i - 1 : i;
}
function safeEnd(text: string, i: number) {
  const c = text.charCodeAt(i);
  return c >= 0xdc00 && c <= 0xdfff ? i + 1 : i;
}

function highlightRange(
  text: string,
  from: number,
  to: number,
  matches: Range[],
): SnippetPart[] {
  const parts: SnippetPart[] = [];
  let cursor = from;
  for (const m of matches) {
    if (m.end <= from || m.start >= to) continue;
    const s = Math.max(m.start, from);
    const e = Math.min(m.end, to);
    if (s < cursor) continue;
    if (s > cursor) parts.push({ text: text.slice(cursor, s), hit: false });
    parts.push({ text: text.slice(s, e), hit: true });
    cursor = e;
  }
  if (cursor < to) parts.push({ text: text.slice(cursor, to), hit: false });
  return parts;
}

function normalizeWhitespace(parts: SnippetPart[]): SnippetPart[] {
  return parts
    .map((p) => ({ ...p, text: p.text.replace(/\s*\n\s*/g, " ") }))
    .filter((p) => p.text.length > 0);
}

/** 整段文字全部返回，只标出命中位置（用于描述）。 */
export function highlightAll(text: string, terms: string[]): SnippetPart[] {
  if (!text) return [];
  return normalizeWhitespace(
    highlightRange(text, 0, text.length, findMatches(text, terms)),
  );
}

/**
 * 截取命中位置附近的片段（用于语音转写这类长文本）。
 * 尽量让不同片段覆盖不同的关键词；片段会对齐到换行（转写的句子边界）。
 */
export function buildSnippet(
  text: string,
  terms: string[],
  { before = 24, after = 60, maxFragments = 2 } = {},
): SnippetPart[] {
  if (!text) return [];
  const matches = findMatches(text, terms);
  if (matches.length === 0) {
    const end = safeEnd(text, Math.min(text.length, before + after));
    const parts = normalizeWhitespace([
      { text: text.slice(0, end), hit: false },
    ]);
    if (end < text.length) parts.push({ text: "…", hit: false });
    return parts;
  }

  const windows: { start: number; end: number }[] = [];
  const covered = new Set<string>();
  const pickNext = () =>
    matches.find(
      (m) =>
        !covered.has(m.term) &&
        windows.every((w) => m.start >= w.end || m.end <= w.start),
    ) ??
    matches.find((m) =>
      windows.every((w) => m.start >= w.end || m.end <= w.start),
    );

  while (windows.length < maxFragments) {
    const m = pickNext();
    if (!m) break;
    let start = Math.max(0, m.start - before);
    let end = Math.min(text.length, m.end + after);
    const nlBefore = text.lastIndexOf("\n", m.start - 1);
    if (nlBefore >= start) start = nlBefore + 1;
    const nlAfter = text.indexOf("\n", m.end);
    if (nlAfter !== -1 && nlAfter < end) end = nlAfter;
    windows.push({ start: safeStart(text, start), end: safeEnd(text, end) });
    for (const other of matches) {
      if (other.start >= start && other.end <= end) covered.add(other.term);
    }
  }

  windows.sort((a, b) => a.start - b.start);
  const parts: SnippetPart[] = [];
  windows.forEach((w, i) => {
    if (i === 0 ? w.start > 0 : true) parts.push({ text: "…", hit: false });
    parts.push(...highlightRange(text, w.start, w.end, matches));
  });
  if (windows[windows.length - 1].end < text.length) {
    parts.push({ text: "…", hit: false });
  }
  return normalizeWhitespace(parts);
}