"""
Worker для sherpa-onnx диаризации (CUDA через onnxruntime-gpu).

Запускается как subprocess из transcribe.py orchestrator. Работает в отдельном venv:
    <skill_root>/venv-sherpa

Использует:
    - pyannote-segmentation-3.0 в ONNX
    - 3D-Speaker eres2net 200k (multilingual) embedding extractor
    - Спектральная кластеризация (FastClustering) для группировки эмбеддингов

Аргументы:
    --input <audio>     путь к аудио (любой формат, конвертируется через ffmpeg в 16kHz mono WAV)
    --out-json <path>   путь сохранения turns JSON
    --num-speakers N    точное число спикеров (отключает кластеризацию по threshold)
    --threshold T       порог кластеризации (default 0.5, чем меньше - тем больше кластеров)
    --provider cuda|cpu (default cuda)

Вывод JSON: list[{"start": float, "end": float, "speaker": "SPEAKER_XX"}].
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import tempfile
import time
from pathlib import Path


SKILL_ROOT = Path(__file__).resolve().parent.parent
MODELS_DIR = SKILL_ROOT / "models"
SEG_MODEL = MODELS_DIR / "sherpa-onnx-pyannote-segmentation-3-0" / "model.onnx"
EMB_MODEL = MODELS_DIR / "3dspeaker_speech_eres2net_base_200k_sv_zh-cn_16k-common.onnx"


def setup_nvidia_dll_path() -> None:
    """Зарегистрировать bin-директории nvidia.* пакетов (CUDA 12 для onnxruntime-gpu)."""
    venv_root = Path(sys.executable).parent.parent
    nvidia_root = venv_root / "Lib" / "site-packages" / "nvidia"
    if not nvidia_root.exists():
        return
    for sub in nvidia_root.iterdir():
        bin_dir = sub / "bin"
        if bin_dir.is_dir():
            if hasattr(os, "add_dll_directory"):
                try:
                    os.add_dll_directory(str(bin_dir))
                except OSError:
                    pass
            os.environ["PATH"] = str(bin_dir) + os.pathsep + os.environ.get("PATH", "")


setup_nvidia_dll_path()

import numpy as np  # noqa: E402
import soundfile as sf  # noqa: E402
import sherpa_onnx  # noqa: E402

# Вывод содержит кириллицу. Без явного переключения печать падает с UnicodeEncodeError
# везде, где консоль не в UTF-8: сборочный агент, чужая локаль.
sys.stdout.reconfigure(encoding='utf-8')
sys.stderr.reconfigure(encoding='utf-8')


def ffmpeg_to_wav16k(input_path: Path, out_wav: Path) -> None:
    """Конвертировать любое аудио в 16kHz mono PCM_S16LE WAV через ffmpeg."""
    cmd = [
        "ffmpeg", "-y", "-i", str(input_path),
        "-vn", "-ac", "1", "-ar", "16000", "-acodec", "pcm_s16le",
        str(out_wav),
    ]
    subprocess.run(cmd, check=True, capture_output=True)


def diarize(args) -> int:
    if not SEG_MODEL.exists():
        print(f"[D] Не найдена модель сегментации: {SEG_MODEL}", file=sys.stderr)
        print(f"[D] Запустите установку: python scripts/setup.py", file=sys.stderr)
        return 1
    if not EMB_MODEL.exists():
        print(f"[D] Не найдена модель эмбеддингов: {EMB_MODEL}", file=sys.stderr)
        print(f"[D] Запустите установку: python scripts/setup.py", file=sys.stderr)
        return 1

    input_path = Path(args.input)
    print(f"[D] Файл: {input_path.name}", flush=True)
    print(f"[D] Provider: {args.provider}", flush=True)
    print(f"[D] Сегментация: {SEG_MODEL.name}", flush=True)
    print(f"[D] Эмбеддинги: {EMB_MODEL.name}", flush=True)

    with tempfile.TemporaryDirectory() as tmp:
        wav_path = Path(tmp) / "audio_16k.wav"
        print(f"[D] ffmpeg -> {wav_path.name}...", flush=True)
        t0 = time.time()
        ffmpeg_to_wav16k(input_path, wav_path)
        print(f"[D]   ffmpeg готов за {time.time() - t0:.1f}с", flush=True)

        samples, sr = sf.read(str(wav_path), dtype="float32")
        if samples.ndim > 1:
            samples = samples.mean(axis=1)
        print(f"[D] Sample rate: {sr}, длительность: {len(samples) / sr:.1f}с", flush=True)

        config = sherpa_onnx.OfflineSpeakerDiarizationConfig(
            segmentation=sherpa_onnx.OfflineSpeakerSegmentationModelConfig(
                pyannote=sherpa_onnx.OfflineSpeakerSegmentationPyannoteModelConfig(
                    model=str(SEG_MODEL),
                ),
                provider=args.provider,
                num_threads=1,
                debug=False,
            ),
            embedding=sherpa_onnx.SpeakerEmbeddingExtractorConfig(
                model=str(EMB_MODEL),
                provider=args.provider,
                num_threads=1,
                debug=False,
            ),
            clustering=sherpa_onnx.FastClusteringConfig(
                num_clusters=args.num_speakers if args.num_speakers else -1,
                threshold=args.threshold,
            ),
            min_duration_on=0.3,
            min_duration_off=0.5,
        )

        if not config.validate():
            print("[D] Конфигурация невалидна", file=sys.stderr)
            return 1

        print("[D] Загрузка моделей...", flush=True)
        t0 = time.time()
        sd = sherpa_onnx.OfflineSpeakerDiarization(config)
        print(f"[D]   Готово за {time.time() - t0:.1f}с", flush=True)

        print("[D] Диаризация...", flush=True)
        t0 = time.time()
        last_pct = -1

        def progress(num_processed: int, num_total: int) -> int:
            nonlocal last_pct
            pct = int(100 * num_processed / max(num_total, 1))
            if pct >= last_pct + 5:
                last_pct = pct
                print(f"[D]   [{pct:3d}%]", flush=True)
            return 0

        result = sd.process(samples, callback=progress).sort_by_start_time()
        elapsed = time.time() - t0

        turns: list[dict] = []
        for r in result:
            turns.append({
                "start": float(r.start),
                "end": float(r.end),
                "speaker": f"SPEAKER_{int(r.speaker):02d}",
            })

        speakers_set = sorted({t["speaker"] for t in turns})
        rtf = elapsed / (len(samples) / sr) if len(samples) > 0 else 0
        print(f"[D] Готово за {elapsed:.1f}с (RTF {rtf:.3f}, {len(turns)} turns, {len(speakers_set)} спикеров)", flush=True)

        Path(args.out_json).write_text(json.dumps(turns, ensure_ascii=False), encoding="utf-8")
        print(f"[D]   -> {args.out_json}", flush=True)

    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(0)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True)
    ap.add_argument("--out-json", required=True)
    ap.add_argument("--num-speakers", type=int, default=None)
    ap.add_argument("--threshold", type=float, default=0.5)
    ap.add_argument("--provider", default="cuda", choices=["cuda", "cpu"])
    args = ap.parse_args()
    return diarize(args)


if __name__ == "__main__":
    import traceback
    try:
        rc = main()
    except SystemExit:
        raise
    except BaseException:
        log_path = SKILL_ROOT / "diarize_sherpa.crash.log"
        with log_path.open("a", encoding="utf-8") as f:
            f.write("=" * 80 + "\n")
            f.write(f"argv: {sys.argv}\n")
            f.write(traceback.format_exc())
        traceback.print_exc()
        rc = 2
    sys.exit(rc)

"""
Установка зависимостей скила transcribe-audio-local.

