transcription/run.py

"""CLI entrypoint для транскрибации совещаний."""

import argparse
import os
import sys
from pathlib import Path

from src.audio_utils import check_ffmpeg, is_audio_file, is_video_file
from src.config import get_profile, load_config
from src.document import build_document
from src.pipeline import run_pipeline


def resolve_device(preferred: str) -> str:
    """Определяет доступное устройство."""
    import torch
    if preferred == "cuda" and torch.cuda.is_available():
        return "cuda"
    if preferred == "mps" and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def print_first_run_info():
    """Показывает информацию о первом запуске и скачивании моделей."""
    print("=" * 60)
    print("ПЕРВЫЙ ЗАПУСК: скачивание моделей")
    print("=" * 60)
    print("При первом запуске будут скачаны модели ИИ (~4–5 GB):")
    print("  • Whisper large-v3      ~3.0 GB  (распознавание речи)")
    print("  • Pyannote диаризация   ~0.4 GB  (разделение спикеров)")
    print("  • Wav2Vec2 alignment    ~1.0 GB  (точные таймкоды)")
    print("")
    print("Это займёт время в зависимости от скорости интернета.")
    print("При последующих запусках модели загружаются с диска.")
    print("=" * 60)
    print("")


def parse_formats(fmt_arg: str | None, config_formats: list) -> list[str]:
    """Парсит строку форматов в список."""
    if fmt_arg:
        return [f.strip().lower() for f in fmt_arg.split(",")]
    return config_formats


def main():
    parser = argparse.ArgumentParser(
        description="Транскрибация совещаний с диаризацией и таймкодами. "
                    "Поддерживает аудио (wav, mp3, m4a, ogg, flac) и видео (mp4, webm, avi, mkv, mov, etc.)."
    )
    parser.add_argument("--input", "-i", required=True, help="Путь к аудио или видео файлу")
    parser.add_argument("--output", "-o", default=None, help="Путь к выходному файлу (если указан, переопределяет формат)")
    parser.add_argument("--profile", "-p", default=None, help="Профиль конфигурации (mac_m4, gpu_8gb, cpu_best)")
    parser.add_argument("--config", "-c", default=None, help="Путь к config.yaml")
    parser.add_argument("--device", "-d", default=None, help="Принудительно: cpu/cuda/mps")
    parser.add_argument("--model", "-m", default=None, help="Принудительно: tiny/base/small/medium/large-v3")
    parser.add_argument("--language", "-l", default=None, help="Язык (ru, en, ...)")
    parser.add_argument("--format", "-f", default=None, help="Форматы через запятую: docx,md,txt (по умолчанию из конфига)")
    
    args = parser.parse_args()
    
    input_path = Path(args.input)
    if not input_path.exists():
        print(f"Ошибка: файл не найден: {args.input}", file=sys.stderr)
        sys.exit(1)
    
    # Проверка ffmpeg для видео
    if is_video_file(args.input):
        if not check_ffmpeg():
            print(
                "Ошибка: для обработки видео нужен ffmpeg.\n"
                "Установите ffmpeg:\n"
                "  Mac:    brew install ffmpeg\n"
                "  Linux:  sudo apt-get install ffmpeg\n"
                "  Windows: https://ffmpeg.org/download.html",
                file=sys.stderr,
            )
            sys.exit(1)
        print(f"[Info] Обнаружен видео файл: {input_path.suffix}")
        print("[Info] Аудио будет извлечено автоматически.")
    elif not is_audio_file(args.input):
        print(f"[Warning] Неизвестный формат: {input_path.suffix}. Попытка обработать...")
    
    # Загрузка конфига
    config = load_config(args.config)
    profile = get_profile(config, args.profile)
    
    # Переопределения из CLI
    if args.device:
        profile["device"] = resolve_device(args.device)
    else:
        profile["device"] = resolve_device(profile.get("device", "cpu"))
    if args.model:
        profile["model"] = args.model
    if args.language:
        profile["language"] = args.language
    
    output_cfg = config.get("output", {})
    config_formats = output_cfg.get("formats", ["docx"])
    formats = parse_formats(args.format, config_formats)
    
    # Определение путей выходных файлов
    output_paths: list[str] = []
    if args.output:
        # Если указан --output, используем его для первого формата
        # Остальные форматы — рядом с тем же именем
        base = Path(args.output)
        output_dir = base.parent
        stem = base.stem
        first_ext = base.suffix.lstrip(".")
        # Первый файл с явным путём
        output_paths.append(str(base))
        # Остальные форматы рядом
        for fmt in formats[1:]:
            output_paths.append(str(output_dir / f"{stem}.{fmt}"))
    else:
        stem = input_path.stem
        output_dir = Path(config.get("paths", {}).get("output_dir", "./output"))
        for fmt in formats:
            output_paths.append(str(output_dir / f"{stem}.{fmt}"))
    
    # Проверка HF токена
    hf_token = os.environ.get("HF_TOKEN") or config.get("hf_token")
    if profile.get("diarize", True) and not hf_token:
        print(
            "Ошибка: для диаризации нужен HuggingFace токен.\n"
            "Установите env HF_TOKEN или укажите hf_token в config.yaml\n"
            "Инструкция: https://huggingface.co/docs/hub/security-tokens",
            file=sys.stderr,
        )
        sys.exit(1)
    
    # Информация о первом запуске
    print_first_run_info()
    
    print(f"Профиль:     {args.profile or config.get('active_profile')}")
    print(f"Устройство:  {profile['device']}")
    print(f"Модель:      {profile['model']}")
    print(f"Язык:        {profile['language']}")
    print(f"Вход:        {args.input}")
    print(f"Форматы:     {', '.join(formats)}")
    print(f"Выход:       {', '.join(output_paths)}")
    print("-" * 40)
    
    # Запуск пайплайна (один раз для всех форматов)
    result = run_pipeline(
        input_path=args.input,
        profile_name=args.profile,
        config_path=args.config,
        output_path=output_paths[0] if output_paths else None,
    )
    
    # Генерация документов для всех форматов
    print("-" * 40)
    print("Генерация документов...")
    for out_path in output_paths:
        build_document(result["segments"], out_path, config)
        print(f"  ✓ {out_path}")
    
    print("-" * 40)
    print(f"Готово! Сохранено {len(output_paths)} файл(а):")
    for p in output_paths:
        print(f"  → {p}")


if __name__ == "__main__":
    main()