ForcePilot/backend/package/yuxi/channels/adapters/bluebubbles/voice.py

from __future__ import annotations

import importlib.util
import os
import shutil
import subprocess
import tempfile
from typing import Any, Literal

from yuxi.channels.adapters.bluebubbles.client import BlueBubblesClient
from yuxi.channels.adapters.bluebubbles.send import send_attachment_bytes
from yuxi.utils.logging_config import logger

AudioFormat = Literal["ogg", "caf", "mp3"]
AUDIO_FORMAT_MIME_MAP: dict[str, str] = {
    "ogg": "audio/ogg",
    "caf": "audio/x-caf",
    "mp3": "audio/mpeg",
}
SYNTHESIS_TARGET = "audio-file"
PREFERRED_AUDIO_FORMAT: AudioFormat = "caf"
SUPPORTED_AUDIO_FILE_FORMATS: list[str] = ["mp3", "caf", "audio/mpeg", "audio/x-caf"]


async def send_tts_voice(
    client: BlueBubblesClient,
    chat_guid: str,
    text: str,
    reply_to_guid: str | None = None,
    fmt: AudioFormat = "caf",
) -> dict[str, Any]:
    audio_data = await _generate_tts_audio(text, fmt)
    if not audio_data:
        return {"success": False, "error": "TTS generation failed"}
    extension_map: dict[str, str] = {"ogg": "voice_message.ogg", "caf": "voice_message.caf", "mp3": "voice_message.mp3"}
    filename = extension_map.get(fmt, "voice_message.caf")
    return await send_attachment_bytes(
        client,
        chat_guid,
        audio_data,
        filename,
        caption="",
        reply_to_guid=reply_to_guid,
    )


async def _generate_tts_audio(text: str, fmt: AudioFormat = "caf") -> bytes:
    try:
        with tempfile.NamedTemporaryFile(suffix=".ogg", delete=False) as tmp:
            tmp_path = tmp.name

        try:
            result = subprocess.run(
                ["espeak", "-w", tmp_path, text],
                capture_output=True,
                timeout=30,
            )
            if result.returncode != 0:
                raise RuntimeError(f"espeak failed: {result.stderr.decode()}")

            if fmt == "ogg":
                with open(tmp_path, "rb") as f:
                    return f.read()

            return await _transcode_to_format(tmp_path, fmt)
        finally:
            os.unlink(tmp_path)
    except FileNotFoundError as e:
        missing = str(e).split()[-1] if str(e) else "espeak"
        logger.warning("[BlueBubbles] %s not found, TTS requires espeak and a transcoder to be installed", missing)
        return b""
    except Exception as e:
        logger.error(f"[BlueBubbles] TTS generation failed: {e}")
        return b""


async def _transcode_to_format(input_path: str, fmt: AudioFormat) -> bytes:
    import asyncio
    import importlib.util

    if importlib.util.find_spec("pydub") is not None:
        from pydub import AudioSegment

        return await asyncio.to_thread(_transcode_pydub, input_path, fmt, AudioSegment)

    logger.warning("[BlueBubbles] pydub not available, falling back to ffmpeg")
    return await _transcode_ffmpeg(input_path, fmt)


def _transcode_pydub(input_path: str, fmt: AudioFormat, AudioSegment) -> bytes:
    audio = AudioSegment.from_file(input_path, format="ogg")
    fmt_map: dict[str, str] = {"caf": "caf", "mp3": "mp3"}
    out_fmt = fmt_map.get(fmt, "caf")

    with tempfile.NamedTemporaryFile(suffix=f".{out_fmt}", delete=False) as out_tmp:
        out_path = out_tmp.name

    try:
        audio.export(out_path, format=out_fmt)
        with open(out_path, "rb") as f:
            return f.read()
    finally:
        try:
            os.unlink(out_path)
        except OSError:
            pass


def check_tts_available() -> dict[str, bool]:
    return {
        "espeak": shutil.which("espeak") is not None,
        "ffmpeg": shutil.which("ffmpeg") is not None,
        "pydub": importlib.util.find_spec("pydub") is not None,
    }


async def _transcode_ffmpeg(input_path: str, fmt: AudioFormat) -> bytes:
    import asyncio

    fmt_map: dict[str, str] = {
        "caf": "caf",
        "mp3": "mp3",
    }
    out_fmt = fmt_map.get(fmt, "caf")

    with tempfile.NamedTemporaryFile(suffix=f".{out_fmt}", delete=False) as out_tmp:
        out_path = out_tmp.name

    try:
        proc = await asyncio.create_subprocess_exec(
            "ffmpeg",
            "-y",
            "-i",
            input_path,
            "-acodec",
            "pcm_s16le" if out_fmt == "caf" else "libmp3lame",
            "-ar",
            "44100" if out_fmt == "caf" else "24000",
            out_path,
            stdout=asyncio.subprocess.DEVNULL,
            stderr=asyncio.subprocess.DEVNULL,
        )
        await proc.wait()
        if proc.returncode != 0:
            raise RuntimeError("ffmpeg transcoding failed")

        with open(out_path, "rb") as f:
            return f.read()
    except FileNotFoundError:
        logger.warning("[BlueBubbles] ffmpeg not found, cannot transcode")
        with open(input_path, "rb") as f:
            return f.read()
    finally:
        try:
            os.unlink(out_path)
        except OSError:
            pass