Что делает:
1. Проверяет системные требования (Python, ffmpeg).
2. Создает venv-whisper и ставит faster-whisper + ctranslate2 (CUDA) + av.
3. Создает venv-sherpa и ставит sherpa_onnx (GPU CUDA) + onnxruntime-gpu + soundfile.
4. Скачивает модели диаризации в models/:
   - sherpa-onnx-pyannote-segmentation-3-0 (~7 MB)
   - 3dspeaker_speech_eres2net_base_200k_sv_zh-cn_16k-common.onnx (~40 MB)

Запуск:
    python scripts/setup.py
    python scripts/setup.py --skip-models       # только venv'ы
    python scripts/setup.py --skip-sherpa       # без диаризации
    python scripts/setup.py --skip-whisper      # пропустить пересоздание venv-whisper

Требования:
    - Python 3.10+ (рекомендуется 3.12)
    - NVIDIA GPU + CUDA 12 + cuDNN 9 (для GPU режима)
    - ffmpeg в PATH
    - Windows x64 или Linux x64
"""

from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
import tarfile
import urllib.request
from pathlib import Path

# Вывод содержит кириллицу. Без явного переключения печать падает с UnicodeEncodeError
# везде, где консоль не в UTF-8: сборочный агент, чужая локаль.
sys.stdout.reconfigure(encoding='utf-8')
sys.stderr.reconfigure(encoding='utf-8')

SKILL_ROOT = Path(__file__).resolve().parent.parent
VENV_WHISPER = SKILL_ROOT / "venv-whisper"
VENV_SHERPA = SKILL_ROOT / "venv-sherpa"
MODELS_DIR = SKILL_ROOT / "models"

IS_WIN = os.name == "nt"
VENV_BIN = "Scripts" if IS_WIN else "bin"
PY_EXE = "python.exe" if IS_WIN else "python"

# URL моделей (k2-fsa официальные релизы)
MODEL_URLS = {
    "sherpa-onnx-pyannote-segmentation-3-0.tar.bz2":
        "https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-segmentation-models/"
        "sherpa-onnx-pyannote-segmentation-3-0.tar.bz2",
    "3dspeaker_speech_eres2net_base_200k_sv_zh-cn_16k-common.onnx":
        "https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-recongition-models/"
        "3dspeaker_speech_eres2net_base_200k_sv_zh-cn_16k-common.onnx",
}

# Пакеты для venv-whisper (faster-whisper + GPU stack)
WHISPER_PACKAGES = [
    "faster-whisper>=1.0",
    "ctranslate2>=4.5",
    "av>=11",
    "huggingface-hub",
    "numpy<3",
    "nvidia-cublas-cu12",
    "nvidia-cudnn-cu12",
    "nvidia-cuda-runtime-cu12",
    "nvidia-cuda-nvrtc-cu12",
]

# Пакеты для venv-sherpa
# sherpa_onnx CUDA wheel ставится отдельно по индексу k2-fsa (см. install_sherpa_gpu)
SHERPA_PACKAGES = [
    "onnxruntime-gpu>=1.18",
    "soundfile>=0.12",
    "numpy<2",
    "nvidia-cublas-cu12",
    "nvidia-cudnn-cu12",
    "nvidia-cuda-runtime-cu12",
    "nvidia-cuda-nvrtc-cu12",
    "nvidia-cufft-cu12",
    "nvidia-nvjitlink-cu12",
]


def step(msg: str) -> None:
    print(f"\n{'=' * 70}\n{msg}\n{'=' * 70}", flush=True)


def info(msg: str) -> None:
    print(f"  {msg}", flush=True)


def run(cmd: list[str], **kwargs) -> int:
    print(f"  $ {' '.join(str(c) for c in cmd)}", flush=True)
    return subprocess.run(cmd, **kwargs).returncode


REQUIRED_GB = 5.0
MIN_DRIVER_VERSION = (525, 0)  # минимум для CUDA 12.0+


def check_gpu() -> tuple[bool, bool]:
    """Probe NVIDIA GPU + версия драйвера. Возвращает (gpu_found, driver_ok).

    Печатает info-строки. Если драйвер старый - выдает инструкцию обновить.
    """
    if not shutil.which("nvidia-smi"):
        return False, False
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total,driver_version", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=10,
        )
        if result.returncode != 0 or not result.stdout.strip():
            return False, False
    except (subprocess.SubprocessError, OSError):
        return False, False

    driver_ok = True
    for line in result.stdout.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 3:
            continue
        name, mem, driver_ver = parts[0], parts[1], parts[2]
        info(f"OK:   GPU: {name} ({mem})")
        info(f"OK:   Драйвер NVIDIA: {driver_ver}")
        try:
            major_minor = tuple(int(x) for x in driver_ver.split(".")[:2])
            if major_minor < MIN_DRIVER_VERSION:
                driver_ok = False
                info(f"WARN: Драйвер {driver_ver} ниже минимального {MIN_DRIVER_VERSION[0]}.x для CUDA 12")
                info("      Обновите драйверы NVIDIA:")
                info("      - Windows: https://www.nvidia.com/Download/index.aspx или GeForce Experience")
                info("      - Linux: см. инструкцию вашего дистрибутива (apt install nvidia-driver-550 и т.п.)")
                info("      Без обновления faster-whisper / sherpa-onnx могут падать с CUDA errors.")
        except ValueError:
            info(f"WARN: Не удалось распарсить версию драйвера '{driver_ver}'")

    return True, driver_ok


def check_requirements(allow_cpu: bool = False) -> tuple[bool, bool]:
    """Возвращает (ok, need_static_ffmpeg)."""
    step("Проверка требований")
    ok = True
    need_static_ffmpeg = False

    py_version = sys.version_info
    if py_version < (3, 10):
        info(f"FAIL: Python {py_version.major}.{py_version.minor} < 3.10")
        ok = False
    else:
        info(f"OK:   Python {py_version.major}.{py_version.minor}.{py_version.micro}")

    if not shutil.which("ffmpeg"):
        info("INFO: ffmpeg не найден в PATH - будет установлен как pip-пакет static-ffmpeg в venv-whisper")
        need_static_ffmpeg = True
    else:
        info(f"OK:   ffmpeg в {shutil.which('ffmpeg')}")

    gpu_found, _driver_ok = check_gpu()
    if not gpu_found:
        if allow_cpu:
            info("WARN: NVIDIA GPU не найден (nvidia-smi). Установка продолжится для CPU-режима.")
            info("      Транскрипция будет работать, но в 10+ раз медленнее. Запуск с --device cpu.")
        else:
            info("FAIL: NVIDIA GPU не найден (nvidia-smi). Скил настроен для CUDA GPU.")
            info("      Если у вас нет GPU - перезапустите setup с флагом --allow-cpu")
            info("      Если GPU есть, но nvidia-smi не работает - установите драйверы NVIDIA + CUDA 12.")
            ok = False

    try:
        free_gb = shutil.disk_usage(SKILL_ROOT).free / (1024 ** 3)
        if free_gb < REQUIRED_GB:
            info(f"FAIL: Свободного места {free_gb:.1f} GB < требуется {REQUIRED_GB} GB")
            ok = False
        else:
            info(f"OK:   Свободно на диске: {free_gb:.1f} GB")
    except OSError as e:
        info(f"WARN: Не удалось проверить свободное место: {e}")

    info(f"Платформа: {'Windows' if IS_WIN else sys.platform}")
    return ok, need_static_ffmpeg


def create_venv(venv_path: Path) -> Path:
    """Создать venv если его нет, вернуть путь к python внутри."""
    py = venv_path / VENV_BIN / PY_EXE
    if py.exists():
        info(f"venv уже существует: {venv_path}")
        return py
    info(f"Создание venv: {venv_path}")
    rc = run([sys.executable, "-m", "venv", str(venv_path)])
    if rc != 0:
        raise RuntimeError(f"Не удалось создать venv {venv_path}")
    rc = run([str(py), "-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel"])
    if rc != 0:
        raise RuntimeError("Не удалось обновить pip")
    return py


def pip_install(py: Path, packages: list[str], extra_args: list[str] | None = None) -> None:
    cmd = [str(py), "-m", "pip", "install"]
    if extra_args:
        cmd += extra_args
    cmd += packages
    rc = run(cmd)
    if rc != 0:
        raise RuntimeError(f"pip install упал на: {packages}")


def install_whisper(skip: bool, install_static_ffmpeg: bool) -> None:
    step("venv-whisper: faster-whisper + CUDA stack")
    if skip:
        info("Пропускаем по флагу --skip-whisper")
        return
    py = create_venv(VENV_WHISPER)
    pip_install(py, WHISPER_PACKAGES)
    if install_static_ffmpeg:
        info("Установка static-ffmpeg (ffmpeg+ffprobe бинарники как pip-пакет)")
        pip_install(py, ["static-ffmpeg"])
    info("OK")


def install_sherpa_gpu(py: Path) -> None:
    """Установить sherpa_onnx с GPU поддержкой.

    Сначала пробуем с PyPI с extra-index-url от k2-fsa (содержит CUDA wheels).
    Если не вышло - инструкция пользователю.
    """
    info("Установка sherpa_onnx (GPU/CUDA)...")
    extra_index = "https://k2-fsa.github.io/sherpa/onnx/cuda.html"
    try:
        cmd = [str(py), "-m", "pip", "install", "sherpa-onnx",
               "-f", extra_index]
        rc = run(cmd)
        if rc == 0:
            info("sherpa_onnx установлен (попробуйте проверить наличие CUDA: см. README)")
            return
    except Exception as e:
        info(f"WARN: {e}")

    info("ВНИМАНИЕ: автоматическая установка sherpa_onnx с CUDA не удалась.")
    info("Установите вручную:")
    info(f"  {py} -m pip install sherpa-onnx")
    info("Или скачайте GPU wheel с https://huggingface.co/csukuangfj2/sherpa-onnx-wheels")
    info("и установите:")
    info(f"  {py} -m pip install <путь-к-wheel.whl>")


def install_sherpa(skip: bool) -> None:
    step("venv-sherpa: sherpa_onnx + onnxruntime-gpu")
    if skip:
        info("Пропускаем по флагу --skip-sherpa")
        return
    py = create_venv(VENV_SHERPA)
    pip_install(py, SHERPA_PACKAGES)
    install_sherpa_gpu(py)
    info("OK")


def download_file(url: str, dest: Path) -> None:
    if dest.exists():
        info(f"Уже скачан: {dest.name} ({dest.stat().st_size / 1e6:.1f} MB)")
        return
    info(f"Скачивание {dest.name}")
    info(f"  url: {url}")
    tmp = dest.with_suffix(dest.suffix + ".tmp")
    try:
        with urllib.request.urlopen(url) as resp, open(tmp, "wb") as f:
            total = int(resp.headers.get("Content-Length", 0))
            downloaded = 0
            chunk = 1 << 16
            last_pct = -1
            while True:
                data = resp.read(chunk)
                if not data:
                    break
                f.write(data)
                downloaded += len(data)
                if total:
                    pct = int(100 * downloaded / total)
                    if pct >= last_pct + 10:
                        last_pct = pct
                        info(f"  [{pct:3d}%] {downloaded / 1e6:.1f} / {total / 1e6:.1f} MB")
        tmp.rename(dest)
    except Exception as e:
        if tmp.exists():
            tmp.unlink()
        raise RuntimeError(f"Не удалось скачать {url}: {e}") from e


def extract_segmentation_archive() -> None:
    """Распаковать sherpa-onnx-pyannote-segmentation-3-0.tar.bz2 в models/."""
    archive = MODELS_DIR / "sherpa-onnx-pyannote-segmentation-3-0.tar.bz2"
    target_dir = MODELS_DIR / "sherpa-onnx-pyannote-segmentation-3-0"
    if (target_dir / "model.onnx").exists():
        info(f"Уже распакован: {target_dir}")
        return
    info(f"Распаковка {archive.name}")
    with tarfile.open(archive, "r:bz2") as tf:
        tf.extractall(MODELS_DIR)
    if not (target_dir / "model.onnx").exists():
        raise RuntimeError(f"После распаковки нет {target_dir / 'model.onnx'}")


def download_models(skip: bool) -> None:
    step("Модели диаризации")
    if skip:
        info("Пропускаем по флагу --skip-models")
        return
    MODELS_DIR.mkdir(exist_ok=True)
    for fname, url in MODEL_URLS.items():
        download_file(url, MODELS_DIR / fname)
    extract_segmentation_archive()
    info("OK")


def main() -> int:
    ap = argparse.ArgumentParser(description="Установка скила transcribe-audio-local")
    ap.add_argument("--skip-whisper", action="store_true", help="Не пересоздавать venv-whisper")
    ap.add_argument("--skip-sherpa", action="store_true", help="Не ставить sherpa (без диаризации)")
    ap.add_argument("--skip-models", action="store_true", help="Не скачивать модели")
    ap.add_argument("--allow-cpu", action="store_true", help="Разрешить установку без GPU (CPU-режим, в 10+ раз медленнее)")
    args = ap.parse_args()

    ok, need_static_ffmpeg = check_requirements(allow_cpu=args.allow_cpu)
    if not ok:
        print("\nПроверка требований не пройдена. Исправьте и повторите.", file=sys.stderr)
        return 1

    try:
        install_whisper(args.skip_whisper, install_static_ffmpeg=need_static_ffmpeg)
        install_sherpa(args.skip_sherpa)
        download_models(args.skip_models)
    except Exception as e:
        print(f"\nОшибка установки: {e}", file=sys.stderr)
        return 1

    step("Готово")
    info(f"Скил: {SKILL_ROOT}")
    info(f"venv-whisper: {VENV_WHISPER}")
    info(f"venv-sherpa:  {VENV_SHERPA}")
    info(f"models:       {MODELS_DIR}")
    info("")
    info("Проверка:")
    info(f"  python {SKILL_ROOT / 'scripts' / 'transcribe.py'} <audio.mp3>")
    info(f"  python {SKILL_ROOT / 'scripts' / 'transcribe.py'} <audio.mp3> --diarize")
    return 0


if __name__ == "__main__":
    sys.exit(main())

"""
Локальная транскрипция аудио через faster-whisper (CUDA) + опц. диаризация (sherpa-onnx GPU).

