#!/usr/bin/env bash
set -euo pipefail

python -m pip install \
  "lhotse==1.32.1" \
  "kaldialign==0.9.3"

K2_WHEEL="$(python - <<'PY'
import platform
import sys

import torch

torch_version = torch.__version__.split("+", 1)[0]
python_tag = f"cp{sys.version_info.major}{sys.version_info.minor}"
machine = platform.machine().lower()
system = platform.system().lower()

wheel_version = "1.24.4.dev20251118"
base = "https://huggingface.co/csukuangfj/k2/resolve/main"

if torch_version != "2.9.1":
    raise SystemExit(
        "ZIPA setup requires the k2 CPU wheel to match torch==2.9.1, "
        f"but found torch=={torch_version}."
    )

if python_tag not in {"cp310", "cp311", "cp312", "cp313", "cp314"}:
    raise SystemExit(f"Unsupported Python tag for ZIPA k2 wheel: {python_tag}")

if system == "linux":
    if machine in {"x86_64", "amd64"}:
        platform_dir = "cpu/1.24.4.dev20250307/linux-x64"
        platform_tag = "manylinux_2_27_x86_64.manylinux_2_28_x86_64"
    elif machine in {"aarch64", "arm64"}:
        platform_dir = "cpu/1.24.4.dev20250307/linux-arm64"
        platform_tag = "manylinux_2_27_aarch64.manylinux_2_28_aarch64"
    else:
        raise SystemExit(f"Unsupported Linux machine for ZIPA k2 wheel: {machine}")
elif system == "darwin" and machine in {"arm64", "aarch64"}:
    platform_dir = "macos"
    platform_tag = "macosx_11_0_arm64"
else:
    raise SystemExit(
        f"Unsupported platform for automatic ZIPA k2 setup: {system}/{machine}"
    )

wheel = (
    f"k2-{wheel_version}+cpu.torch{torch_version}-"
    f"{python_tag}-{python_tag}-{platform_tag}.whl"
)
print(f"{base}/{platform_dir}/{wheel}")
PY
)"

python -m pip install "${K2_WHEEL}"

python - <<'PY'
import importlib

for module in ("k2", "lhotse", "kaldialign"):
    importlib.import_module(module)

print("ZIPA optional dependencies are installed.")
PY