#!/bin/bash
# tqLLM installer: curl -fsSL https://tqllm.com/install.sh | bash
#   macOS on Apple Silicon: the native Metal build (Homebrew, then pip, then a source build; see --help)
#   Linux with an NVIDIA GPU (x86-64 or arm64 / DGX Spark, also inside WSL 2): the Docker image
# Options go after `bash -s --`, e.g. `curl -fsSL https://tqllm.com/install.sh | bash -s -- --check`.

# ===================================== Linux + NVIDIA =====================================
linux_main() {
  set -euo pipefail
  VERSION="1.4.4"; DRY=0
  while [ $# -gt 0 ]; do
    case "$1" in
      --dry-run|--check) DRY=1 ;;
      --version) VERSION="$2"; shift ;;
      -h|--help) echo "tqLLM installer (Linux + NVIDIA): --check / --dry-run prints the commands, --version X.Y.Z"; exit 0 ;;
      *) echo "unknown option: $1" >&2; exit 2 ;;
    esac
    shift
  done
  OS="$(uname -s)"; ARCH="$(uname -m)"
  say "tqLLM $VERSION installer ($OS $ARCH)"
  install_linux
}
say()  { printf '\033[1m==>\033[0m %s\n' "$*"; }
warn() { printf '\033[33mwarning:\033[0m %s\n' "$*" >&2; }
die()  { printf '\033[31merror:\033[0m %s\n' "$*" >&2; exit 1; }
run()  { printf '  $ %s\n' "$*"; [ "$DRY" = 1 ] || "$@"; }
install_linux() {
  case "$ARCH" in x86_64|aarch64|arm64) ;; *) die "Linux on $ARCH is not supported (x86-64 or arm64 with an NVIDIA GPU)." ;; esac
  if [ -r /proc/version ] && grep -qi microsoft /proc/version; then say "Windows (WSL 2) detected"; fi
  if command -v nvidia-smi >/dev/null 2>&1; then
    say "GPU: $(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null | head -1)"
  else
    warn "nvidia-smi not found: tqLLM on Linux needs an NVIDIA GPU and its driver."
  fi
  command -v docker >/dev/null 2>&1 || die "Docker is not installed. Install Docker and the NVIDIA Container Toolkit, then run this again:
  https://docs.docker.com/engine/install/  and  https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html"
  if ! docker info 2>/dev/null | grep -qi nvidia; then
    warn "Docker does not list the NVIDIA runtime: install the NVIDIA Container Toolkit if 'docker run --gpus all' fails."
  fi
  IMAGE="ghcr.io/turbo-tan/tqllm:$VERSION"
  say "Pulling $IMAGE"
  run docker pull "$IMAGE"
  cat <<EOF

Done. Put a model in ./models and start the server, for example Qwen3.8-27B for a 24 GB card:
  pip install -U huggingface_hub
  hf download YTan2000/Qwen3.8-27B-TQ3_4S Qwen3.8-27B-TQ3_4S-v2.gguf --local-dir models
  docker run --rm --gpus all -p 127.0.0.1:8086:8086 -v "\$PWD/models:/models:ro" --entrypoint tqserve \\
    $IMAGE --model /models/Qwen3.8-27B-TQ3_4S-v2.gguf --host 0.0.0.0 --port 8086
Then open http://127.0.0.1:8086 (monitor) or point an app at http://127.0.0.1:8086/v1.
Models and settings per GPU: https://tqllm.com/install
EOF
}

case "$(uname -s)" in
  Linux) linux_main "$@"; exit 0 ;;
  Darwin) ;;
  MINGW*|MSYS*|CYGWIN*) echo "error: on Windows, run this inside WSL 2 (Ubuntu) with Docker Desktop's WSL integration: https://tqllm.com/install#windows" >&2; exit 1 ;;
  *) echo "error: $(uname -s) is not supported: tqLLM runs on macOS (Apple Silicon) and Linux with an NVIDIA GPU." >&2; exit 1 ;;
esac

# ===================================== macOS (Apple Silicon) =====================================
set -u

REPO="turbo-tan/tqllm"
TAP_FORMULA="turbo-tan/tap/tqllm"
MLX_PIN="mlx==0.32.4"
FALLBACK_VERSION="1.4.4"
PREFIX="${TQLLM_PREFIX:-$HOME/.tqllm}"
METHOD="auto"; VERSION=""; WHEELS=""; INDEX_URL="${PIP_INDEX_URL:-}"; CHECK=0

