vibevoice-tts

/**
 * @file processor.ts
 * @description VibeVoice TTS formatting pipeline.
 *
 * Pipeline stages:
 *   1. Try JSON input parse (structured input)
 *   2. Clean input text (smart quotes, emoji, HTML, URLs, etc.)
 *   3. Detect multi-speaker format or force single/multi mode
 *   4. Clamp speaker IDs to valid range (1-4)
 *   5. Build speaker metadata with voice/emotion annotations
 *   6. Format script output
 *   7. Collect warnings
 *
 * VibeVoice does NOT support SSML — output is plain text or script format.
 * Emotion is an annotation for voice reference selection, never an inline tag.
 */

import type {
  VibeVoiceResult,
  SpeakerInfo,
  ScriptLine,
  EmotionLabel,
  EmotionMapping,
} from "./types";
import {
  cleanText,
  detectSpeakers,
  tryParseJsonInput,
  formatScript,
  applyEmotionMapping,
} from "./formatter";

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const SPEAKER_MIN = 1;
const SPEAKER_MAX = 4;

/**
 * Default emotion-to-slot mapping (PsiPi method).
 * Slot 1: neutral/calm — slot 2: happy/sad — slot 3: angry/excited.
 */
const DEFAULT_EMOTION_MAP: Readonly<Record<string, number>> = {
  neutral: 1,
  calm: 1,
  happy: 2,
  sad: 2,
  angry: 3,
  excited: 3,
};

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/**
 * Process input text through the VibeVoice formatting pipeline.
 *
 * @param text    - The raw input text (or JSON array of speaker lines)
 * @param params  - Formatting parameters
 * @returns A VibeVoiceResult with cleaned text, script, and speaker metadata
 */
