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

MODEL_REVISION="541d1f99c6b0c3cd0b11a95167540bb8edefd82b"
MODEL_SHA256="3cbdc85877e668ca7b82d0d56770eb1fac76691f55d6b97545e8d61ca588d10d"
MODEL_FILENAME="parakeet-tdt-0.6b-v3.nemo"

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
REPOSITORY_ROOT="$(cd -- "$SCRIPT_DIR/../.." && pwd)"
INSTALL_ROOT="$HOME/.local/share/pi-dictate"
UNIT_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user"
MODEL_DIR="$INSTALL_ROOT/models/parakeet-tdt-0.6b-v3"
MODEL_PATH="$MODEL_DIR/$MODEL_FILENAME"
DICTIONARY_SOURCE="$SCRIPT_DIR/coding-terms.txt"
DICTIONARY_PATH="$INSTALL_ROOT/config/coding-terms.txt"
DRY_RUN=false
MODEL_SOURCE=""
ACTION=install

usage() {
  cat <<'EOF'
Usage: install.sh [--dry-run] [--model-source PATH]
       install.sh uninstall [--dry-run]

Installs the pinned service environment, model snapshot, and systemd user units.
Without --model-source, installation downloads the exact pinned model revision.
EOF
}

run() {
  printf '+ '
  printf '%q ' "$@"
  printf '\n'
  if ! $DRY_RUN; then
    "$@"
  fi
}

is_owned_unit() {
  local destination="$1"
  local source="$2"
  [[ -L "$destination" && "$(readlink -f -- "$destination")" == "$(readlink -f -- "$source")" ]] ||
    { [[ -f "$destination" ]] && grep -qxF '# Managed by pi-dictate install.sh' "$destination"; }
}

install_owned_unit() {
  local source="$1"
  local destination="$2"
  if [[ -L "$destination" && "$(readlink -f -- "$destination")" == "$(readlink -f -- "$source")" ]]; then
    run rm "$destination"
  elif [[ -e "$destination" || -L "$destination" ]]; then
    if [[ ! -f "$destination" ]] || ! grep -qxF '# Managed by pi-dictate install.sh' "$destination"; then
      printf 'Refusing to replace unmanaged systemd unit: %s\n' "$destination" >&2
      return 1
    fi
  fi
  run install -m 0644 "$source" "$destination"
}

remove_owned_unit() {
  local destination="$1"
  local source="$2"
  if [[ -L "$destination" && "$(readlink -f -- "$destination")" == "$(readlink -f -- "$source")" ]]; then
    run rm "$destination"
  elif [[ -f "$destination" ]] && grep -qxF '# Managed by pi-dictate install.sh' "$destination"; then
    run rm "$destination"
  elif [[ -e "$destination" || -L "$destination" ]]; then
    printf 'Leaving unmanaged unit in place: %s\n' "$destination" >&2
  fi
}

while (($#)); do
  case "$1" in
    uninstall) ACTION=uninstall ;;
    --dry-run) DRY_RUN=true ;;
    --model-source)
      shift
      (($#)) || { printf '%s\n' '--model-source requires a path' >&2; exit 2; }
      MODEL_SOURCE="$1"
      ;;
    -h|--help) usage; exit 0 ;;
    *) printf 'Unknown argument: %s\n' "$1" >&2; usage >&2; exit 2 ;;
  esac
  shift
done

socket_source="$REPOSITORY_ROOT/systemd/user/pi-dictate.socket"
service_source="$REPOSITORY_ROOT/systemd/user/pi-dictate.service"
socket_link="$UNIT_DIR/pi-dictate.socket"
service_link="$UNIT_DIR/pi-dictate.service"

if [[ "$ACTION" == uninstall ]]; then
  for pair in "$socket_link:$socket_source" "$service_link:$service_source"; do
    destination="${pair%%:*}"
    source="${pair#*:}"
    if [[ ( -e "$destination" || -L "$destination" ) ]] && ! is_owned_unit "$destination" "$source"; then
      printf 'Refusing to disable or remove unmanaged systemd unit: %s\n' "$destination" >&2
      exit 1
    fi
  done
  run systemctl --user disable --now pi-dictate.socket
  remove_owned_unit "$socket_link" "$socket_source"
  remove_owned_unit "$service_link" "$service_source"
  run systemctl --user daemon-reload
  run rm -rf "$INSTALL_ROOT"
  exit 0