say()  { printf '%s\n' "$*"; }
step() { printf '\n==> %s\n' "$*"; }
warn() { printf 'warning: %s\n' "$*" >&2; }
die()  { printf 'error: %s\n' "$*" >&2; exit 1; }
usage() {
  cat <<'EOF'
tqLLM installer (Apple Silicon Macs): tries Homebrew, then pip, then a source build.
  --check            report what this Mac has and which route will work; install nothing
  --method M         brew | pip | source (only that route; default: each in turn)
  --version V        release to install (default: the latest GitHub release)
  --wheels DIR       install Python packages only from DIR: for networks that block PyPI
  --index-url URL    a PyPI mirror (for example your company's); same as PIP_INDEX_URL
  --prefix DIR       where the pip / source routes live (default ~/.tqllm)
EOF
}

while [ $# -gt 0 ]; do
  case "$1" in
    --check) CHECK=1 ;;
    --method) shift; METHOD="${1:-}" ;;
    --version) shift; VERSION="${1:-}" ;;
    --wheels) shift; WHEELS="${1:-}" ;;
    --index-url) shift; INDEX_URL="${1:-}" ;;
    --prefix) shift; PREFIX="${1:-}" ;;
    -h|--help) usage; exit 0 ;;
    *) die "unknown option: $1 (see --help)" ;;
  esac
  shift
done
case "$METHOD" in auto|brew|pip|source) ;; *) die "--method must be brew, pip or source" ;; esac
[ -z "$WHEELS" ] || [ -d "$WHEELS" ] || die "--wheels: no such folder: $WHEELS"
[ -z "$WHEELS" ] || WHEELS="$(cd "$WHEELS" && pwd)"

# ---------------------------------------------------------------------------------------------------------------------
# What this Mac has
# ---------------------------------------------------------------------------------------------------------------------
OS="$(uname -s)"; ARCH="$(uname -m)"
[ "$ARCH" = arm64 ] || die "tqLLM needs Apple Silicon (arm64); this Mac is $ARCH"
MACOS="$(sw_vers -productVersion 2>/dev/null || echo 0)"
[ "${MACOS%%.*}" -ge 15 ] 2>/dev/null || die "tqLLM needs macOS 15 or newer (this Mac: $MACOS)"
MEM_GB=$(( $(/usr/sbin/sysctl -n hw.memsize 2>/dev/null || echo 0) / 1073741824 ))
CHIP="$(/usr/sbin/sysctl -n machdep.cpu.brand_string 2>/dev/null || echo unknown)"

mkdir -p "$PREFIX" || die "cannot create $PREFIX"
LOG="$PREFIX/install.log"
: > "$LOG"
run() { printf '+ %s\n' "$*" >> "$LOG"; "$@" >> "$LOG" 2>&1; }

HAVE_BREW=0; command -v brew >/dev/null 2>&1 && HAVE_BREW=1
HAVE_CLT=0; xcode-select -p >/dev/null 2>&1 && HAVE_CLT=1
HAVE_GIT=0; command -v git >/dev/null 2>&1 && HAVE_GIT=1

PY=""
for c in python3.12 python3.13 python3.11 python3.10 /opt/homebrew/bin/python3.12 /opt/homebrew/bin/python3 python3; do
  p="$(command -v "$c" 2>/dev/null)" || continue
  if "$p" -c 'import sys; sys.exit(0 if (3, 10) <= sys.version_info[:2] <= (3, 13) else 1)' 2>/dev/null; then PY="$p"; break; fi
done

# PyPI reachability: a file download that is redirected (to a company mirror's web page) or refused does not install
pypi_state() {
  local code loc
  code="$(curl -s -o /dev/null -m 15 -w '%{http_code}' https://pypi.org/simple/mlx/ 2>/dev/null)"
  [ "$code" = 200 ] || { echo "unreachable (pypi.org: ${code:-no answer})"; return; }
  local url
  url="$(curl -s -m 15 https://pypi.org/simple/mlx/ 2>/dev/null | grep -o 'https://files.pythonhosted.org[^"#]*mlx-0.32.4[^"#]*macosx[^"#]*\.whl' | head -1)"
  [ -n "$url" ] || { echo "ok"; return; }
  code="$(curl -s -o /dev/null -I -m 15 -w '%{http_code} %{redirect_url}' "$url" 2>/dev/null)"
  case "$code" in
    200*) echo "ok" ;;
    30*) loc="${code#* }"; case "$loc" in https://files.pythonhosted.org*|"") echo "ok" ;; *) echo "redirected to ${loc%%\?*} (a company mirror or filter)" ;; esac ;;
    *) echo "blocked (package download: ${code%% *})" ;;
  esac
}
GH_OK=0; curl -s -o /dev/null -m 15 "https://github.com/$REPO" 2>/dev/null && GH_OK=1

