vibevoice-tts

/**
 * @file preparer.ts
 * @description Podcast script preparation pipeline for VibeVoice TTS.
 *
 * Handles the full preparation workflow:
 *   - Abbreviation expansion (first-use or all)
 *   - Direction tag extraction ([emphasis], [slower], etc.)
 *   - Pause marker conversion ([pause 1s] → empty lines)
 *   - Sentence splitting & paragraph breaking
 *   - Script chunking (~500 words per chunk)
 *   - TTS parameter recommendations
 *
 * VibeVoice specifics:
 *   - Empty lines = natural pauses
 *   - No SSML support — plain text only
 *   - Max ~500 words per chunk for best quality
 *   - 150-170 WPM is standard podcast speaking rate
 */

import type {
  ScriptChunk,
  TTSRecommendations,
  PauseMarker,
  PrepareResult,
  DirectionTag,
} from "./types";

// ---------------------------------------------------------------------------
// Abbreviation Expansion
// ---------------------------------------------------------------------------

const ABBREVIATION_MAP: Readonly<Record<string, string>> = {
  "e.g.": "for example",
  "i.e.": "that is",
  "etc.": "et cetera",
  "vs.": "versus",
  "Mr.": "Mister",
  "Mrs.": "Missus",
  "Ms.": "Miss or Missus",
  "Dr.": "Doctor",
  "Prof.": "Professor",
  "St.": "Saint",
  "Jr.": "Junior",
  "Sr.": "Senior",
  "Inc.": "Incorporated",
  "Ltd.": "Limited",
  "Co.": "Company",
  "Corp.": "Corporation",
  AI: "A I",
  TTS: "text to speech",
  API: "A P I",
  URL: "U R L",
  HTTP: "H T T P",
  HTTPS: "H T T P S",
  JSON: "jay son",
  CLI: "C L I",
  UI: "U I",
  UX: "U X",
};

/**
 * Expand common abbreviations for TTS clarity.
 * `firstUseOnly`: expand only the first occurrence of each abbreviation.
 * `all`: expand every occurrence.
 */
export function expandAbbreviations(
  text: string,
  options: { firstUseOnly?: boolean; all?: boolean } = {},
): string {
  const { firstUseOnly = true, all = false } = options;

  if (all) {
    let result = text;
    for (const [abbr, expansion] of Object.entries(ABBREVIATION_MAP)) {
      const escaped = abbr.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
      result = result.replace(new RegExp(escaped, "g"), expansion);
    }
    return result;
  }

  // First-use only: expand first occurrence of each abbreviation
  let result = text;
  for (const [abbr, expansion] of Object.entries(ABBREVIATION_MAP)) {
    const escaped = abbr.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
    const regex = new RegExp(escaped);
    const firstMatch = regex.exec(result);
    if (firstMatch) {
      result =
        result.slice(0, firstMatch.index) +
        expansion +
        result.slice(firstMatch.index + abbr.length);
    }
  }
  return result;
}

// ---------------------------------------------------------------------------
// Direction Tag Extraction
// ---------------------------------------------------------------------------

const RE_DIRECTION_TAG = /\[(emphasis|slower|faster|whisper|louder|softer)\]/gi;

/**
 * Extract direction tags from text and return them as user-facing hints.
 * Unlike emotion annotations, direction tags are preserved as metadata
 * to guide the human operator on TTS parameter adjustments.
 */
export function extractDirectionTags(text: string): {
  cleanedText: string;
  directions: Array<{ tag: DirectionTag; position: number; context: string }>;
} {
  const directions: Array<{
    tag: DirectionTag;
    position: number;
    context: string;
  }> = [];
  const re = new RegExp(RE_DIRECTION_TAG.source, RE_DIRECTION_TAG.flags);
  let match: RegExpExecArray | null;

  while ((match = re.exec(text)) !== null) {
    const start = Math.max(0, match.index - 30);
    const end = Math.min(text.length, match.index + match[0].length + 30);
    directions.push({
      tag: match[1].toLowerCase() as DirectionTag,
      position: match.index,
      context: text.slice(start, end).trim(),
    });
  }

  // Strip direction tags from the cleaned text
  const cleaned = text.replace(RE_DIRECTION_TAG, "").trim();

  return { cleanedText: cleaned, directions };
}