fi

command -v uv >/dev/null || { printf '%s\n' 'uv is required' >&2; exit 1; }
command -v systemctl >/dev/null || { printf '%s\n' 'systemctl is required' >&2; exit 1; }

run mkdir -p "$INSTALL_ROOT" "$MODEL_DIR" "$(dirname "$DICTIONARY_PATH")" "$UNIT_DIR"
run chmod 0700 "$INSTALL_ROOT" "$(dirname "$MODEL_DIR")" "$MODEL_DIR" "$(dirname "$DICTIONARY_PATH")"
if $DRY_RUN; then
  printf '+ UV_PROJECT_ENVIRONMENT=%q uv sync --frozen --no-dev --no-editable --project %q\n' "$INSTALL_ROOT/venv" "$SCRIPT_DIR"
else
  UV_PROJECT_ENVIRONMENT="$INSTALL_ROOT/venv" uv sync --frozen --no-dev --no-editable --project "$SCRIPT_DIR"
fi

MODEL_TEMP="$MODEL_DIR/.${MODEL_FILENAME}.install.$$"
cleanup_model_temp() {
  if ! $DRY_RUN; then
    rm -f -- "$MODEL_TEMP"
  fi
}
trap cleanup_model_temp EXIT

if [[ -n "$MODEL_SOURCE" ]]; then
  [[ -f "$MODEL_SOURCE" ]] || { printf 'Model source not found: %s\n' "$MODEL_SOURCE" >&2; exit 1; }
  run cp --reflink=auto "$MODEL_SOURCE" "$MODEL_TEMP"
else
  if $DRY_RUN; then
    printf '+ download nvidia/parakeet-tdt-0.6b-v3@%s to %q\n' "$MODEL_REVISION" "$MODEL_TEMP"
  else
    MODEL_REVISION="$MODEL_REVISION" MODEL_PATH="$MODEL_TEMP" \
      "$INSTALL_ROOT/venv/bin/python" - <<'PY'
import os
import shutil
from pathlib import Path
from huggingface_hub import hf_hub_download

source = hf_hub_download(
    repo_id="nvidia/parakeet-tdt-0.6b-v3",
    filename="parakeet-tdt-0.6b-v3.nemo",
    revision=os.environ["MODEL_REVISION"],
)
destination = Path(os.environ["MODEL_PATH"])
destination.parent.mkdir(parents=True, exist_ok=True)
shutil.copyfile(source, destination)
PY
  fi
fi

if ! $DRY_RUN; then
  actual_sha256="$(sha256sum "$MODEL_TEMP" | cut -d' ' -f1)"
  if [[ "$actual_sha256" != "$MODEL_SHA256" ]]; then
    rm -f "$MODEL_TEMP"
    printf 'Model checksum mismatch: expected %s, got %s\n' "$MODEL_SHA256" "$actual_sha256" >&2
    exit 1
  fi
fi
run chmod 0400 "$MODEL_TEMP"
run mv -f "$MODEL_TEMP" "$MODEL_PATH"
trap - EXIT

if [[ ! -e "$DICTIONARY_PATH" ]]; then
  run cp "$DICTIONARY_SOURCE" "$DICTIONARY_PATH"
  run chmod 0600 "$DICTIONARY_PATH"
else
  printf 'Keeping existing coding dictionary: %s\n' "$DICTIONARY_PATH"
fi

install_owned_unit "$socket_source" "$socket_link"
install_owned_unit "$service_source" "$service_link"
run systemctl --user daemon-reload
run systemctl --user enable --now pi-dictate.socket

if $DRY_RUN; then
  printf 'Would install pi-dictate with model revision %s.\n' "$MODEL_REVISION"
else
  printf 'Installed pi-dictate with model revision %s.\n' "$MODEL_REVISION"
  printf 'Check it with: systemctl --user status pi-dictate.socket pi-dictate.service\n'
fi