Архитектура: orchestrator + 2 subprocess (изоляция CUDA-DLL ctranslate2 vs onnxruntime).
При --diarize транскрипция и диаризация запускаются параллельно.

Использование:
    python transcribe.py <input_path> [--output-dir DIR] [--model MODEL]
                         [--language ru] [--diarize]
                         [--num-speakers N] [--threshold T]

Зависимости:
    venv-whisper (рядом со скилом): faster-whisper, ctranslate2-CUDA, av, huggingface-hub.
    venv-sherpa (рядом со скилом, для --diarize): sherpa_onnx CUDA, onnxruntime-gpu, soundfile.
    ffmpeg в PATH.

Файлы вывода:
    <base> - транскрипция.md       MD c таймкодами по сегментам (без спикеров)
    <base> - транскрипция.txt      Plain text
    <base> - со спикерами.md       MD c [Спикер, MM:SS] (только при --diarize)
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time

# Подключаем ffmpeg из static-ffmpeg если он установлен (когда системного ffmpeg нет в PATH).
# subprocess'ы наследуют PATH, поэтому worker'ам ffmpeg тоже будет доступен.
try:
    from static_ffmpeg import add_paths as _sff_add_paths
    _sff_add_paths()
except ImportError:
    pass
import traceback
from pathlib import Path