// ---------------------------------------------------------------------------
// Pause Marker Conversion
// ---------------------------------------------------------------------------

const RE_PAUSE_MARKER = /\[pause\s*(\d+(?:\.\d+)?)s?\]/gi;

/**
 * Convert explicit [pause Xs] markers to empty lines.
 * VibeVoice interprets empty lines as natural pauses.
 */
export function convertPauseMarkers(text: string): {
  convertedText: string;
  pauses: PauseMarker[];
} {
  const pauses: PauseMarker[] = [];
  const lines = text.split("\n");
  const result: string[] = [];
  let lineIndex = 0;

  for (const line of lines) {
    const pauseMatches = [...line.matchAll(RE_PAUSE_MARKER)];

    if (pauseMatches.length > 0) {
      for (const m of pauseMatches) {
        pauses.push({
          atLine: lineIndex,
          seconds: parseFloat(m[1]),
          reason: "explicit",
        });
      }

      // Replace pause markers with empty string, trim
      const converted = line.replace(RE_PAUSE_MARKER, "").trim();
      if (converted.length > 0) {
        result.push(converted);
      }
      // Insert empty line for the pause
      result.push("");
      lineIndex++;
    } else {
      result.push(line);
    }
    lineIndex++;
  }

  return { convertedText: result.join("\n"), pauses };
}

// ---------------------------------------------------------------------------
// Expression Tag Conversion
// ---------------------------------------------------------------------------

/**
 * Convert bracket-based expression markers to VibeVoice expression tags.
 * E.g., `[laughs]` → `<laughs>`, `[sighs]` → `<sighs>`.
 *
 * Supports ANY expression — not just known tags. The model can invent
 * new ones freely: `[roar]` → `<roar>`, `[sing softly]` → `<sing softly>`.
 */
export function convertExpressionBrackets(text: string): {
  convertedText: string;
  expressions: Array<{ tag: string; position: number; original: string }>;
} {
  const expressions: Array<{
    tag: string;
    position: number;
    original: string;
  }> = [];

  // Match [word] or [phrase with spaces] that looks like an expression
  // Exclude known non-expression brackets: [pause Xs], [emphasis], etc.
  const RE_EXPRESSION = /\[([a-zA-Z][a-zA-Z0-9\s-]*)\](?!\s*\d)/g;

  let result = text;

  // We need to process all matches and track positions
  const re = new RegExp(RE_EXPRESSION.source, RE_EXPRESSION.flags);
  const matches: Array<{ match: RegExpExecArray; index: number }> = [];
  let match: RegExpExecArray | null;
  while ((match = re.exec(text)) !== null) {
    matches.push({ match, index: match.index });
  }

  // Process in reverse order to preserve positions
  for (let i = matches.length - 1; i >= 0; i--) {
    const { match: m, index } = matches[i];
    const tag = m[1].trim();

    // Skip if it looks like a direction tag or pause marker
    if (
      /^(?:pause|emphasis|slower|faster|louder|softer|whisper|URL|section|chapter)/i.test(
        tag,
      )
    ) {
      continue;
    }

    expressions.push({
      tag,
      position: index,
      original: m[0],
    });

    // Replace [tag] with <tag>
    const before = result.slice(0, index);
    const after = result.slice(index + m[0].length);
    result = before + `<${tag}>` + after;
  }

  // Reverse expressions array since we processed in reverse
  expressions.reverse();

  return { convertedText: result, expressions };
}
// ---------------------------------------------------------------------------
// Sentence Splitting
// ---------------------------------------------------------------------------

const ABBREV_EXCEPTIONS = new Set([
  "e.g.",
  "i.e.",
  "etc.",
  "mr.",
  "mrs.",
  "ms.",
  "dr.",
  "prof.",
  "st.",
  "jr.",
  "sr.",
  "vs.",
  "inc.",
  "ltd.",
  "co.",
  "corp.",
]);

/**
 * Split a paragraph into sentences.
 * Handles common edge cases: abbreviations, decimal numbers, ellipsis.
 */
