ForcePilot/backend/package/yuxi/channel/extensions/qqbot/audio.py

from __future__ import annotations

import asyncio
import logging
import subprocess
import tempfile
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


class QQBotAudio:
    def __init__(self, ffmpeg_path: str = "ffmpeg"):
        self._ffmpeg = ffmpeg_path
        self._audio_format_policy = {
            "stt_direct_formats": ["wav", "pcm"],
            "upload_direct_formats": ["silk", "amr"],
            "transcode_enabled": True,
        }

    async def silk_to_wav(self, silk_data: bytes) -> bytes:
        with tempfile.NamedTemporaryFile(suffix=".silk", delete=False) as silk_file:
            silk_file.write(silk_data)
            silk_path = silk_file.name

        wav_path = silk_path + ".wav"

        try:
            proc = await asyncio.create_subprocess_exec(
                self._ffmpeg,
                "-y",
                "-f",
                "silk",
                "-i",
                silk_path,
                "-acodec",
                "pcm_s16le",
                "-ar",
                "24000",
                "-ac",
                "1",
                wav_path,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            await proc.communicate()

            if proc.returncode != 0:
                logger.warning("ffmpeg SILK→WAV conversion failed, returncode=%d", proc.returncode)
                return silk_data

            wav_data = Path(wav_path).read_bytes()
            return wav_data
        except FileNotFoundError:
            logger.warning("ffmpeg not found, returning raw data")
            return silk_data
        finally:
            Path(silk_path).unlink(missing_ok=True)
            Path(wav_path).unlink(missing_ok=True)

    async def wav_to_silk(self, wav_data: bytes) -> bytes:
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as wav_file:
            wav_file.write(wav_data)
            wav_path = wav_file.name

        silk_path = wav_path + ".silk"

        try:
            proc = await asyncio.create_subprocess_exec(
                self._ffmpeg,
                "-y",
                "-i",
                wav_path,
                "-acodec",
                "silk",
                "-ar",
                "24000",
                "-ac",
                "1",
                silk_path,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            await proc.communicate()

            if proc.returncode != 0:
                logger.warning("ffmpeg WAV→SILK conversion failed, returncode=%d", proc.returncode)
                return wav_data

            silk_data = Path(silk_path).read_bytes()
            return silk_data
        except FileNotFoundError:
            logger.warning("ffmpeg not found, returning raw data")
            return wav_data
        finally:
            Path(wav_path).unlink(missing_ok=True)
            Path(silk_path).unlink(missing_ok=True)

    def is_silk(self, data: bytes) -> bool:
        return data[:2] == b"#!"

    def is_direct_upload_format(self, ext: str) -> bool:
        return ext.lower().lstrip(".") in self._audio_format_policy["upload_direct_formats"]

    def can_transcode(self) -> bool:
        try:
            subprocess.run([self._ffmpeg, "-version"], capture_output=True, check=False)
            return True
        except FileNotFoundError:
            return False