WellNuo/services/sherpaTTS.ts

/**
 * Sherpa TTS Service - Native implementation for offline Text-to-Speech
 * Uses react-native-sherpa-onnx-offline-tts with Piper VITS models
 */

import { Platform, NativeModules, NativeEventEmitter } from 'react-native';
import * as FileSystem from 'expo-file-system';
import { Asset } from 'expo-asset';

// Available Piper neural voices
export interface PiperVoice {
  id: string;
  name: string;
  description: string;
  modelDir: string;
  onnxFile: string;
  gender: 'male' | 'female';
  accent: string;
}

export const AVAILABLE_VOICES: PiperVoice[] = [
  {
    id: 'lessac',
    name: 'Lessac',
    description: 'American Female (Natural)',
    modelDir: 'vits-piper-en_US-lessac-medium',
    onnxFile: 'en_US-lessac-medium.onnx',
    gender: 'female',
    accent: 'US',
  },
  {
    id: 'ryan',
    name: 'Ryan',
    description: 'American Male (Natural)',
    modelDir: 'vits-piper-en_US-ryan-medium',
    onnxFile: 'en_US-ryan-medium.onnx',
    gender: 'male',
    accent: 'US',
  },
  {
    id: 'alba',
    name: 'Alba',
    description: 'British Female',
    modelDir: 'vits-piper-en_GB-alba-medium',
    onnxFile: 'en_GB-alba-medium.onnx',
    gender: 'female',
    accent: 'UK',
  },
];

interface SherpaTTSState {
  initialized: boolean;
  initializing: boolean;
  error: string | null;
  currentVoice: PiperVoice;
}

// Check if native module is available
const TTSManager = NativeModules.TTSManager;
const NATIVE_MODULE_AVAILABLE = !!TTSManager;

let ttsManagerEmitter: NativeEventEmitter | null = null;
if (NATIVE_MODULE_AVAILABLE) {
  ttsManagerEmitter = new NativeEventEmitter(TTSManager);
}

let currentState: SherpaTTSState = {
  initialized: false,
  initializing: false,
  error: null,
  currentVoice: AVAILABLE_VOICES[0],
};

// State listeners
const stateListeners: ((state: SherpaTTSState) => void)[] = [];

function updateState(updates: Partial<SherpaTTSState>) {
  currentState = { ...currentState, ...updates };
  notifyListeners();
}

function notifyListeners() {
  stateListeners.forEach(listener => listener({ ...currentState }));
}

export function addStateListener(listener: (state: SherpaTTSState) => void) {
  stateListeners.push(listener);
  listener({ ...currentState });
  return () => {
    const index = stateListeners.indexOf(listener);
    if (index >= 0) stateListeners.splice(index, 1);
  };
}

export function getState(): SherpaTTSState {
  return { ...currentState };
}

/**
 * Copy bundled TTS model assets to document directory for native access
 * NOTE: Temporarily disabled dynamic requires for Metro bundler compatibility
 */