# Вывод содержит кириллицу. Без явного переключения печать падает с UnicodeEncodeError
# везде, где консоль не в UTF-8: сборочный агент, чужая локаль.
sys.stdout.reconfigure(encoding='utf-8')
sys.stderr.reconfigure(encoding='utf-8')


SKILL_ROOT = Path(__file__).resolve().parent.parent
WHISPER_VENV_PYTHON = SKILL_ROOT / "venv-whisper" / ("Scripts" if os.name == "nt" else "bin") / ("python.exe" if os.name == "nt" else "python")
SHERPA_VENV_PYTHON = SKILL_ROOT / "venv-sherpa" / ("Scripts" if os.name == "nt" else "bin") / ("python.exe" if os.name == "nt" else "python")
SHERPA_WORKER = Path(__file__).resolve().parent / "diarize_sherpa.py"

AUDIO_EXTS = {".m4a", ".mp3", ".wav", ".ogg", ".flac", ".aac", ".wma", ".opus"}


def _setup_cuda_dll_path() -> None:
    """Добавить bin-директории nvidia.* пакетов в PATH и DLL search (Windows, CUDA 12)."""
    for name in ("nvidia.cublas", "nvidia.cudnn", "nvidia.cuda_runtime", "nvidia.cuda_nvrtc"):
        try:
            mod = __import__(name, fromlist=[""])
            bin_dir = os.path.join(mod.__path__[0], "bin")
            if os.path.isdir(bin_dir):
                if hasattr(os, "add_dll_directory"):
                    os.add_dll_directory(bin_dir)
                os.environ["PATH"] = bin_dir + os.pathsep + os.environ.get("PATH", "")
        except ImportError:
            pass


