"""Download official SenseVoiceSmall and Silero VAD sherpa-onnx models."""

from __future__ import annotations

import argparse
import hashlib
import shutil
import tarfile
import tempfile
import urllib.request
from pathlib import Path

MODEL_URL = (
    "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/"
    "sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17.tar.bz2"
)
MODEL_SHA256 = "f6b2a72ebcb1ac7a764d4cfccd886e6bcb2a95c4657c2199d0ba95ed4b9ea71a"
REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_TARGET = REPOSITORY_ROOT / "artifacts" / "models" / "sensevoice-small"
VAD_URL = (
    "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/"
    "silero_vad.onnx"
)
VAD_SHA256 = "9e2449e1087496d8d4caba907f23e0bd3f78d91fa552479bb9c23ac09cbb1fd6"
DEFAULT_VAD_TARGET = (
    REPOSITORY_ROOT / "artifacts" / "models" / "silero-vad" / "silero_vad.onnx"
)


def _download(url: str, destination: Path) -> None:
    with urllib.request.urlopen(url, timeout=60) as response, destination.open("wb") as output:
        shutil.copyfileobj(response, output, length=1024 * 1024)


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def install(target: Path) -> None:
    target = target.resolve()
    if (target / "model.onnx").is_file() and (target / "tokens.txt").is_file():
        print(f"SenseVoiceSmall is already installed at {target}")
        return
    target.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="sensevoice-install-") as temporary_directory:
        temporary_root = Path(temporary_directory)
        archive = temporary_root / "sensevoice-small.tar.bz2"
        print(f"Downloading the official SenseVoiceSmall model to {archive} ...", flush=True)
        _download(MODEL_URL, archive)
        actual_hash = _sha256(archive)
        if actual_hash != MODEL_SHA256:
            raise RuntimeError(
                f"SenseVoiceSmall archive checksum mismatch: {actual_hash} != {MODEL_SHA256}"
            )
        extracted = temporary_root / "extracted"
        extracted.mkdir()
        with tarfile.open(archive, "r:bz2") as package:
            package.extractall(extracted, filter="data")
        model_file = next(extracted.rglob("model.onnx"), None)
        if model_file is None or not (model_file.parent / "tokens.txt").is_file():
            raise RuntimeError("SenseVoiceSmall archive does not contain the expected model files")
        for item in model_file.parent.iterdir():
            destination = target / item.name
            if item.is_dir():
                shutil.copytree(item, destination, dirs_exist_ok=True)
            else:
                shutil.copy2(item, destination)
    print(f"SenseVoiceSmall installed at {target}")


def install_vad(target: Path) -> None:
    target = target.resolve()
    if target.is_file() and _sha256(target) == VAD_SHA256:
        print(f"Silero VAD is already installed at {target}")
        return
    target.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="silero-vad-install-") as temporary_directory:
        downloaded = Path(temporary_directory) / "silero_vad.onnx"
        print(f"Downloading the official Silero VAD model to {downloaded} ...", flush=True)
        _download(VAD_URL, downloaded)
        actual_hash = _sha256(downloaded)
        if actual_hash != VAD_SHA256:
            raise RuntimeError(
                f"Silero VAD checksum mismatch: {actual_hash} != {VAD_SHA256}"
            )
        shutil.copy2(downloaded, target)
    print(f"Silero VAD installed at {target}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--target", type=Path, default=DEFAULT_TARGET)
    parser.add_argument("--vad-target", type=Path, default=DEFAULT_VAD_TARGET)
    parser.add_argument("--skip-vad", action="store_true")
    args = parser.parse_args()
    install(args.target)
    if not args.skip_vad:
        install_vad(args.vad_target)


if __name__ == "__main__":
    main()