async function copyModelToDocuments(voice: PiperVoice): Promise<string | null> {
  // TEMP: Skip dynamic requires - TTS models will be loaded differently
  console.log('[SherpaTTS] copyModelToDocuments temporarily disabled');
  return null;

  /* DISABLED - dynamic requires don't work with Metro bundler
  try {
    const destDir = `${FileSystem.documentDirectory}tts-models/${voice.modelDir}`;
    const onnxPath = `${destDir}/${voice.onnxFile}`;

    // Check if already copied
    const onnxInfo = await FileSystem.getInfoAsync(onnxPath);
    if (onnxInfo.exists) {
      console.log('[SherpaTTS] Model already exists at:', destDir);
      return destDir;
    }

    console.log('[SherpaTTS] Copying model to documents directory...');

    // Create destination directory
    await FileSystem.makeDirectoryAsync(destDir, { intermediates: true });

    // For Expo, we need to copy from assets
    // The models are in assets/tts-models/
    const assetBase = `../assets/tts-models/${voice.modelDir}`;

    // Copy main ONNX file
    const onnxAsset = Asset.fromModule(require(`../assets/tts-models/${voice.modelDir}/${voice.onnxFile}`));
    await onnxAsset.downloadAsync();
    if (onnxAsset.localUri) {
      await FileSystem.copyAsync({
        from: onnxAsset.localUri,
        to: onnxPath,
      });
    }

    // Copy tokens.txt
    const tokensAsset = Asset.fromModule(require(`../assets/tts-models/${voice.modelDir}/tokens.txt`));
    await tokensAsset.downloadAsync();
    if (tokensAsset.localUri) {
      await FileSystem.copyAsync({
        from: tokensAsset.localUri,
        to: `${destDir}/tokens.txt`,
      });
    }

    // Copy espeak-ng-data directory
    // This is more complex - need to copy entire directory
    const espeakDestDir = `${destDir}/espeak-ng-data`;
    await FileSystem.makeDirectoryAsync(espeakDestDir, { intermediates: true });

    // For now, we'll use the bundle path directly on iOS
    console.log('[SherpaTTS] Model copied successfully');
    return destDir;

  } catch (error) {
    console.error('[SherpaTTS] Error copying model:', error);
    return null;
  }
  */
}

/**
 * Get the path to bundled models (iOS bundle or Android assets)
 */
function getBundledModelPath(voice: PiperVoice): string | null {
  if (Platform.OS === 'ios') {
    // On iOS, assets are in the main bundle
    const bundlePath = NativeModules.RNFSManager?.MainBundlePath || '';
    if (!bundlePath) {
      // Try to construct path from FileSystem
      // Models should be copied during pod install or prebuild
      return null;
    }
    return `${bundlePath}/assets/tts-models/${voice.modelDir}`;
  } else if (Platform.OS === 'android') {
    // On Android, assets are extracted to files dir
    return `${FileSystem.documentDirectory}tts-models/${voice.modelDir}`;
  }
  return null;
}

/**
 * Initialize Sherpa TTS with a specific voice model
 */
export async function initializeSherpaTTS(voice?: PiperVoice): Promise<boolean> {
  if (!NATIVE_MODULE_AVAILABLE) {
    console.log('[SherpaTTS] Native module not available (Expo Go mode)');
    updateState({
      initialized: false,
      error: 'Native module not available - use native build'
    });
    return false;
  }

  if (currentState.initializing) {
    console.log('[SherpaTTS] Already initializing...');
    return false;
  }

  const selectedVoice = voice || currentState.currentVoice;
  updateState({ initializing: true, error: null });

  try {
    console.log('[SherpaTTS] Initializing with voice:', selectedVoice.name);

    // Get model paths
    // For native build, models should be in the app bundle
    // We use FileSystem.bundleDirectory on iOS
    let modelBasePath: string;

    if (Platform.OS === 'ios') {
      // iOS: Models are copied to bundle during build
      // Access via MainBundle
      const mainBundle = await FileSystem.getInfoAsync(FileSystem.bundleDirectory || '');
      if (mainBundle.exists) {
        modelBasePath = `${FileSystem.bundleDirectory}assets/tts-models/${selectedVoice.modelDir}`;
      } else {
        // Fallback: try document directory
        modelBasePath = `${FileSystem.documentDirectory}tts-models/${selectedVoice.modelDir}`;
      }
    } else {
      // Android: Extract from assets to document directory
      modelBasePath = `${FileSystem.documentDirectory}tts-models/${selectedVoice.modelDir}`;
    }

    // Check if model exists
    const modelPath = `${modelBasePath}/${selectedVoice.onnxFile}`;
    const tokensPath = `${modelBasePath}/tokens.txt`;
    const dataDirPath = `${modelBasePath}/espeak-ng-data`;

    console.log('[SherpaTTS] Model paths:', { modelPath, tokensPath, dataDirPath });

    // Create config JSON for native module
    const config = JSON.stringify({
      modelPath,
      tokensPath,
      dataDirPath,
    });

    // Initialize native TTS
    // Sample rate 22050, mono channel
    TTSManager.initializeTTS(22050, 1, config);

    updateState({
      initialized: true,
      initializing: false,
      currentVoice: selectedVoice,
      error: null
    });

    console.log('[SherpaTTS] Initialized successfully');
    return true;

  } catch (error) {
    const errorMessage = error instanceof Error ? error.message : 'Unknown error';
    console.error('[SherpaTTS] Initialization error:', errorMessage);
    updateState({
      initialized: false,
      initializing: false,
      error: errorMessage
    });
    return false;
  }
}