def format_ts(seconds: float) -> str:
    s = int(seconds)
    return f"{s // 60:02d}:{s % 60:02d}"


# =====================================================================
# WORKER: транскрипция (запускается как subprocess в venv-whisper)
# =====================================================================

def worker_transcribe(args) -> int:
    _setup_cuda_dll_path()
    from faster_whisper import WhisperModel  # noqa

    print(f"[T] Модель: {args.model} ({args.device}/{args.compute_type})", flush=True)
    t0 = time.time()
    model = WhisperModel(args.model, device=args.device, compute_type=args.compute_type)
    print(f"[T] Модель загружена за {time.time() - t0:.1f}с", flush=True)

    print(f"[T] Транскрипция (word_timestamps={args.need_words})...", flush=True)
    t0 = time.time()
    segments_iter, info = model.transcribe(
        args.input,
        language=args.language,
        beam_size=5,
        vad_filter=True,
        vad_parameters={"min_silence_duration_ms": 500},
        condition_on_previous_text=True,
        word_timestamps=args.need_words,
    )
    print(f"[T] Язык: {info.language} (p={info.language_probability:.2f}) длительность {info.duration:.1f}с", flush=True)

    segments: list[dict] = []
    last_pct = -1
    for seg in segments_iter:
        words = []
        if args.need_words and seg.words:
            for w in seg.words:
                words.append({"start": float(w.start), "end": float(w.end), "text": w.word})
        segments.append({
            "start": float(seg.start),
            "end": float(seg.end),
            "text": seg.text.strip(),
            "words": words,
        })
        pct = int(100 * seg.end / max(info.duration, 1))
        if pct >= last_pct + 5:
            last_pct = pct
            print(f"[T] [{pct:3d}%] {format_ts(seg.start)}: {seg.text.strip()[:80]}", flush=True)

    info_dict = {
        "language": info.language,
        "language_probability": float(info.language_probability),
        "duration": float(info.duration),
    }

    out = {"segments": segments, "info": info_dict}
    Path(args.out_json).write_text(json.dumps(out, ensure_ascii=False), encoding="utf-8")
    print(f"[T] Готово за {time.time() - t0:.1f}с ({len(segments)} сегментов) -> {args.out_json}", flush=True)
    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(0)


# =====================================================================
# MERGE / OUTPUT
# =====================================================================

def assign_speaker(midpoint: float, turns: list[dict]) -> str:
    for t in turns:
        if t["start"] <= midpoint <= t["end"]:
            return t["speaker"]
    if not turns:
        return "UNKNOWN"
    nearest = min(turns, key=lambda t: min(abs(t["start"] - midpoint), abs(t["end"] - midpoint)))
    return nearest["speaker"]


def smooth_word_speakers(words: list[dict], min_run: float = 1.2) -> None:
    """Сгладить короткие пробивки спикера в словах: если короткий run X между двумя длинными run Y,
    переаттрибутировать слова run X на спикера Y. Изменяет words на месте.
    """
    if not words:
        return
    runs: list[tuple[int, int, str, float]] = []
    i = 0
    n = len(words)
    while i < n:
        j = i
        spk = words[i]["speaker"]
        while j + 1 < n and words[j + 1]["speaker"] == spk:
            j += 1
        duration = words[j]["end"] - words[i]["start"]
        runs.append((i, j, spk, duration))
        i = j + 1

    changed = True
    while changed:
        changed = False
        for k in range(1, len(runs) - 1):
            ri, rj, rspk, rdur = runs[k]
            li, lj, lspk, ldur = runs[k - 1]
            ni, nj, nspk, ndur = runs[k + 1]
            if rdur < min_run and lspk == nspk and lspk != rspk:
                for x in range(ri, rj + 1):
                    words[x]["speaker"] = lspk
                runs[k - 1] = (li, rj, lspk, words[rj]["end"] - words[li]["start"])
                runs.pop(k)
                if k < len(runs):
                    nri, nrj, nrspk, _ = runs[k]
                    if nrspk == lspk:
                        runs[k - 1] = (li, nrj, lspk, words[nrj]["end"] - words[li]["start"])
                        runs.pop(k)
                changed = True
                break


