ERP/bot/app/audio_enhancer.py

"""Audio preprocessing & denoising pipeline for procurement bot.

Workflow:
  Raw Voice Bytes (Bale)
    → enhance_voice() [Optional Audio Enhancement / Denoising API]
        - If disabled or missing URL → zero overhead pass-through
        - Zero transcoding by default (raw bytes sent as-is)
        - Transcoding to target format ONLY if explicitly configured
        - Graceful fail-soft on timeout or error (transparent raw audio fallback)
        - Enhanced bytes saved to data/voices/{YYYYMMDD}/{trace_id}_{message_id}_enhanced.{ext}
        - Observability events recorded (audio.enhance.result with ok/size/latency_ms/path)
    → Multi-STT Providers
"""
import asyncio
import logging
import subprocess
import tempfile
import time
from pathlib import Path

import httpx

from app import observability as obs
from app.config import Settings

logger = logging.getLogger("bot.audio_enhancer")

_CONVERT_TIMEOUT = 60


def _transcode(raw: bytes, src_ext: str, target_fmt: str) -> bytes:
    """Convert audio to target format using ffmpeg only when explicitly requested."""
    target_fmt = target_fmt.lstrip(".")
    with tempfile.TemporaryDirectory() as tmp:
        src = Path(tmp) / f"in{src_ext or '.ogg'}"
        dst = Path(tmp) / f"out.{target_fmt}"
        src.write_bytes(raw)
        proc = subprocess.run(
            ["ffmpeg", "-y", "-i", str(src), "-vn", str(dst)],
            capture_output=True,
            timeout=_CONVERT_TIMEOUT,
        )
        if proc.returncode != 0 or not dst.exists():
            raise RuntimeError(f"ffmpeg conversion to {target_fmt} failed: {proc.stderr[-300:]!r}")
        return dst.read_bytes()


def _resolve_extension(filename: str | None, mime_type: str | None) -> str:
    ext = Path(filename or "").suffix
    if not ext and mime_type:
        clean_mime = mime_type.split(";")[0].strip().lower()
        ext = {
            "audio/ogg": ".ogg",
            "audio/opus": ".ogg",
            "audio/mpeg": ".mp3",
            "audio/mp4": ".m4a",
            "audio/wav": ".wav",
            "audio/x-wav": ".wav",
            "audio/webm": ".webm",
            "audio/flac": ".flac",
        }.get(clean_mime, ".ogg")
    return ext or ".ogg"


async def enhance_voice(
    settings: Settings,
    audio_bytes: bytes,
    mime_type: str | None = None,
    filename: str | None = None,
    trace_id: str | None = None,
    message_id: int | None = None,
) -> tuple[bytes, dict]:
    """Enhance and denoise voice audio before STT.

    Returns:
        (processed_audio_bytes, metadata)
    Guarantees:
        - Never raises an unhandled exception.
        - On any failure or timeout, gracefully returns original audio_bytes (fallback).
        - Enhanced bytes are saved to daily folder with _enhanced suffix when successful.
        - Observability events are captured.
    """
    if not settings.AUDIO_ENHANCE_ENABLED or not (settings.AUDIO_ENHANCE_BASE_URL or "").strip():
        return audio_bytes, {"enhanced": False, "skipped": True}

    base_url = settings.AUDIO_ENHANCE_BASE_URL.rstrip("/")
    api_key = settings.AUDIO_ENHANCE_API_KEY or ""
    model = settings.AUDIO_ENHANCE_MODEL or ""
    timeout = float(settings.AUDIO_ENHANCE_TIMEOUT_SECONDS or 60.0)
    target_format = (settings.AUDIO_ENHANCE_TARGET_FORMAT or "").strip().lower()
    method = (settings.AUDIO_ENHANCE_METHOD or "auto").strip().lower()

    orig_ext = _resolve_extension(filename, mime_type)
    t0 = time.perf_counter()

    # Condition 2: Zero transcoding unless explicitly requested by configuration
    payload_bytes = audio_bytes
    upload_ext = orig_ext
    if target_format and not orig_ext.endswith(f".{target_format}"):
        try:
            payload_bytes = await asyncio.to_thread(_transcode, audio_bytes, orig_ext, target_format)
            upload_ext = f".{target_format}"
        except Exception as exc:
            logger.warning("Target format conversion failed: %s; using raw bytes", exc)
            payload_bytes = audio_bytes
            upload_ext = orig_ext

    headers: dict[str, str] = {}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            upload_filename = f"voice{upload_ext}"
            req_mime = mime_type or "audio/ogg"

            if method == "raw_stream":
                headers["Content-Type"] = req_mime
                resp = await client.post(base_url, content=payload_bytes, headers=headers)
            else:
                # Default / multipart upload (compatible with standard voice enhancement APIs)
                data = {"model": model} if model else {}
                files = {"file": (upload_filename, payload_bytes, req_mime)}
                resp = await client.post(base_url, files=files, data=data, headers=headers)

        ms = int((time.perf_counter() - t0) * 1000)

        if resp.status_code == 200 and len(resp.content) > 0:
            enhanced_bytes = resp.content

            # Determine response audio extension
            resp_ct = (resp.headers.get("content-type") or "").split(";")[0].strip().lower()
            out_ext = {
                "audio/ogg": ".ogg",
                "audio/opus": ".ogg",
                "audio/mpeg": ".mp3",
                "audio/mp4": ".m4a",
                "audio/wav": ".wav",
                "audio/x-wav": ".wav",
            }.get(resp_ct, upload_ext)

            saved_path = obs.save_enhanced_voice(
                trace_id, message_id or 0, enhanced_bytes, ext=out_ext
            )

            await obs.event(
                trace_id,
                "audio.enhance.result",
                ok=True,
                latency_ms=ms,
                model=model,
                payload={
                    "path": saved_path,
                    "size": len(enhanced_bytes),
                    "original_size": len(audio_bytes),
                    "provider": model or "audio_enhancer",
                },
            )
            logger.info(
                "Audio enhancement succeeded in %d ms (size: %d -> %d bytes)",
                ms,
                len(audio_bytes),
                len(enhanced_bytes),
            )
            return enhanced_bytes, {
                "enhanced": True,
                "path": saved_path,
                "size": len(enhanced_bytes),
                "latency_ms": ms,
            }

        # Non-200 or empty response -> graceful fallback
        err_msg = f"HTTP {resp.status_code}: {resp.text[:200]}"
        logger.warning("Audio enhancement failed: %s; falling back to raw audio", err_msg)
        await obs.event(
            trace_id,
            "audio.enhance.result",
            ok=False,
            latency_ms=ms,
            model=model,
            payload={"error": err_msg, "fallback": True},
        )
        return audio_bytes, {"enhanced": False, "fallback": True, "error": err_msg}

    except Exception as exc:
        ms = int((time.perf_counter() - t0) * 1000)
        err_msg = str(exc)[:300]
        logger.warning("Audio enhancement exception: %s; falling back to raw audio", err_msg)
        await obs.event(
            trace_id,
            "audio.enhance.result",
            ok=False,
            latency_ms=ms,
            model=model,
            payload={"error": err_msg, "fallback": True},
        )
        return audio_bytes, {"enhanced": False, "fallback": True, "error": err_msg}