export function splitIntoSentences(paragraph: string): string[] {
  const sentences: string[] = [];
  let current = "";
  const chars = paragraph.split("");
  let i = 0;

  while (i < chars.length) {
    current += chars[i];

    if (/[.!?]/.test(chars[i])) {
      const rest = paragraph.slice(i + 1).trimStart();
      const nextChar = rest[0];

      // Not a sentence boundary if followed by lowercase (likely abbreviation)
      if (
        nextChar &&
        nextChar === nextChar.toLowerCase() &&
        /[a-z]/.test(nextChar)
      ) {
        i++;
        continue;
      }

      // Not a sentence boundary if part of ellipsis
      if (
        i + 1 < chars.length &&
        chars[i + 1] === "." &&
        i + 2 < chars.length &&
        chars[i + 2] === "."
      ) {
        i++;
        continue;
      }

      // Check for common abbreviations
      const before = current.toLowerCase();
      if ([...ABBREV_EXCEPTIONS].some((a: string) => before.endsWith(a))) {
        i++;
        continue;
      }

      // Sentence boundary
      sentences.push(current.trim());
      current = "";
    }

    i++;
  }

  if (current.trim().length > 0) {
    sentences.push(current.trim());
  }

  return sentences.filter((s) => s.length > 0);
}

/**
 * Split long paragraphs into smaller chunks of max N sentences.
 * Returns an array of shorter paragraphs.
 */
export function splitParagraphs(
  text: string,
  maxSentences: number = 3,
): string[] {
  const paragraphs = text.split(/\n\n+/);
  const result: string[] = [];

  for (const para of paragraphs) {
    const trimmed = para.trim();
    if (trimmed.length === 0) continue;

    const sentences = splitIntoSentences(trimmed);
    if (sentences.length <= maxSentences) {
      result.push(trimmed);
    } else {
      for (let i = 0; i < sentences.length; i += maxSentences) {
        const group = sentences.slice(i, i + maxSentences);
        result.push(group.join(" "));
      }
    }
  }

  return result;
}

// ---------------------------------------------------------------------------
// Script Chunking
// ---------------------------------------------------------------------------

/**
 * Count words in text (handles multi-language).
 */
export function countWords(text: string): number {
  // Split on whitespace, filter empty
  return text.split(/\s+/).filter((w) => w.length > 0).length;
}

/**
 * Split a prepared script into chunks of approximately `targetWords` words.
 * Chunks are aligned with paragraph boundaries when possible.
 */
export function chunkScript(
  text: string,
  targetWords: number = 500,
): ScriptChunk[] {
  const paragraphs = text.split(/\n\n+/).filter((p) => p.trim().length > 0);
  const chunks: ScriptChunk[] = [];
  let currentText = "";
  let currentWords = 0;
  let chunkIndex = 1;
  const speakersInChunk = new Set<number>();

  for (const para of paragraphs) {
    const paraWords = countWords(para);

    // If adding this paragraph would exceed target and we already have content
    if (currentWords > 0 && currentWords + paraWords > targetWords) {
      // Finalize current chunk
      const trimmed = currentText.trim();
      chunks.push({
        index: chunkIndex++,
        text: trimmed,
        wordCount: currentWords,
        speakers: [...speakersInChunk].sort((a, b) => a - b),
      });
      currentText = "";
      currentWords = 0;
      speakersInChunk.clear();
    }

    // Extract speaker IDs from this paragraph
    const speakerMatches = para.matchAll(/Speaker\s+(\d+)/gi);
    for (const m of speakerMatches) {
      speakersInChunk.add(parseInt(m[1], 10));
    }

    // Add paragraph to current chunk
    if (currentText.length > 0) {
      currentText += "\n\n" + para;
    } else {
      currentText = para;
    }
    currentWords += paraWords;
  }

  // Finalize last chunk
  if (currentText.trim().length > 0) {
    chunks.push({
      index: chunkIndex,
      text: currentText.trim(),
      wordCount: currentWords,
      speakers: [...speakersInChunk].sort((a, b) => a - b),
    });
  }

  return chunks;
}

// ---------------------------------------------------------------------------
// TTS Recommendations
// ---------------------------------------------------------------------------

/**
 * Generate TTS parameter recommendations based on script analysis.
 */