def merge_with_speakers(segments: list[dict], turns: list[dict]) -> list[dict]:
    all_words: list[dict] = []
    word_segments_no_words: list[dict] = []
    for seg in segments:
        words = seg.get("words") or []
        if not words:
            mid = (seg["start"] + seg["end"]) / 2
            spk = assign_speaker(mid, turns)
            word_segments_no_words.append({
                "start": seg["start"], "end": seg["end"], "speaker": spk, "text": seg["text"],
            })
            continue
        for w in words:
            mid = (w["start"] + w["end"]) / 2
            spk = assign_speaker(mid, turns)
            all_words.append({"start": w["start"], "end": w["end"], "text": w["text"], "speaker": spk})

    smooth_word_speakers(all_words, min_run=1.2)

    utterances: list[dict] = []
    current: dict | None = None
    for w in all_words:
        if current and current["speaker"] == w["speaker"]:
            current["end"] = w["end"]
            current["text"] += w["text"]
        else:
            if current:
                utterances.append(current)
            current = {"start": w["start"], "end": w["end"], "speaker": w["speaker"], "text": w["text"]}
    if current:
        utterances.append(current)

    utterances.extend(word_segments_no_words)
    utterances.sort(key=lambda u: u["start"])

    merged: list[dict] = []
    for u in utterances:
        u["text"] = u["text"].strip()
        if not u["text"]:
            continue
        if merged and merged[-1]["speaker"] == u["speaker"] and u["start"] - merged[-1]["end"] < 1.5:
            merged[-1]["end"] = u["end"]
            merged[-1]["text"] += " " + u["text"]
        else:
            merged.append(u)
    return merged


def write_basic_md(path: Path, input_name: str, info: dict, segments: list[dict], model_name: str) -> None:
    lines = [
        f"# Транскрипция: {input_name}",
        "",
        f"- Модель: `{model_name}`",
        f"- Язык: {info['language']} (p={info['language_probability']:.2f})",
        f"- Длительность: {format_ts(info['duration'])} ({info['duration']:.1f}с)",
        "",
        "---",
        "",
    ]
    for seg in segments:
        lines.append(f"**[{format_ts(seg['start'])}]** {seg['text']}")
        lines.append("")
    path.write_text("\n".join(lines), encoding="utf-8")


def write_plain(path: Path, segments: list[dict]) -> None:
    path.write_text(" ".join(s["text"] for s in segments), encoding="utf-8")


def write_speakers_md(path: Path, input_name: str, info: dict, utterances: list[dict], model_name: str) -> None:
    speakers_set = sorted({u["speaker"] for u in utterances})
    lines = [
        f"# Транскрипция со спикерами: {input_name}",
        "",
        f"- Модель: `{model_name}`",
        f"- Диаризация: `sherpa-onnx (pyannote-segmentation-3.0 + 3D-Speaker)`",
        f"- Длительность: {format_ts(info['duration'])} ({info['duration']:.1f}с)",
        f"- Найдено спикеров: {len(speakers_set)} ({', '.join(speakers_set)})",
        "",
        "---",
        "",
    ]
    for u in utterances:
        lines.append(f"**[{u['speaker']}, {format_ts(u['start'])}]** {u['text']}")
        lines.append("")
    path.write_text("\n".join(lines), encoding="utf-8")


# =====================================================================
# ORCHESTRATOR
# =====================================================================

def spawn_worker(mode: str, env_extra: dict[str, str], cli: list[str], python_exe: Path) -> subprocess.Popen:
    cmd = [str(python_exe), __file__, "--worker", mode] + cli
    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    env["PYTHONIOENCODING"] = "utf-8"
    env.update(env_extra)
    return subprocess.Popen(cmd, env=env)


def spawn_sherpa_diarize(env_extra: dict[str, str], cli: list[str]) -> subprocess.Popen:
    """Запустить sherpa-onnx диаризацию из отдельного venv-sherpa."""
    if not SHERPA_VENV_PYTHON.exists():
        raise RuntimeError(
            f"venv-sherpa не найден: {SHERPA_VENV_PYTHON}. "
            "Запустите установку: python scripts/setup.py"
        )
    cmd = [str(SHERPA_VENV_PYTHON), str(SHERPA_WORKER)] + cli
    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    env["PYTHONIOENCODING"] = "utf-8"
    env.update(env_extra)
    return subprocess.Popen(cmd, env=env)