if [ -z "$VERSION" ]; then
  VERSION="$(curl -s -m 15 "https://api.github.com/repos/$REPO/releases/latest" 2>/dev/null | grep -m1 '"tag_name"' | sed -E 's/.*"v?([^"]+)".*/\1/')"
  [ -n "$VERSION" ] || VERSION="$FALLBACK_VERSION"
fi
TARBALL="https://github.com/$REPO/archive/refs/tags/v$VERSION.tar.gz"

if [ -n "$WHEELS" ]; then PYPI="not used (--wheels $WHEELS)"
elif [ -n "$INDEX_URL" ]; then PYPI="mirror $INDEX_URL"
else PYPI="$(pypi_state)"; fi

step "This Mac"
say "  chip $CHIP, ${MEM_GB} GB, macOS $MACOS"
say "  Homebrew: $([ $HAVE_BREW = 1 ] && echo yes || echo no)   Xcode command line tools: $([ $HAVE_CLT = 1 ] && echo yes || echo no)   git: $([ $HAVE_GIT = 1 ] && echo yes || echo no)"
say "  Python 3.10-3.13: ${PY:-none found}"
say "  GitHub: $([ $GH_OK = 1 ] && echo reachable || echo unreachable)   PyPI packages: $PYPI"
say "  tqLLM release: $VERSION"
[ "$MEM_GB" -ge 24 ] || warn "the 27B model needs about 20 GB of memory; ${MEM_GB} GB is likely too little"

PYPI_OK=1
case "$PYPI" in ok*|mirror*|"not used"*) ;; *) PYPI_OK=0 ;; esac

recommend() {
  if [ $GH_OK = 0 ]; then echo "none: GitHub is unreachable (the source tarball and the Homebrew tap live there)"; return; fi
  if [ $PYPI_OK = 0 ]; then
    echo "none until Python packages are available: pass --index-url <your PyPI mirror> or --wheels <folder with the mlx 0.32.4 wheel>"; return
  fi
  if [ $HAVE_BREW = 1 ] && [ -z "$WHEELS" ] && [ -z "$INDEX_URL" ]; then echo "brew"; return; fi
  if [ -n "$PY" ] && [ $HAVE_CLT = 1 ]; then echo "pip"; return; fi
  if [ $HAVE_CLT = 0 ]; then echo "none until the Xcode command line tools are installed: xcode-select --install"; return; fi
  echo "none until Python 3.10-3.13 is installed (python.org or: brew install python@3.12)"
}
REC="$(recommend)"
say "  recommended route: $REC"
if [ $CHECK = 1 ]; then exit 0; fi
case "$REC" in none*) [ "$METHOD" = auto ] && die "${REC#none*: }" ;; esac

# ---------------------------------------------------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------------------------------------------------
pip_opts() {
  if [ -n "$WHEELS" ]; then printf '%s\n' --no-index --find-links "$WHEELS"
  elif [ -n "$INDEX_URL" ]; then printf '%s\n' --index-url "$INDEX_URL"; fi
}
link_cli() {  # put <venv>/bin/tqllm on PATH via ~/.local/bin
  mkdir -p "$HOME/.local/bin" && ln -sf "$PREFIX/venv/bin/tqllm" "$HOME/.local/bin/tqllm"
  case ":$PATH:" in *":$HOME/.local/bin:"*) ;; *) say "  add ~/.local/bin to your PATH:  echo 'export PATH=\"\$HOME/.local/bin:\$PATH\"' >> ~/.zshrc" ;; esac
}
verify() {  # $1: the tqllm command, $2: its Python; the command answers and its package carries tqserve
  "$1" --help >/dev/null 2>&1 || { warn "$1 --help failed"; return 1; }
  "$2" -c 'import tqllm, pathlib, sys; sys.exit(0 if (pathlib.Path(tqllm.__file__).parent / "bin" / "tqserve").exists() else 1)' \
    || { warn "tqserve is missing from the installed package"; return 1; }
}
make_venv() {
  [ -n "$PY" ] || { warn "no Python 3.10-3.13"; return 1; }
  [ $HAVE_CLT = 1 ] || { warn "the Xcode command line tools are missing (xcode-select --install)"; return 1; }
  rm -rf "$PREFIX/venv" && run "$PY" -m venv "$PREFIX/venv" || return 1
  local o; o=$(pip_opts | tr '\n' ' ')
  # shellcheck disable=SC2086
  run "$PREFIX/venv/bin/python" -m pip install $o --upgrade pip >/dev/null 2>&1 || true
  # shellcheck disable=SC2086
  run "$PREFIX/venv/bin/python" -m pip install $o "$MLX_PIN" "huggingface_hub>=0.24" "scikit-build-core>=0.10" "cmake>=3.25"
}

