src/audio.py

import os
import time
import tempfile
import warnings
import asyncio
import edge_tts
from transformers import pipeline

# -------------------------------------------------------------------------
# Speech-to-Text (STT): Whisper-small running locally on CPU
# - ~242MB model, good Hindi accuracy, no API dependency
# -------------------------------------------------------------------------
_stt_pipeline = None


def _get_stt_pipeline():
    """Lazily loads the Whisper-small pipeline (only once on first call)."""
    global _stt_pipeline
    if _stt_pipeline is None:
        print("Loading Whisper-small model for STT (this takes ~30s on first run)...")
        # Suppress the duplicate logits-processor warning that appears
        # when language/task are passed alongside the model's own processor setup.
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message="A custom logits processor of type",
                category=UserWarning,
            )
            _stt_pipeline = pipeline(
                task="automatic-speech-recognition",
                model="openai/whisper-small",
                device="cpu",
                generate_kwargs={"language": "hindi", "task": "transcribe"},
            )
        print("Whisper-small ready.")
    return _stt_pipeline


def speech_to_text(audio_filepath: str) -> str:
    """
    Transcribes a Hindi audio recording into Hindi text (Devanagari).

    Args:
        audio_filepath: Path to the audio file (wav/mp3/webm/ogg).
    Returns:
        Transcribed Hindi string, or an error message beginning with 'Error:'.
    """
    if not audio_filepath:
        return ""
    if not os.path.exists(audio_filepath):
        return f"Error: audio file not found at {audio_filepath}"

    try:
        pipe = _get_stt_pipeline()
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message="A custom logits processor of type",
                category=UserWarning,
            )
            result = pipe(audio_filepath)
        transcript = result.get("text", "").strip()
        if not transcript:
            return "Error: Could not transcribe audio - please speak clearly and try again."
        return transcript
    except Exception as exc:
        print(f"[STT Error] {exc}")
        return f"Error: Could not transcribe audio ({exc})"


# -------------------------------------------------------------------------
# Text-to-Speech (TTS): Edge-TTS – high-quality neural Hindi voice
# - Saves to /tmp/ so it works on read-only deployments like HF Spaces
# -------------------------------------------------------------------------
async def _async_tts(text: str, output_path: str):
    """Async helper to run the Edge TTS communication."""
    communicate = edge_tts.Communicate(text, "hi-IN-SwaraNeural")
    await communicate.save(output_path)


def text_to_speech(text: str) -> str | None:
    """
    Converts Hindi text to a high-quality spoken MP3 file using edge-tts.

    Saves into /tmp/ to ensure write access regardless of the deployment
    environment (HF Spaces /app directory is read-only).

    Args:
        text: Hindi text in Devanagari.
    Returns:
        Absolute path to the saved MP3 file, or None on failure.
    """
    if not text or text.startswith("Error:"):
        return None
    
    max_retries = 3
    for attempt in range(max_retries):
        try:
            # Use a named temp file so Gradio can serve it correctly
            tmp = tempfile.NamedTemporaryFile(
                suffix=".mp3", delete=False, dir="/tmp"
            )
            output_path = tmp.name
            tmp.close()
            
            # Run the async edge-tts code synchronously
            asyncio.run(_async_tts(text, output_path))
            return output_path
            
        except Exception as exc:
            err = str(exc)
            if "429" in err and attempt < max_retries - 1:
                wait = 2 ** attempt
                print(f"[TTS] Rate-limited, retrying in {wait}s... (attempt {attempt + 1}/{max_retries})")
                time.sleep(wait)
            else:
                print(f"[TTS Error] {exc}")
                return None