#!/usr/bin/env sh
set -eu

RESTART_SERVICE=false
RESTART_IF_CHANGED=false
START_SERVICE=false
CHANGED=false

mark_changed() {
  CHANGED=true
}

kill_stale_kare_runtime() {
  if ! command -v pgrep >/dev/null 2>&1; then
    return 0
  fi

  stale_pids="$(pgrep -f 'model-gateway-service-launcher|crun_uni-model-infer' 2>/dev/null || true)"
  if [ -z "$stale_pids" ]; then
    return 0
  fi

  kill $stale_pids 2>/dev/null || true
  sleep 1
  stale_pids="$(pgrep -f 'model-gateway-service-launcher|crun_uni-model-infer' 2>/dev/null || true)"
  if [ -n "$stale_pids" ]; then
    kill -9 $stale_pids 2>/dev/null || true
  fi
  mark_changed
}

usage() {
  cat <<'EOF'
Usage:
  sudo uni-model-infer-repair-kare [--restart|--restart-if-changed|--start]

Repairs openKylin KARE generated model-gateway systemd units. Some graphical
installers wrap normal system services under /opt/kare, which breaks the
dedicated model-gateway service user. This helper restores the native unit and
backs up KARE-generated model-gateway units before removing them from the active
systemd search path.
EOF
}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --restart)
      RESTART_SERVICE=true
      shift
      ;;
    --restart-if-changed)
      RESTART_IF_CHANGED=true
      shift
      ;;
    --start)
      START_SERVICE=true
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown argument: $1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

if [ "$(id -u)" != "0" ]; then
  echo "Please run as root: sudo uni-model-infer-repair-kare" >&2
  exit 1
fi

native_unit=""
for candidate in \
  /opt/model-gateway/share/model-gateway/systemd/model-gateway.service \
  /lib/systemd/system/model-gateway.service \
  /usr/lib/systemd/system/model-gateway.service; do
  if [ -f "$candidate" ]; then
    native_unit="$candidate"
    break
  fi
done

link_list="$(mktemp)"
for root in /etc/systemd/system /run/systemd/system; do
  [ -d "$root" ] || continue
  find "$root" -type l -name model-gateway.service -print >> "$link_list" 2>/dev/null || true
done
while IFS= read -r link_path; do
  [ -n "$link_path" ] || continue
  link_target="$(readlink -f "$link_path" 2>/dev/null || true)"
  case "$link_target" in
    /opt/kare/*)
      rm -f "$link_path"
      mark_changed
      ;;
  esac
done < "$link_list"
rm -f "$link_list"

if [ -L /etc/systemd/system/model-gateway.service ]; then
  link_target="$(readlink -f /etc/systemd/system/model-gateway.service 2>/dev/null || true)"
  case "$link_target" in
    /opt/kare/*)
      rm -f /etc/systemd/system/model-gateway.service
      mark_changed
      ;;
  esac
elif [ -f /etc/systemd/system/model-gateway.service ] &&
     grep -q '/opt/kare/' /etc/systemd/system/model-gateway.service 2>/dev/null; then
  stamp="$(date +%Y%m%d%H%M%S)"
  install -d -m 0700 /root/model-gateway-kare-backup/"$stamp"
  cp -a /etc/systemd/system/model-gateway.service \
    /root/model-gateway-kare-backup/"$stamp"/etc_systemd_system_model-gateway.service
  rm -f /etc/systemd/system/model-gateway.service
  mark_changed
fi

if [ -n "$native_unit" ]; then
  install -d -m 0755 /etc/systemd/system
  if [ ! -f /etc/systemd/system/model-gateway.service ] ||
     [ -L /etc/systemd/system/model-gateway.service ] ||
     ! cmp -s "$native_unit" /etc/systemd/system/model-gateway.service 2>/dev/null; then
    install -m 0644 "$native_unit" /etc/systemd/system/model-gateway.service
    mark_changed
  fi
fi

install -d -m 0755 /usr/share/applications /usr/share/icons/hicolor/scalable/apps
if [ -f /opt/model-gateway/share/model-gateway/desktop/uni-model-infer.desktop ]; then
  install -m 0644 /opt/model-gateway/share/model-gateway/desktop/uni-model-infer.desktop \
    /usr/share/applications/uni-model-infer.desktop
fi
if [ -f /opt/model-gateway/share/model-gateway/icons/uni-model-infer.svg ]; then
  install -m 0644 /opt/model-gateway/share/model-gateway/icons/uni-model-infer.svg \
    /usr/share/icons/hicolor/scalable/apps/uni-model-infer.svg
fi

if [ -d /opt/kare ]; then
  stamp="$(date +%Y%m%d%H%M%S)"
  backup_dir="/root/model-gateway-kare-backup/$stamp"
  found_file="$(mktemp)"
  find /opt/kare -path '*/systemd/system/model-gateway.service*' -print > "$found_file" 2>/dev/null || true
  if [ -s "$found_file" ]; then
    install -d -m 0700 "$backup_dir"
    while IFS= read -r path; do
      [ -e "$path" ] || continue
      rel="${path#/opt/kare/}"
      safe_rel="$(printf '%s' "$rel" | tr '/' '_')"
      cp -a "$path" "$backup_dir/$safe_rel"
      rm -f "$path"
      mark_changed
    done < "$found_file"
  fi
  rm -f "$found_file"
