WellNuo/services/sherpaTTS.ts

/**
 * Sherpa TTS Service - Offline Neural Text-to-Speech
 * Uses Piper VITS models for high-quality offline speech synthesis
 * Model is bundled with the app (no download needed)
 */

import TTSManager from 'react-native-sherpa-onnx-offline-tts';
import RNFS from 'react-native-fs';
import { Platform } from 'react-native';
import AsyncStorage from '@react-native-async-storage/async-storage';

// 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',
    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',
  },
];

const VOICE_STORAGE_KEY = '@wellnuo_selected_voice';
let currentVoice: PiperVoice = AVAILABLE_VOICES[0]; // Default to Lessac

interface TTSConfig {
  modelPath: string;
  tokensPath: string;
  dataDirPath: string;
}

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

let isInitialized = false;
let currentState: SherpaTTSState = {
  initialized: false,
  initializing: false,
  error: null,
};

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

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 };
}

/**
 * Get the path to bundled model files for a specific voice
 * On iOS, assets are in the main bundle
 */
function getBundledModelPath(voice: PiperVoice = currentVoice): string {
  if (Platform.OS === 'ios') {
    // On iOS, bundled assets are in MainBundle
    return `${RNFS.MainBundlePath}/assets/tts-models/${voice.modelDir}`;
  } else {
    // On Android, assets need to be copied from APK to filesystem first
    return `${RNFS.DocumentDirectoryPath}/tts-models/${voice.modelDir}`;
  }
}

/**
 * Load saved voice preference from storage
 */
export async function loadSavedVoice(): Promise<PiperVoice> {
  try {
    const savedVoiceId = await AsyncStorage.getItem(VOICE_STORAGE_KEY);
    if (savedVoiceId) {
      const voice = AVAILABLE_VOICES.find(v => v.id === savedVoiceId);
      if (voice) {
        currentVoice = voice;
        return voice;
      }
    }
  } catch (error) {
    console.log('[SherpaTTS] Failed to load saved voice:', error);
  }
  return currentVoice;
}

/**
 * Save voice preference to storage
 */
async function saveVoicePreference(voiceId: string): Promise<void> {
  try {
    await AsyncStorage.setItem(VOICE_STORAGE_KEY, voiceId);
  } catch (error) {
    console.log('[SherpaTTS] Failed to save voice preference:', error);
  }
}

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

/**
 * Change to a different voice (requires re-initialization)
 */
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;
  }

  // If same voice, no need to reinitialize
  if (voice.id === currentVoice.id && isInitialized) {
    return true;
  }

  // Deinitialize current voice
  if (isInitialized) {
    deinitialize();
  }

  // Switch to new voice
  currentVoice = voice;
  await saveVoicePreference(voiceId);

  // Initialize with new voice
  return await initializeSherpaTTS();
}

/**
 * Copy bundled model to document directory (needed for Android)
 */
async function ensureModelAvailable(voice: PiperVoice = currentVoice): Promise<string> {
  const modelDir = getBundledModelPath(voice);

  // Check if model exists at expected location
  const onnxPath = `${modelDir}/${voice.onnxFile}`;
  const exists = await RNFS.exists(onnxPath);

  if (exists) {
    console.log('[SherpaTTS] Model found at:', modelDir);
    return modelDir;
  }

  // If not found, try alternative paths for development
  const altPaths = [
    `${RNFS.MainBundlePath}/tts-models/${voice.modelDir}`,
    `${RNFS.DocumentDirectoryPath}/tts-models/${voice.modelDir}`,
  ];

  for (const altPath of altPaths) {
    const altOnnxPath = `${altPath}/${voice.onnxFile}`;
    if (await RNFS.exists(altOnnxPath)) {
      console.log('[SherpaTTS] Model found at alternative path:', altPath);
      return altPath;
    }
  }

  throw new Error(`TTS model for ${voice.name} not found. Please ensure the model is bundled with the app.`);
}

/**
 * Initialize Sherpa TTS with bundled Piper voice model
 */
export async function initializeSherpaTTS(): Promise<boolean> {
  if (isInitialized) {
    console.log('[SherpaTTS] Already initialized');
    return true;
  }

  try {
    currentState = { ...currentState, initializing: true, error: null };
    notifyListeners();

    // Load saved voice preference
    await loadSavedVoice();
    console.log('[SherpaTTS] Using voice:', currentVoice.name);

    // Get model directory
    const modelDir = await ensureModelAvailable(currentVoice);

    // Find the ONNX model file
    const files = await RNFS.readDir(modelDir);
    console.log('[SherpaTTS] Model files:', files.map(f => f.name));

    const onnxFile = files.find(f => f.name.endsWith('.onnx'));
    const tokensFile = files.find(f => f.name === 'tokens.txt');
    const espeakDir = files.find(f => f.name === 'espeak-ng-data');

    if (!onnxFile || !tokensFile || !espeakDir) {
      throw new Error('Model files incomplete. Missing: ' +
        [!onnxFile && 'onnx', !tokensFile && 'tokens', !espeakDir && 'espeak-ng-data']
          .filter(Boolean).join(', '));
    }

    // Configure TTS engine
    const config: TTSConfig = {
      modelPath: onnxFile.path,
      tokensPath: tokensFile.path,
      dataDirPath: espeakDir.path,
    };

    console.log('[SherpaTTS] Initializing with config:', config);

    // Initialize Sherpa TTS
    await TTSManager.initialize(JSON.stringify(config));

    isInitialized = true;
    currentState = {
      initialized: true,
      initializing: false,
      error: null,
    };
    notifyListeners();

    console.log('[SherpaTTS] Initialization complete!');
    return true;
  } catch (error) {
    const errorMessage = error instanceof Error ? error.message : 'Unknown error';
    console.error('[SherpaTTS] Initialization failed:', errorMessage);

    currentState = {
      initialized: false,
      initializing: false,
      error: errorMessage,
    };
    notifyListeners();

    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 (!isInitialized) {
    console.warn('[SherpaTTS] Not initialized, falling back to system TTS');
    options?.onError?.(new Error('Sherpa TTS not initialized'));
    return;
  }

  const { speed = 1.0, speakerId = 0, onStart, onDone, onError } = options || {};

  try {
    onStart?.();
    await TTSManager.generateAndPlay(text, speakerId, speed);
    onDone?.();
  } catch (error) {
    console.error('[SherpaTTS] Speak error:', error);
    onError?.(error instanceof Error ? error : new Error('TTS failed'));
  }
}

/**
 * Stop current speech
 */
export function stop(): void {
  if (isInitialized) {
    TTSManager.stopPlaying();
  }
}

/**
 * Deinitialize TTS engine
 */
export function deinitialize(): void {
  if (isInitialized) {
    TTSManager.deinitialize();
    isInitialized = false;
    currentState = {
      initialized: false,
      initializing: false,
      error: null,
    };
    notifyListeners();
  }
}

/**
 * Check if Sherpa TTS is available and initialized
 */
export function isAvailable(): boolean {
  return isInitialized;
}

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