vibevoice-tts

/**
 * @file formatter.ts
 * @description Text cleaning and VibeVoice script formatting utilities.
 *
 * VibeVoice TTS format specifics:
 *   - No SSML support — plain text only
 *   - Multi-speaker: "Speaker N:" prefix per line
 *   - Emotion via voice reference selection, not inline tags
 *   - Text cleaning pipeline removes smart quotes, emoji, HTML, URLs, etc.
 */

import type { ScriptLine } from "./types";

// ---------------------------------------------------------------------------
// Regex patterns for text cleaning
// ---------------------------------------------------------------------------

const RE_CURLY_SINGLE = /[\u2018\u2019]/g;
const RE_CURLY_DOUBLE = /[\u201c\u201d]/g;
const RE_CHINESE_QUOTES = /[\u300c\u300d\u300e\u300f\u300a\u300b\u3014\u3015]/g;
const RE_EMOJI = /[\u{1F300}-\u{1FFFF}]/gu;
const RE_URL = /https?:\/\/\S+/g;
const RE_ELLIPSIS_CHAR = /\u2026/g;
const RE_ANNOTATION_BRACKET =
  /\[(?!(?:neutral|happy|sad|angry|excited|calm|thoughtful|whisper|serious|energetic|warm|cold|mysterious|dramatic|pause|emphasis|slower|faster|louder|softer)\b)[^\]]*\]/gi;
const RE_ANNOTATION_PAREN =
  /\((?!(?:neutral|happy|sad|angry|excited|calm|thoughtful|whisper|serious|energetic|warm|cold|mysterious|dramatic)\b)[^)]*\)/gi;

/**
 * HTML/XML tags — BUT NOT VibeVoice expression tags.
 * Expression tags are simple: <word> or <phrase with spaces>.
 * Real HTML has attributes, slashes, numbers, etc.
 * This regex skips expression-like tags and removes real markup.
 */
const RE_HTML_TAGS = /<(?!\/?[a-zA-Z][a-zA-Z0-9\s-]*>)[^>]*>/g;
const RE_H_WHITESPACE = /[^\S\n]+/g;
const RE_LINE_TRIM = /^[ \t]+|[ \t]+$/gm;
const RE_EXCESS_BLANKS = /\n{3,}/g;

// ---------------------------------------------------------------------------
// Regex patterns for speaker detection
// ---------------------------------------------------------------------------

const RE_SPEAKER_PREFIX = /^(Speaker\s*(\d+))\s*:\s*/i;
const RE_SPEAKER_SHORT = /^S\s*(\d+)\s*:\s*/i;
const RE_EMOTION_BRACKET =
  /^\[(neutral|happy|sad|angry|excited|calm|thoughtful|whisper|serious|energetic|warm|cold|mysterious|dramatic)\]\s*/i;
const RE_EMOTION_PAREN =
  /^\((neutral|happy|sad|angry|excited|calm|thoughtful|whisper|serious|energetic|warm|cold|mysterious|dramatic)\)\s*/i;
// ---------------------------------------------------------------------------
// Text Cleaning Pipeline
// ---------------------------------------------------------------------------

/**
 * Clean input text according to the VibeVoice text-cleaning specification.
 * See the function body for the ordered pipeline steps.
 */
function cleanTextImpl(text: string): string {
  let result = text;

  // 1. Smart/curly quotes -> straight quotes
  result = result.replace(RE_CURLY_SINGLE, "'");
  result = result.replace(RE_CURLY_DOUBLE, '"');

  // 2. Chinese quotation marks -> remove
  result = result.replace(RE_CHINESE_QUOTES, "");

  // 3. Emoji -> strip
  result = result.replace(RE_EMOJI, "");

  // 4. HTML/XML tags -> strip
  result = result.replace(RE_HTML_TAGS, "");

  // 5. Non-speech annotations -> remove (before URL so [URL] is safe)
  result = result.replace(RE_ANNOTATION_BRACKET, "");
  result = result.replace(RE_ANNOTATION_PAREN, "");

  // 6. Ellipsis character -> three dots
  result = result.replace(RE_ELLIPSIS_CHAR, "...");

  // 7. URLs -> [URL] placeholder
  result = result.replace(RE_URL, "[URL]");

  // 8. Tab characters -> spaces
  result = result.replace(/\t/g, " ");

  // 9. Collapse multiple horizontal whitespace -> single space
  result = result.replace(RE_H_WHITESPACE, " ");

  // 10. Trim leading/trailing whitespace per line
  result = result.replace(RE_LINE_TRIM, "");

  // 11. Collapse excess blank lines (3+ -> 2)
  result = result.replace(RE_EXCESS_BLANKS, "\n\n");

  return result.trim();
}

// ---------------------------------------------------------------------------
// JSON Input Parser
// ---------------------------------------------------------------------------

/**
 * Try to parse text as a JSON array of speaker-line objects.
 *
 * Expected format:
 * ```json
 * [
 *   { "speaker": 1, "text": "Hello!" },
 *   { "speaker": 2, "text": "Hi there!" }
 * ]
 * ```
 */
function tryParseJsonInputImpl(text: string): ScriptLine[] | null {
  const trimmed = text.trim();
  if (trimmed.length === 0 || trimmed[0] !== "[") {
    return null;
  }

  let parsed: unknown;
  try {
    parsed = JSON.parse(trimmed);
  } catch {
    return null;
  }

  if (!Array.isArray(parsed) || parsed.length === 0) {
    return null;
  }

  const result: ScriptLine[] = [];
  for (const item of parsed) {
    if (
      typeof item !== "object" ||
      item === null ||
      typeof item.speaker !== "number" ||
      typeof item.text !== "string"
    ) {
      return null;
    }
    result.push({ speaker: item.speaker, text: item.text });
  }

  return result;
}