export function processText(
  text: string,
  params: {
    mode: "single" | "multi" | "auto";
    speakers: number;
    voices: ReadonlyArray<string>;
    emotions: ReadonlyArray<string>;
    clean: boolean;
    output: "plain" | "script" | "json";
    emotion_map?: EmotionMapping;
  },
): VibeVoiceResult {
  const warnings: string[] = [];
  let scriptLines: ScriptLine[];

  // --------------------------------------------------
  // Stage 1: Try JSON input parsing (structured input)
  // --------------------------------------------------
  const jsonLines = tryParseJsonInput(text);
  if (jsonLines !== null) {
    // JSON input is already structured — clean individual text fields
    if (params.clean) {
      scriptLines = jsonLines.map((line) => ({
        speaker: line.speaker,
        text: cleanText(line.text),
      }));
    } else {
      scriptLines = [...jsonLines];
    }

    // Warn if input was JSON but mode says single — probably fine, move on
  } else {
    // --------------------------------------------------
    // Stage 2: Clean text
    // --------------------------------------------------
    const cleaned = params.clean ? cleanText(text) : text;

    if (cleaned.length === 0) {
      warnings.push("Input text is empty after cleaning.");
      return buildEmptyResult(text, cleaned, warnings);
    }

    // --------------------------------------------------
    // Stage 3: Detect or force speaker mode
    // --------------------------------------------------
    if (params.mode === "single") {
      // Single speaker — every non-empty line is speaker 1
      const lines = splitIntoLines(cleaned);
      scriptLines = lines.map((l) => ({ speaker: 1, text: l }));
    } else {
      // "auto" or "multi": try to detect "Speaker N:" prefixes
      const detected = detectSpeakers(cleaned);

      if (detected !== null) {
        scriptLines = detected;
      } else if (params.mode === "multi") {
        warnings.push(
          'Mode is "multi" but no "Speaker N:" prefixes found in text. ' +
            "Falling back to single-speaker mode.",
        );
        const lines = splitIntoLines(cleaned);
        scriptLines = lines.map((l) => ({ speaker: 1, text: l }));
      } else {
        // "auto" with no prefixes — single speaker
        const lines = splitIntoLines(cleaned);
        scriptLines = lines.map((l) => ({ speaker: 1, text: l }));
      }
    }
  }

  // --------------------------------------------------
  // Stage 3.5: Emotion annotation mapping (PsiPi method)
  // --------------------------------------------------
  const emotionMap: Readonly<Record<string, number>> =
    params.emotion_map ?? DEFAULT_EMOTION_MAP;
  scriptLines = applyEmotionMapping(scriptLines, emotionMap);

  // --------------------------------------------------
  // Stage 4: Clamp speaker IDs to valid range
  // --------------------------------------------------
  scriptLines = clampSpeakerIds(scriptLines, warnings);

  // --------------------------------------------------
  // Stage 5: Build speaker metadata
  // --------------------------------------------------
  const uniqueSpeakers = collectUniqueSpeakers(scriptLines);
  const speakers = buildSpeakerMetadata(
    params.speakers,
    uniqueSpeakers,
    scriptLines,
    params.voices,
    params.emotions,
    emotionMap,
  );

  // Validate voice/emotion array lengths
  if (params.voices.length > 0 && params.voices.length < params.speakers) {
    warnings.push(
      `voices array length (${params.voices.length}) is less than ` +
        `speakers (${params.speakers}). Missing voices will be undefined.`,
    );
  }
  if (params.emotions.length > 0 && params.emotions.length < params.speakers) {
    warnings.push(
      `emotions array length (${params.emotions.length}) is less than ` +
        `speakers (${params.speakers}). Missing emotions will be undefined.`,
    );
  }

  // --------------------------------------------------
  // Stage 6: Format output
  // --------------------------------------------------
  const script = formatScript(scriptLines, params.output);

  return {
    original: text,
    cleaned: params.clean ? cleanText(text) : text,
    script,
    speakers,
    warnings,
  };
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

/**
 * Split text into non-empty lines, trimming whitespace.
 * Preserves line breaks as semantic boundaries for speaker assignment.
 */
function splitIntoLines(text: string): string[] {
  return text
    .split("\n")
    .map((l) => l.trim())
    .filter((l) => l.length > 0);
}

/**
 * Clamp all speaker IDs to the valid range [SPEAKER_MIN, SPEAKER_MAX].
 * Emits a warning for each out-of-range speaker.
 */
function clampSpeakerIds(
  lines: ScriptLine[],
  warnings: string[],
): ScriptLine[] {
  return lines.map((line) => {
    if (line.speaker >= SPEAKER_MIN && line.speaker <= SPEAKER_MAX) {
      return line;
    }
    const clamped = Math.max(SPEAKER_MIN, Math.min(SPEAKER_MAX, line.speaker));
    warnings.push(
      `Speaker ID ${line.speaker} is out of range (${SPEAKER_MIN}-${SPEAKER_MAX}). Clamped to ${clamped}.`,
    );
    return { speaker: clamped, text: line.text };
  });
}

/**
 * Collect unique speaker IDs from script lines, sorted ascending.
 */
function collectUniqueSpeakers(lines: ReadonlyArray<ScriptLine>): number[] {
  const ids = new Set<number>();
  for (const line of lines) {
    ids.add(line.speaker);
  }
  return [...ids].sort((a, b) => a - b);
}

/**
 * Build SpeakerInfo array for all declared speakers.
 *
 * @param declaredCount - The number of speakers declared in params
 * @param activeIds     - IDs that actually appear in the script
 * @param lines         - All script lines (for line counting)
 * @param voices        - Voice names per speaker
 * @param emotions      - Emotion labels per speaker
 */
function buildSpeakerMetadata(
  declaredCount: number,
  activeIds: number[],
  lines: ReadonlyArray<ScriptLine>,
  voices: ReadonlyArray<string>,
  emotions: ReadonlyArray<string>,
  emotionMap: Readonly<Record<string, number>>,
): SpeakerInfo[] {
  // Count lines per speaker
  const lineCounts = new Map<number, number>();
  for (const line of lines) {
    lineCounts.set(line.speaker, (lineCounts.get(line.speaker) ?? 0) + 1);
  }

  // Determine which IDs to include:
  // all declared speakers (1..declaredCount) plus any active IDs beyond that
  const includedIds = new Set<number>();
  for (let i = 1; i <= declaredCount; i++) {
    includedIds.add(i);
  }
  for (const id of activeIds) {
    includedIds.add(id);
  }

  const result: SpeakerInfo[] = [];
  const sortedIds = [...includedIds].sort((a, b) => a - b);

  for (const id of sortedIds) {
    const voice = voices[id - 1];
    const emotionRaw = emotions[id - 1];
    // Validate emotion is a known label; fall back to the reverse lookup
    // of the emotion map (PsiPi method) when no explicit label was given.
    const mappedEmotion = Object.entries(emotionMap).find(
      ([, slot]) => slot === id,
    )?.[0];
    const emotion = isValidEmotion(emotionRaw)
      ? emotionRaw
      : isValidEmotion(mappedEmotion)
        ? mappedEmotion
        : undefined;

    result.push({
      id,
      voice: voice !== undefined && voice.length > 0 ? voice : undefined,
      emotion,
      lines: lineCounts.get(id) ?? 0,
    });
  }

  return result;
}

/**
 * Check whether a string is a valid EmotionLabel value.
 */
function isValidEmotion(label: string | undefined): label is EmotionLabel {
  if (label === undefined || label.length === 0) return false;
  const valid: EmotionLabel[] = [
    "neutral",
    "happy",
    "sad",
    "angry",
    "excited",
    "calm",
    "thoughtful",
    "whisper",
    "serious",
    "energetic",
    "warm",
    "cold",
    "mysterious",
    "dramatic",
  ];
  return (valid as ReadonlyArray<string>).includes(label);
}

/**
 * Build an empty result (for edge cases like empty text).
 */
/**
 * Build an empty result (for edge cases like empty text).
 */
function buildEmptyResult(
  original: string,
  cleaned: string,
  warnings: string[],
): VibeVoiceResult {
  return {
    original,
    cleaned,
    script: "",
    speakers: [],
    warnings,
  };
}