ForcePilot/backend/package/yuxi/channels/adapters/feishu/audio_transcode.py

from __future__ import annotations

import asyncio
import logging
import os
import subprocess
import tempfile
from dataclasses import dataclass

logger = logging.getLogger(__name__)

DEFAULT_FFMPEG_TIMEOUT_S = 120.0
DEFAULT_MAX_DURATION_SECS = float(os.environ.get("MEDIA_FFMPEG_MAX_AUDIO_DURATION_SECS", "300"))
FEISHU_OGG_OPUS_SAMPLE_RATE = 16000
FEISHU_OGG_OPUS_BITRATE = "24k"
FEISHU_OGG_OPUS_CHANNELS = 1

SUPPORTED_INPUT_FORMATS = {
    ".mp3",
    ".wav",
    ".m4a",
    ".aac",
    ".flac",
    ".ogg",
    ".opus",
    ".wma",
    ".aiff",
    ".aif",
    ".alac",
    ".ape",
    ".webm",
    ".mp4",
    ".mov",
    ".avi",
    ".mkv",
    ".3gp",
}


@dataclass
class TranscodeResult:
    success: bool
    data: bytes | None = None
    duration_s: float = 0.0
    error: str = ""


def _ffmpeg_available() -> bool:
    try:
        result = subprocess.run(
            ["ffmpeg", "-version"],
            capture_output=True,
            timeout=5,
        )
        return result.returncode == 0
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return False


async def probe_audio_duration(file_data: bytes) -> float:
    if not _ffmpeg_available():
        return 0.0

    loop = asyncio.get_running_loop()

    def _probe() -> float:
        with tempfile.NamedTemporaryFile(suffix=".tmp", delete=False) as tmp:
            tmp.write(file_data)
            tmp_path = tmp.name

        try:
            result = subprocess.run(
                [
                    "ffprobe",
                    "-v",
                    "error",
                    "-show_entries",
                    "format=duration",
                    "-of",
                    "default=noprint_wrappers=1:nokey=1",
                    tmp_path,
                ],
                capture_output=True,
                text=True,
                timeout=30,
            )
            if result.returncode == 0 and result.stdout.strip():
                return float(result.stdout.strip())
        except (ValueError, subprocess.TimeoutExpired, OSError):
            pass
        finally:
            try:
                os.unlink(tmp_path)
            except OSError:
                pass
        return 0.0

    return await loop.run_in_executor(None, _probe)


async def transcode_to_ogg_opus(
    input_data: bytes,
    *,
    sample_rate: int = FEISHU_OGG_OPUS_SAMPLE_RATE,
    bitrate: str = FEISHU_OGG_OPUS_BITRATE,
    channels: int = FEISHU_OGG_OPUS_CHANNELS,
    timeout: float = DEFAULT_FFMPEG_TIMEOUT_S,
    max_duration_s: float = DEFAULT_MAX_DURATION_SECS,
) -> TranscodeResult:
    if not _ffmpeg_available():
        return TranscodeResult(success=False, error="ffmpeg not available")

    duration = await probe_audio_duration(input_data)
    if duration > max_duration_s:
        return TranscodeResult(
            success=False,
            duration_s=duration,
            error=f"Audio duration {duration:.1f}s exceeds maximum {max_duration_s:.0f}s",
        )

    loop = asyncio.get_running_event_loop()

    def _transcode() -> TranscodeResult:
        with tempfile.NamedTemporaryFile(suffix=".tmp", delete=False) as in_tmp:
            in_tmp.write(input_data)
            input_path = in_tmp.name

        output_fd, output_path = tempfile.mkstemp(suffix=".ogg")
        os.close(output_fd)

        try:
            cmd = [
                "ffmpeg",
                "-y",
                "-i",
                input_path,
                "-c:a",
                "libopus",
                "-ar",
                str(sample_rate),
                "-b:a",
                bitrate,
                "-ac",
                str(channels),
                "-f",
                "ogg",
                "-map_metadata",
                "-1",
                output_path,
            ]

            result = subprocess.run(
                cmd,
                capture_output=True,
                timeout=timeout,
            )

            if result.returncode != 0:
                stderr = (result.stderr or b"").decode("utf-8", errors="replace")[-500:]
                return TranscodeResult(success=False, duration_s=duration, error=f"ffmpeg failed: {stderr}")

            with open(output_path, "rb") as f:
                output_data = f.read()

            if not output_data:
                return TranscodeResult(success=False, duration_s=duration, error="Transcode produced empty output")

            return TranscodeResult(success=True, data=output_data, duration_s=duration)
        except subprocess.TimeoutExpired:
            return TranscodeResult(success=False, duration_s=duration, error="FFmpeg transcode timed out")
        except Exception as e:
            return TranscodeResult(success=False, duration_s=duration, error=str(e))
        finally:
            for path in (input_path, output_path):
                try:
                    os.unlink(path)
                except OSError:
                    pass

    return await loop.run_in_executor(None, _transcode)


def guess_audio_input_format(filename: str, mime_type: str = "") -> str | None:
    if filename:
        ext = os.path.splitext(filename)[1].lower()
        if ext in SUPPORTED_INPUT_FORMATS:
            return ext
    if mime_type:
        type_map = {
            "audio/mpeg": ".mp3",
            "audio/mp3": ".mp3",
            "audio/wav": ".wav",
            "audio/x-wav": ".wav",
            "audio/mp4": ".m4a",
            "audio/aac": ".aac",
            "audio/flac": ".flac",
            "audio/ogg": ".ogg",
            "audio/opus": ".opus",
            "audio/x-ms-wma": ".wma",
            "audio/aiff": ".aiff",
            "video/mp4": ".mp4",
            "video/webm": ".webm",
        }
        return type_map.get(mime_type)
    return None


async def detect_and_transcode(
    input_data: bytes,
    *,
    filename: str = "",
    mime_type: str = "",
    max_duration_s: float = DEFAULT_MAX_DURATION_SECS,
) -> TranscodeResult:
    fmt = guess_audio_input_format(filename, mime_type)
    if fmt is None:
        return TranscodeResult(success=False, error=f"Unsupported audio format: filename={filename}, mime={mime_type}")

    if fmt in (".ogg", ".opus"):
        duration = await probe_audio_duration(input_data)
        if duration > max_duration_s:
            return TranscodeResult(
                success=False,
                duration_s=duration,
                error=f"Audio duration {duration:.1f}s exceeds maximum {max_duration_s:.0f}s",
            )
        return TranscodeResult(success=True, data=input_data, duration_s=duration)

    return await transcode_to_ogg_opus(input_data, max_duration_s=max_duration_s)