// ---------------------------------------------------------------------------
// Speaker Detection
// ---------------------------------------------------------------------------

/**
 * Detect "Speaker N:" prefixes in text and parse into script lines.
 *
 * Recognised formats:
 *   - "Speaker 1: Hello"  (full form, case-insensitive)
 *   - "S1: Hello"         (short form)
 *   - "S 2: Hello"        (short form with space)
 *
 * Lines without a prefix continue the previous speaker's turn.
 * Returns null if NO speaker prefix is found anywhere in the text.
 */
function detectSpeakersImpl(text: string): ScriptLine[] | null {
  const lines = text.split("\n");
  const result: ScriptLine[] = [];
  let currentSpeaker = 1;
  let hasAnyPrefix = false;

  for (const rawLine of lines) {
    const trimmed = rawLine.trim();
    if (trimmed.length === 0) continue;

    // Try full "Speaker N:" prefix first (group 2 = just digits)
    const fullMatch = RE_SPEAKER_PREFIX.exec(trimmed);

    if (fullMatch !== null) {
      hasAnyPrefix = true;
      currentSpeaker = parseInt(fullMatch[2], 10);
      const rest = trimmed.slice(fullMatch[0].length).trim();
      if (rest.length > 0) {
        result.push({ speaker: currentSpeaker, text: rest });
      }
      continue;
    }

    // Try short "S N:" prefix (group 1 = just digits)
    const shortMatch = RE_SPEAKER_SHORT.exec(trimmed);
    if (shortMatch !== null) {
      hasAnyPrefix = true;
      currentSpeaker = parseInt(shortMatch[1], 10);
      const rest = trimmed.slice(shortMatch[0].length).trim();
      if (rest.length > 0) {
        result.push({ speaker: currentSpeaker, text: rest });
      }
      continue;
    }

    // Continuation of the current speaker's turn
    result.push({ speaker: currentSpeaker, text: trimmed });
  }

  return hasAnyPrefix ? result : null;
}

// ---------------------------------------------------------------------------
// Emotion Annotation (PsiPi method)
// ---------------------------------------------------------------------------
/**
 * Strip a leading [emotion] or (emotion) annotation from a line.
 *
 * Returns the detected emotion label (lowercased) or null when the line
 * carries no annotation. The annotation itself is always removed from the
 * returned text — VibeVoice never receives inline emotion tags.
 */
function extractEmotionAnnotationImpl(line: string): {
  emotion: string | null;
  text: string;
} {
  const bracket = RE_EMOTION_BRACKET.exec(line);
  if (bracket !== null) {
    return {
      emotion: bracket[1].toLowerCase(),
      text: line.slice(bracket[0].length).trim(),
    };
  }
  const paren = RE_EMOTION_PAREN.exec(line);
  if (paren !== null) {
    return {
      emotion: paren[1].toLowerCase(),
      text: line.slice(paren[0].length).trim(),
    };
  }
  return { emotion: null, text: line };
}

/**
 * Reassign speaker slots based on inline emotion annotations (PsiPi method).
 *
 * For each line starting with an [emotion] / (emotion) annotation, the
 * annotation is stripped and the line is reassigned to the speaker slot
 * mapped to that emotion. Lines without annotations keep their speaker.
 * Emotions missing from the map also keep their original speaker.
 */
function applyEmotionMappingImpl(
  lines: ReadonlyArray<ScriptLine>,
  emotionMap: Readonly<Record<string, number>>,
): ScriptLine[] {
  return lines.map((line) => {
    const { emotion, text } = extractEmotionAnnotationImpl(line.text);
    if (emotion === null) {
      return line;
    }
    const slot = emotionMap[emotion];
    if (slot === undefined) {
      // Annotation present but unmapped — still strip it from the output.
      return { speaker: line.speaker, text };
    }
    return { speaker: slot, text };
  });
}
// ---------------------------------------------------------------------------
// Script Formatting
// ---------------------------------------------------------------------------

/**
 * Format an array of script lines into the requested output representation.
 */
function formatScriptImpl(
  lines: ReadonlyArray<ScriptLine>,
  output: "plain" | "script" | "json",
): string {
  if (output === "json") {
    return JSON.stringify(lines, null, 2);
  }

  if (output === "plain") {
    return lines.map((l) => l.text).join("\n");
  }

  return lines.map((l) => `Speaker ${l.speaker}: ${l.text}`).join("\n");
}

// ---------------------------------------------------------------------------
// Re-exports as const values (workaround for TS export resolution)
// ---------------------------------------------------------------------------

export const cleanText: (text: string) => string = cleanTextImpl;
export const tryParseJsonInput: (text: string) => ScriptLine[] | null =
  tryParseJsonInputImpl;
export const detectSpeakers: (text: string) => ScriptLine[] | null =
  detectSpeakersImpl;
export const formatScript: (
  lines: ReadonlyArray<ScriptLine>,
  output: "plain" | "script" | "json",
) => string = formatScriptImpl;
export const extractEmotionAnnotation: (line: string) => {
  emotion: string | null;
  text: string;
} = extractEmotionAnnotationImpl;
export const applyEmotionMapping: (
  lines: ReadonlyArray<ScriptLine>,
  emotionMap: Readonly<Record<string, number>>,
) => ScriptLine[] = applyEmotionMappingImpl;