def orchestrate(args) -> int:
    input_path = Path(args.input)
    if not input_path.exists():
        print(f"Файл не найден: {input_path}", file=sys.stderr)
        return 1

    ext = input_path.suffix.lower()
    if ext not in AUDIO_EXTS:
        print(f"Расширение {ext} не поддерживается. Допустимы: {', '.join(sorted(AUDIO_EXTS))}", file=sys.stderr)
        return 1

    if not WHISPER_VENV_PYTHON.exists():
        print(f"venv-whisper не найден: {WHISPER_VENV_PYTHON}", file=sys.stderr)
        print("Запустите установку: python scripts/setup.py", file=sys.stderr)
        return 1

    if args.output_dir:
        output_dir = Path(args.output_dir)
    else:
        output_dir = input_path.parent / "Транскрипция" / input_path.stem
    output_dir.mkdir(parents=True, exist_ok=True)

    base = input_path.stem
    transcribe_json = output_dir / f"{base}.transcribe.json"
    turns_json = output_dir / f"{base}.turns.json"

    print(f"Файл: {input_path.name}", flush=True)
    print(f"Каталог: {output_dir}", flush=True)

    transcribe_cli = [
        "--input", str(input_path),
        "--out-json", str(transcribe_json),
        "--model", args.model,
        "--language", args.language,
        "--device", args.device,
        "--compute-type", args.compute_type,
    ]
    if args.diarize:
        transcribe_cli.append("--need-words")

    transcribe_proc = spawn_worker("transcribe", {}, transcribe_cli, WHISPER_VENV_PYTHON)

    diarize_proc = None
    if args.diarize:
        sherpa_cli = [
            "--input", str(input_path),
            "--out-json", str(turns_json),
            "--provider", args.device,
            "--threshold", str(args.threshold),
        ]
        if args.num_speakers is not None:
            sherpa_cli += ["--num-speakers", str(args.num_speakers)]
        print(f"Диаризация: sherpa-onnx (pyannote-segmentation-3.0 + 3D-Speaker)", flush=True)
        diarize_proc = spawn_sherpa_diarize({}, sherpa_cli)

    rc_t = transcribe_proc.wait()
    if rc_t != 0:
        if diarize_proc is not None:
            diarize_proc.terminate()
        print(f"Транскрипция упала с кодом {rc_t}", file=sys.stderr)
        return rc_t

    payload = json.loads(transcribe_json.read_text(encoding="utf-8"))
    segments = payload["segments"]
    info = payload["info"]

    transcript_md = output_dir / f"{base} - транскрипция.md"
    plain_txt = output_dir / f"{base} - транскрипция.txt"
    write_basic_md(transcript_md, input_path.name, info, segments, args.model)
    write_plain(plain_txt, segments)
    print(f"  MD:    {transcript_md}", flush=True)
    print(f"  Plain: {plain_txt}", flush=True)

    if diarize_proc is not None:
        rc_d = diarize_proc.wait()
        if rc_d != 0:
            print(f"Диаризация упала с кодом {rc_d} (транскрипция сохранена)", file=sys.stderr)
            return rc_d
        turns = json.loads(turns_json.read_text(encoding="utf-8"))
        utterances = merge_with_speakers(segments, turns)
        speakers_md = output_dir / f"{base} - со спикерами.md"
        write_speakers_md(speakers_md, input_path.name, info, utterances, args.model)
        print(f"  Spk:   {speakers_md}", flush=True)

    if not args.keep_intermediate:
        for f in (transcribe_json, turns_json):
            try:
                if f.exists():
                    f.unlink()
            except OSError:
                pass

    return 0


# =====================================================================
# MAIN
# =====================================================================

def main() -> int:
    if len(sys.argv) >= 3 and sys.argv[1] == "--worker":
        mode = sys.argv[2]
        worker_args = sys.argv[3:]
        ap = argparse.ArgumentParser()
        if mode == "transcribe":
            ap.add_argument("--input", required=True)
            ap.add_argument("--out-json", required=True)
            ap.add_argument("--model", default="mobiuslabsgmbh/faster-whisper-large-v3-turbo")
            ap.add_argument("--language", default="ru")
            ap.add_argument("--device", default="cuda")
            ap.add_argument("--compute-type", default="float16")
            ap.add_argument("--need-words", action="store_true")
            return worker_transcribe(ap.parse_args(worker_args))
        print(f"Неизвестный worker mode: {mode}", file=sys.stderr)
        return 1

    ap = argparse.ArgumentParser(description="Локальная транскрипция аудио через faster-whisper + опц. диаризация sherpa-onnx")
    ap.add_argument("input", help="Путь к аудиофайлу (m4a/mp3/wav/ogg/flac/aac/wma/opus)")
    ap.add_argument("--output-dir", default=None, help="Каталог вывода (default: <каталог входа>/Транскрипция/<имя>/)")
    ap.add_argument("--model", default="mobiuslabsgmbh/faster-whisper-large-v3-turbo")
    ap.add_argument("--language", default="ru")
    ap.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
    ap.add_argument("--compute-type", default="float16")
    ap.add_argument("--diarize", action="store_true", help="Включить диаризацию (sherpa-onnx, параллельно)")
    ap.add_argument("--num-speakers", type=int, default=None, help="Точное число спикеров (без авто-кластеризации)")
    ap.add_argument("--threshold", type=float, default=0.5, help="Порог кластеризации sherpa-onnx (default 0.5)")
    ap.add_argument("--keep-intermediate", action="store_true", help="Не удалять промежуточные JSON")
    args = ap.parse_args()
    return orchestrate(args)


if __name__ == "__main__":
    try:
        rc = main()
    except SystemExit:
        raise
    except BaseException:
        log_path = SKILL_ROOT / "transcribe.crash.log"
        with log_path.open("a", encoding="utf-8") as f:
            f.write("=" * 80 + "\n")
            f.write(f"argv: {sys.argv}\n")
            f.write(traceback.format_exc())
            f.write("\n")
        print(f"\nКраш записан в {log_path}", file=sys.stderr)
        traceback.print_exc()
        rc = 2
    sys.exit(rc)

"""
Smoke-тест после установки скила transcribe-audio-local.

Что проверяет:
1. venv-whisper существует и в нем установлен faster-whisper
2. venv-sherpa существует и в нем установлены sherpa_onnx + onnxruntime
3. Модели sherpa-onnx и 3D-Speaker лежат в models/
4. ffmpeg доступен
5. (опционально) тестовый прогон на коротком сгенерированном WAV

Запуск:
    python scripts/verify.py             # все проверки кроме прогона
    python scripts/verify.py --full      # + тестовый прогон (создает 3-сек тон и транскрибирует)

Выход: 0 если все проверки прошли, 1 если есть FAIL.
"""

from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
import tempfile
import wave
from pathlib import Path

# Вывод содержит кириллицу. Без явного переключения печать падает с UnicodeEncodeError
# везде, где консоль не в UTF-8: сборочный агент, чужая локаль.
sys.stdout.reconfigure(encoding='utf-8')
sys.stderr.reconfigure(encoding='utf-8')