fi

if [ -e /opt/kare/opt/model-gateway ]; then
  rm -rf /opt/kare/opt/model-gateway 2>/dev/null || true
  mark_changed
fi
for stale_path in \
  /opt/kare/usr/share/applications/uni-model-infer.desktop \
  /opt/kare/usr/share/icons/hicolor/scalable/apps/uni-model-infer.svg \
  /opt/kare/usr/bin/uni-model-infer-open \
  /opt/kare/usr/bin/uni-model-infer-repair-kare \
  /opt/kare/usr/bin/model-gateway-grant-access \
  /opt/kare/usr/bin/model-gateway-resolve-model \
  /opt/kare/usr/bin/model-gateway-set-llama-server \
  /opt/kare/usr/bin/model-gateway-set-run-user; do
  if [ -e "$stale_path" ]; then
    rm -f "$stale_path" 2>/dev/null || true
    mark_changed
  fi
done

kill_stale_kare_runtime

install -d -m 0755 /etc/sudoers.d
cat > /etc/sudoers.d/model-gateway-grant-access <<'EOF'
# Allow the model-gateway service to resolve server-local GGUF paths and grant
# itself read access without collecting sudo passwords in the Web UI.
%model-gateway ALL=(root) NOPASSWD: /usr/bin/model-gateway-grant-access --path *
%model-gateway ALL=(root) NOPASSWD: /usr/bin/model-gateway-resolve-model --name * --size *
%model-gateway ALL=(root) NOPASSWD: /usr/bin/model-gateway-resolve-model --name * --size * --mtime *
EOF
chmod 0440 /etc/sudoers.d/model-gateway-grant-access
if command -v visudo >/dev/null 2>&1; then
  visudo -cf /etc/sudoers.d/model-gateway-grant-access >/dev/null
fi

if command -v systemctl >/dev/null 2>&1 && [ -d /run/systemd/system ]; then
  systemctl daemon-reload || true
  systemctl reset-failed model-gateway >/dev/null 2>&1 || true
  systemctl reenable model-gateway >/dev/null 2>&1 || systemctl enable model-gateway >/dev/null 2>&1 || true
  if [ "$RESTART_SERVICE" = "true" ] ||
     { [ "$RESTART_IF_CHANGED" = "true" ] && [ "$CHANGED" = "true" ]; }; then
    systemctl restart model-gateway || true
  elif [ "$START_SERVICE" = "true" ]; then
    systemctl start model-gateway || true
  fi
fi

exit 0
