umb/bot/utils/voice.py

import asyncio
import logging
import os
import re
import shutil
import time
from pathlib import Path

from faster_whisper import WhisperModel
from aiogram import Bot

from config import PROXY_ENABLED, PROXY_URL, WHISPER_MODEL_SIZE, WHISPER_MIN_FREE_SPACE_BYTES

logger = logging.getLogger(__name__)

_model: WhisperModel | None = None
_model_lock = asyncio.Lock()

BASE_DIR = Path(__file__).resolve().parent.parent.parent
MODELS_DIR = BASE_DIR / "models" / "faster-whisper"
VOICE_DIR = BASE_DIR / "bot" / "data" / "voice"


def _setup_proxy() -> None:
    if PROXY_ENABLED and PROXY_URL:
        os.environ["ALL_PROXY"] = PROXY_URL
        logger.info("Proxy set for model download: %s...", PROXY_URL[:30])


def _get_model_path() -> Path:
    return MODELS_DIR / WHISPER_MODEL_SIZE


def _is_model_downloaded() -> bool:
    """Check whether the Whisper model files are already present locally."""
    model_path = _get_model_path()
    model_bin = model_path / "model.bin"
    return model_path.is_dir() and model_bin.is_file()


def _get_free_disk_space_bytes(path: Path) -> int:
    """Return free disk space in bytes for the filesystem containing path."""
    path.mkdir(parents=True, exist_ok=True)
    return shutil.disk_usage(path).free


def _load_model() -> WhisperModel:
    """Synchronous model load/download. Must run in a thread."""
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    VOICE_DIR.mkdir(parents=True, exist_ok=True)

    model_path = _get_model_path()
    model_bin = model_path / "model.bin"

    if not model_path.is_dir() or not model_bin.is_file():
        _setup_proxy()
        logger.info(
            "Whisper model '%s' not found locally. Starting download to %s...",
            WHISPER_MODEL_SIZE,
            MODELS_DIR,
        )
        start = time.monotonic()
        model = WhisperModel(
            WHISPER_MODEL_SIZE,
            device="cpu",
            cpu_threads=4,
            compute_type="int8",
            download_root=str(MODELS_DIR),
        )
        elapsed = time.monotonic() - start
        logger.info(
            "Whisper model '%s' downloaded and loaded in %.1fs",
            WHISPER_MODEL_SIZE,
            elapsed,
        )
        return model

    logger.info(
        "Loading Whisper model '%s' from %s...",
        WHISPER_MODEL_SIZE,
        model_path,
    )
    return WhisperModel(
        str(model_path),
        device="cpu",
        cpu_threads=4,
        compute_type="int8",
    )


async def preload_model() -> None:
    """Download/load the Whisper model at bot startup if it is not already loaded."""
    global _model
    async with _model_lock:
        if _model is not None:
            logger.info("Whisper model is already loaded, skipping preload")
            return

        if _is_model_downloaded():
            logger.info(
                "Whisper model '%s' found locally at %s",
                WHISPER_MODEL_SIZE,
                _get_model_path(),
            )
        else:
            free_space = _get_free_disk_space_bytes(MODELS_DIR)
            required_space = WHISPER_MIN_FREE_SPACE_BYTES
            logger.info(
                "Free disk space: %.2f GB, required: %.2f GB",
                free_space / (1024 ** 3),
                required_space / (1024 ** 3),
            )
            if free_space < required_space:
                raise RuntimeError(
                    f"Not enough disk space to download Whisper model '{WHISPER_MODEL_SIZE}'. "
                    f"Free: {free_space / (1024 ** 3):.2f} GB, "
                    f"required: {required_space / (1024 ** 3):.2f} GB"
                )
            logger.info(
                "Whisper model '%s' will be downloaded at startup",
                WHISPER_MODEL_SIZE,
            )

        loop = asyncio.get_running_loop()
        _model = await loop.run_in_executor(None, _load_model)
        logger.info("Whisper model '%s' is ready", WHISPER_MODEL_SIZE)


async def _get_model() -> WhisperModel:
    """Thread-safe lazy initializer for the Whisper model."""
    global _model
    if _model is None:
        async with _model_lock:
            if _model is None:
                loop = asyncio.get_running_loop()
                _model = await loop.run_in_executor(None, _load_model)
                logger.info("Whisper model '%s' loaded on demand", WHISPER_MODEL_SIZE)
    return _model


async def download_voice(bot: Bot, file_id: str) -> str:
    VOICE_DIR.mkdir(parents=True, exist_ok=True)
    file = await bot.get_file(file_id)
    safe_id = re.sub(r"[^A-Za-z0-9_-]", "_", file_id)
    path = str(VOICE_DIR / f"{safe_id}.ogg")
    await bot.download_file(file.file_path, destination=path)
    logger.info("Voice downloaded: %s", path)
    return path


async def convert_to_wav(input_path: str) -> str:
    output_path = str(Path(input_path).with_suffix(".wav"))
    proc = await asyncio.create_subprocess_exec(
        "ffmpeg",
        "-i", input_path,
        "-ar", "16000",
        "-ac", "1",
        "-c:a", "pcm_s16le",
        output_path,
        "-y",
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    _, stderr = await proc.communicate()
    if proc.returncode != 0:
        logger.error("ffmpeg error: %s", stderr.decode(errors="replace"))
        raise RuntimeError("ffmpeg conversion failed")
    logger.info("Converted to WAV: %s", output_path)
    return output_path


async def transcribe_audio(file_path: str) -> str:
    model = await _get_model()
    loop = asyncio.get_running_loop()

    def _transcribe() -> str:
        segments, _info = model.transcribe(file_path, language="ru", beam_size=5)
        return " ".join(seg.text for seg in segments).strip()

    text = await loop.run_in_executor(None, _transcribe)
    logger.info("Transcription result (%d chars): %s...", len(text), text[:100])
    return text


def normalize_text(text: str) -> str:
    text = text.strip()
    if not text:
        return text
    text = text[0].upper() + text[1:]
    if text[-1] not in ".!?…":
        text += "."
    text = re.sub(r"\s+", " ", text)
    return text


async def cleanup_files(*paths: str | None) -> None:
    for path in paths:
        if path and os.path.exists(path):
            try:
                await asyncio.to_thread(os.unlink, path)
                logger.debug("Cleaned up: %s", path)
            except Exception as exc:
                logger.warning("Cleanup failed for %s: %s", path, exc)