SKILL_ROOT = Path(__file__).resolve().parent.parent
VENV_BIN = "Scripts" if os.name == "nt" else "bin"
PY_EXE = "python.exe" if os.name == "nt" else "python"

VENV_WHISPER_PY = SKILL_ROOT / "venv-whisper" / VENV_BIN / PY_EXE
VENV_SHERPA_PY = SKILL_ROOT / "venv-sherpa" / VENV_BIN / PY_EXE
MODELS_DIR = SKILL_ROOT / "models"
SEG_MODEL = MODELS_DIR / "sherpa-onnx-pyannote-segmentation-3-0" / "model.onnx"
EMB_MODEL = MODELS_DIR / "3dspeaker_speech_eres2net_base_200k_sv_zh-cn_16k-common.onnx"

results: list[tuple[str, bool, str]] = []


def check(name: str, condition: bool, detail: str = "") -> None:
    results.append((name, condition, detail))
    status = "OK  " if condition else "FAIL"
    print(f"  [{status}] {name}" + (f": {detail}" if detail else ""), flush=True)


def check_venv(name: str, py_exe: Path, test_import: str) -> bool:
    if not py_exe.exists():
        check(f"venv-{name}", False, f"не найден: {py_exe}")
        return False
    check(f"venv-{name}", True, str(py_exe))
    try:
        result = subprocess.run(
            [str(py_exe), "-c", test_import],
            capture_output=True, text=True, timeout=60,
        )
        if result.returncode == 0:
            check(f"venv-{name} imports", True, result.stdout.strip())
            return True
        check(f"venv-{name} imports", False, result.stderr.strip().splitlines()[-1] if result.stderr else "ошибка импорта")
        return False
    except (subprocess.SubprocessError, OSError) as e:
        check(f"venv-{name} imports", False, str(e))
        return False


def create_silent_wav(path: Path, duration_sec: float = 3.0, sample_rate: int = 16000) -> None:
    n_samples = int(duration_sec * sample_rate)
    silence = b"\x00\x00" * n_samples
    with wave.open(str(path), "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sample_rate)
        w.writeframes(silence)


def run_transcribe_test() -> bool:
    print("\n  Тестовый прогон транскрипции на 3-сек silent WAV...", flush=True)
    with tempfile.TemporaryDirectory() as tmp:
        wav_path = Path(tmp) / "test.wav"
        create_silent_wav(wav_path)
        out_dir = Path(tmp) / "out"
        out_dir.mkdir()
        try:
            result = subprocess.run(
                [str(VENV_WHISPER_PY),
                 str(SKILL_ROOT / "scripts" / "transcribe.py"),
                 str(wav_path),
                 "--output-dir", str(out_dir),
                 "--model", "openai/whisper-tiny",
                 "--compute-type", "int8"],
                capture_output=True, text=True, timeout=180,
            )
            if result.returncode == 0:
                check("smoke-test транскрипция", True, "транскрибация silent WAV прошла")
                return True
            check("smoke-test транскрипция", False, f"rc={result.returncode}: {result.stderr[-200:]}")
            return False
        except subprocess.TimeoutExpired:
            check("smoke-test транскрипция", False, "таймаут (>3 мин)")
            return False
        except OSError as e:
            check("smoke-test транскрипция", False, str(e))
            return False


def main() -> int:
    ap = argparse.ArgumentParser(description="Smoke-тест скила transcribe-audio-local")
    ap.add_argument("--full", action="store_true", help="Включить тестовый прогон транскрипции (~1-2 мин)")
    args = ap.parse_args()

    print(f"Скил: {SKILL_ROOT}\n", flush=True)
    print("Проверка установки:", flush=True)

    ffmpeg_path = shutil.which("ffmpeg")
    if ffmpeg_path:
        check("ffmpeg", True, f"в PATH: {ffmpeg_path}")
    else:
        # Может быть установлен через static-ffmpeg в venv-whisper
        try:
            r = subprocess.run(
                [str(VENV_WHISPER_PY), "-c",
                 "from static_ffmpeg import add_paths; add_paths(); "
                 "import shutil; print(shutil.which('ffmpeg') or 'NOT_FOUND')"],
                capture_output=True, text=True, timeout=30,
            )
            if r.returncode == 0 and "NOT_FOUND" not in r.stdout:
                check("ffmpeg", True, f"через static-ffmpeg в venv-whisper: {r.stdout.strip()}")
            else:
                check("ffmpeg", False, "не найден ни в PATH, ни через static-ffmpeg")
        except (subprocess.SubprocessError, OSError) as e:
            check("ffmpeg", False, f"не найден в PATH, static-ffmpeg тоже недоступен: {e}")

    check_venv("whisper", VENV_WHISPER_PY,
               "import faster_whisper; print('faster-whisper', faster_whisper.__version__)")

    check_venv("sherpa", VENV_SHERPA_PY,
               "import sherpa_onnx, onnxruntime; print('sherpa_onnx', sherpa_onnx.__version__, '| onnxruntime', onnxruntime.__version__)")

    check("модель сегментации (pyannote-3.0)", SEG_MODEL.exists(),
          f"{SEG_MODEL.stat().st_size / 1e6:.1f} MB" if SEG_MODEL.exists() else "не найдена")

    check("модель эмбеддингов (3D-Speaker)", EMB_MODEL.exists(),
          f"{EMB_MODEL.stat().st_size / 1e6:.1f} MB" if EMB_MODEL.exists() else "не найдена")

    if args.full and VENV_WHISPER_PY.exists():
        run_transcribe_test()

    print("", flush=True)
    failed = [name for name, ok, _ in results if not ok]
    if failed:
        print(f"FAIL: {len(failed)} проверок не прошли: {', '.join(failed)}", flush=True)
        print("Запустите 'python scripts/setup.py' для (пере)установки.", flush=True)
        return 1
    print(f"OK: все {len(results)} проверок прошли. Скил готов к работе.", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())