/**
 * Speak text using Sherpa TTS
 */
export async function speak(
  text: string,
  options?: {
    speed?: number;
    speakerId?: number;
    onStart?: () => void;
    onDone?: () => void;
    onError?: (error: Error) => void;
  }
): Promise<void> {
  if (!NATIVE_MODULE_AVAILABLE || !currentState.initialized) {
    options?.onError?.(new Error('Sherpa TTS not initialized'));
    return;
  }

  const speed = options?.speed ?? 1.0;
  const speakerId = options?.speakerId ?? 0;

  try {
    options?.onStart?.();

    await TTSManager.generateAndPlay(text, speakerId, speed);

    options?.onDone?.();
  } catch (error) {
    const err = error instanceof Error ? error : new Error('TTS playback failed');
    console.error('[SherpaTTS] Speak error:', err);
    options?.onError?.(err);
  }
}

/**
 * Stop current speech playback
 */
export function stop(): void {
  if (NATIVE_MODULE_AVAILABLE && currentState.initialized) {
    try {
      TTSManager.deinitialize();
      // Re-initialize after stop to be ready for next speech
      setTimeout(() => {
        if (currentState.currentVoice) {
          initializeSherpaTTS(currentState.currentVoice);
        }
      }, 100);
    } catch (error) {
      console.error('[SherpaTTS] Stop error:', error);
    }
  }
}

/**
 * Deinitialize and free resources
 */
export function deinitialize(): void {
  if (NATIVE_MODULE_AVAILABLE) {
    try {
      TTSManager.deinitialize();
    } catch (error) {
      console.error('[SherpaTTS] Deinitialize error:', error);
    }
  }
  updateState({ initialized: false, error: null });
}

/**
 * Check if Sherpa TTS is available (native module loaded)
 */
export function isAvailable(): boolean {
  return NATIVE_MODULE_AVAILABLE && currentState.initialized;
}

/**
 * Get current voice
 */
export function getCurrentVoice(): PiperVoice {
  return currentState.currentVoice;
}

/**
 * Set and switch to a different voice
 */
export async function setVoice(voiceId: string): Promise<boolean> {
  const voice = AVAILABLE_VOICES.find(v => v.id === voiceId);
  if (!voice) {
    console.error('[SherpaTTS] Voice not found:', voiceId);
    return false;
  }

  // Deinitialize current and reinitialize with new voice
  deinitialize();
  return initializeSherpaTTS(voice);
}

/**
 * Add listener for volume updates during playback
 */
export function addVolumeListener(callback: (volume: number) => void): (() => void) | null {
  if (!ttsManagerEmitter) return null;

  const subscription = ttsManagerEmitter.addListener('VolumeUpdate', (event) => {
    callback(event.volume);
  });

  return () => subscription.remove();
}

/**
 * Load saved voice preference
 */
export async function loadSavedVoice(): Promise<PiperVoice> {
  // For now, just return default voice
  // Could add AsyncStorage persistence later
  return AVAILABLE_VOICES[0];
}

export default {
  initialize: initializeSherpaTTS,
  speak,
  stop,
  deinitialize,
  isAvailable,
  addStateListener,
  getState,
  getVoices: () => AVAILABLE_VOICES,
  getCurrentVoice,
  setVoice,
  loadSavedVoice,
  addVolumeListener,
};