export function generateTTSRecommendations(
  text: string,
  chunkSize: number = 500,
): TTSRecommendations {
  const wordCount = countWords(text);
  const hasMultipleSpeakers = /Speaker\s+\d+/gi.test(text);
  const hasDialogue = /Speaker\s+1:.*Speaker\s+2:/gis.test(text);

  return {
    // 0.6-0.7 for narration, 0.8-0.9 for dialogue
    temperature: hasDialogue ? 0.85 : 0.65,
    // 0.95x feels more natural than 1.0x
    speed: 0.95,
    // 150-170 WPM standard podcast rate
    wpm: hasDialogue ? 160 : 155,
    chunkSize,
  };
}

// ---------------------------------------------------------------------------
// Full Preparation Pipeline
// ---------------------------------------------------------------------------

/**
 * Run the full podcast script preparation pipeline.
 *
 * Steps:
 *   1. Expand abbreviations
 *   2. Extract direction tags (preserve as metadata)
 *   3. Convert pause markers to empty lines
 *   4. Split long paragraphs
 *   5. Chunk the script
 *   6. Generate TTS recommendations
 */
export function prepareScript(
  text: string,
  params: {
    chunkSize?: number;
    paragraphPauses?: boolean;
    sentencePauses?: boolean;
    maxSentencesPerParagraph?: number;
    expandAbbreviations?: boolean;
    preserveDirections?: boolean;
    output?: "plain" | "script" | "json";
  } = {},
): PrepareResult {
  const {
    chunkSize = 500,
    paragraphPauses = true,
    sentencePauses = false,
    maxSentencesPerParagraph = 3,
    expandAbbreviations: doExpand = true,
    preserveDirections = true,
    output = "script",
  } = params;

  const warnings: string[] = [];
  let processed = text;

  // Step 1: Expand abbreviations
  if (doExpand) {
    processed = expandAbbreviations(processed, { firstUseOnly: true });
  }

  // Step 2: Extract direction tags
  const directionResult = preserveDirections
    ? extractDirectionTags(processed)
    : {
        cleanedText: processed,
        directions: [] as Array<{
          tag: DirectionTag;
          position: number;
          context: string;
        }>,
      };
  processed = directionResult.cleanedText;

  if (directionResult.directions.length > 0) {
    warnings.push(
      `Found ${directionResult.directions.length} direction tag(s): ` +
        directionResult.directions.map((d) => `[${d.tag}]`).join(", ") +
        ". These are metadata hints — adjust TTS params accordingly.",
    );
  }

  // Step 3: Convert pause markers to empty lines
  const pauseResult = convertPauseMarkers(processed);
  processed = pauseResult.convertedText;
  const pauses: PauseMarker[] = [...pauseResult.pauses];

  // Step 4: Split long paragraphs
  if (sentencePauses) {
    const splitParas = splitParagraphs(processed, maxSentencesPerParagraph);
    processed = splitParas.join("\n\n");
  }

  // Ensure paragraph pauses are present (empty lines between paragraphs)
  if (paragraphPauses) {
    // Normalize: ensure double newlines between paragraphs
    processed = processed.replace(/\n{3,}/g, "\n\n");
  }

  // Step 5: Chunk the script
  const chunks = chunkScript(processed, chunkSize);

  // Step 6: Generate TTS recommendations
  const recommendations = generateTTSRecommendations(processed, chunkSize);

  // Calculate totals
  const totalWords = countWords(processed);
  const estimatedMinutes = totalWords / recommendations.wpm;

  // Warnings
  if (totalWords === 0) {
    warnings.push("Prepared script is empty after processing.");
  }
  if (chunks.length > 1) {
    warnings.push(
      `Script split into ${chunks.length} chunks. Generate each separately and concatenate with 200ms crossfade.`,
    );
  }
  if (totalWords > 4000) {
    warnings.push(
      `Very long script (${totalWords} words, ~${estimatedMinutes.toFixed(1)} min). ` +
        "Consider splitting into multiple episodes for better quality.",
    );
  }

  // Build the full prepared script
  let preparedScript = processed;
  if (output === "json") {
    preparedScript = JSON.stringify(chunks, null, 2);
  }

  return {
    original: text,
    preparedScript,
    chunks,
    pauses,
    recommendations,
    warnings,
    totalWords,
    estimatedMinutes,
  };
}