#!/usr/bin/env python3
"""Скачивает модели распознавания и диаризации в папку models.

Запускается один раз после распаковки. Около 900 МБ.
"""
import os
import sys
import tarfile
from pathlib import Path

# Встраиваемая сборка Python идёт без хранилища сертификатов, поэтому urllib
# на HTTPS падает. Берём requests: он ходит через certifi, который уже в пакете.
import requests

BASE = Path(os.environ.get("TALKSCORE_ASR_HOME") or Path(__file__).resolve().parent)
MODELS = BASE / "models"

SEG_URL = ("https://github.com/k2-fsa/sherpa-onnx/releases/download/"
           "speaker-segmentation-models/sherpa-onnx-pyannote-segmentation-3-0.tar.bz2")
EMB_URL = ("https://github.com/k2-fsa/sherpa-onnx/releases/download/"
           "speaker-recongition-models/nemo_en_titanet_small.onnx")


def setup_console() -> None:
    """Windows-консоль по умолчанию не в UTF-8, иначе русский текст рассыпается."""
    for stream in (sys.stdout, sys.stderr):
        try:
            stream.reconfigure(encoding="utf-8", errors="replace")
        except (AttributeError, ValueError):
            pass


def download(url: str, dst: Path, label: str) -> None:
    if dst.exists() and dst.stat().st_size > 0:
        print(f"  {label}: уже на месте")
        return
    dst.parent.mkdir(parents=True, exist_ok=True)
    tmp = dst.with_suffix(dst.suffix + ".part")
    with requests.get(url, stream=True, timeout=60) as r:
        r.raise_for_status()
        total = int(r.headers.get("Content-Length", 0))
        done = 0
        with tmp.open("wb") as out:
            for chunk in r.iter_content(chunk_size=1 << 20):
                out.write(chunk)
                done += len(chunk)
                if total:
                    sys.stdout.write(
                        f"\r  {label}: {done / total * 100:5.1f}%  "
                        f"({done / 1e6:.0f}/{total / 1e6:.0f} МБ)")
                    sys.stdout.flush()
    # Переименование в самом конце: прерванная закачка не оставит битый файл,
    # который на следующем запуске сочли бы готовым.
    tmp.replace(dst)
    print()


def fetch_gigaam(target: Path) -> None:
    """Качает модель распознавания с Hugging Face.

    Ускоренный транспорт xet на части сетей не проходит, поэтому при неудаче
    повторяем обычным способом - он медленнее, но работает везде.
    """
    os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
    from huggingface_hub import snapshot_download

    # Только вариант e2e_rnnt: он единственный ставит пунктуацию и заглавные.
    kwargs = dict(repo_id="istupakov/gigaam-v3-onnx", local_dir=str(target),
                  allow_patterns=["*e2e_rnnt*", "config.json"])
    try:
        snapshot_download(**kwargs)
    except Exception as exc:  # noqa: BLE001 - причина может быть любой, важен откат
        print(f"\n  Быстрая загрузка не удалась ({type(exc).__name__}), пробую обычную")
        os.environ["HF_HUB_DISABLE_XET"] = "1"
        snapshot_download(**kwargs)


def main() -> int:
    setup_console()
    MODELS.mkdir(parents=True, exist_ok=True)

    print("Диаризация:")
    download(EMB_URL, MODELS / "nemo_en_titanet_small.onnx", "голосовые отпечатки")

    seg_dir = MODELS / "sherpa-onnx-pyannote-segmentation-3-0"
    if (seg_dir / "model.onnx").is_file():
        print("  сегментация: уже на месте")
    else:
        archive = MODELS / "seg.tar.bz2"
        download(SEG_URL, archive, "сегментация")
        with tarfile.open(archive, "r:bz2") as tar:
            tar.extractall(MODELS)
        archive.unlink()

    print("\nРаспознавание (GigaAM v3, около 850 МБ):")
    target = MODELS / "gigaam"
    if (target / "config.json").is_file():
        print("  уже на месте")
    else:
        fetch_gigaam(target)

    print("\nГотово. Теперь запустите start.bat")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except requests.RequestException as exc:
        print(f"\nОшибка сети: {exc}")
        print("Проверьте подключение и доступность github.com и huggingface.co")
        sys.exit(1)