route_brew() {
  [ $HAVE_BREW = 1 ] || { warn "Homebrew is not installed (https://brew.sh)"; return 1; }
  [ -z "$WHEELS$INDEX_URL" ] || { warn "Homebrew downloads from PyPI itself; --wheels / --index-url need the pip or source route"; return 1; }
  say "  brew install $TAP_FORMULA (builds tqserve; a few minutes)"
  run brew install "$TAP_FORMULA" || { brew list tqllm >/dev/null 2>&1 && run brew upgrade "$TAP_FORMULA"; } || return 1
  verify "$(brew --prefix)/bin/tqllm" "$(brew --prefix tqllm)/libexec/bin/python"
}
route_pip() {
  say "  Python venv in $PREFIX/venv, then pip builds tqserve from the v$VERSION source"
  make_venv || return 1
  local o; o=$(pip_opts | tr '\n' ' ')
  run curl -fL -m 600 -o "$PREFIX/tqllm-$VERSION.tar.gz" "$TARBALL" || return 1
  # shellcheck disable=SC2086
  run "$PREFIX/venv/bin/python" -m pip install $o --no-build-isolation --no-deps "$PREFIX/tqllm-$VERSION.tar.gz" || return 1
  verify "$PREFIX/venv/bin/tqllm" "$PREFIX/venv/bin/python" && link_cli
}
route_source() {
  [ $HAVE_GIT = 1 ] || { warn "git is missing (xcode-select --install)"; return 1; }
  say "  git clone v$VERSION into $PREFIX/src, CMake build of tqserve against the venv's MLX"
  make_venv || return 1
  rm -rf "$PREFIX/src" && run git clone --depth 1 --branch "v$VERSION" "https://github.com/$REPO.git" "$PREFIX/src" || return 1
  local mlxdir; mlxdir="$("$PREFIX/venv/bin/python" -m mlx --cmake-dir 2>>"$LOG")" || return 1
  local cm="$PREFIX/venv/bin/cmake"; [ -x "$cm" ] || cm=cmake
  run "$cm" -S "$PREFIX/src" -B "$PREFIX/src/build" -DCMAKE_BUILD_TYPE=Release -DTQLLM_CUDA=OFF -DTQLLM_MLX=ON -DBUILD_TESTING=OFF -DCMAKE_PREFIX_PATH="$mlxdir" || return 1
  run "$cm" --build "$PREFIX/src/build" --target tqserve -j "$(/usr/sbin/sysctl -n hw.perflevel0.physicalcpu 2>/dev/null || echo 8)" || return 1
  # the tqllm package from the same source with this tqserve inside it (what the wheel build installs), and its command
  local sp; sp="$("$PREFIX/venv/bin/python" -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" || return 1
  rm -rf "$sp/tqllm" && cp -R "$PREFIX/src/python/tqllm" "$sp/tqllm" && mkdir -p "$sp/tqllm/bin" && cp "$PREFIX/src/build/tqserve" "$sp/tqllm/bin/" || return 1
  printf '#!%s\nimport sys\nfrom tqllm.cli import main\nsys.exit(main())\n' "$PREFIX/venv/bin/python" > "$PREFIX/venv/bin/tqllm" && chmod +x "$PREFIX/venv/bin/tqllm"
  verify "$PREFIX/venv/bin/tqllm" "$PREFIX/venv/bin/python" && link_cli
}

if [ "$METHOD" = auto ]; then ROUTES="brew pip source"; else ROUTES="$METHOD"; fi
for r in $ROUTES; do
  step "Installing with $r"
  if "route_$r"; then
    step "Done: tqLLM $VERSION ($r)"
    say "  tqllm serve      first run downloads the model and drafter (~17 GB), then serves on http://127.0.0.1:8086"
    say "  tqllm claude     Claude Code on the local server (also: tqllm codex, tqllm opencode)"
    exit 0
  fi
  warn "the $r route did not work; details: $LOG"
  tail -5 "$LOG" | sed 's/^/    /' >&2
done
die "no route worked. Run with --check to see why, or send $LOG with an issue: https://github.com/$REPO/issues"