#!/usr/bin/env bash
#
# SWAT+ Editor installer / updater for Linux.
#
#   curl -fsSL https://install.swat.tamu.edu/spe.sh | bash
#
# Installs the SWAT+ Editor application directory to
# ~/.local/share/SWATPlus/SWATPlusEditor and wires it into the desktop.
# Re-running it updates in place using a file manifest:
#
#   * a file the package didn't change  -> left alone (your edits survive)
#   * a file the package changed         -> downloaded fresh in full (replaced)
#   * a file you added (not in manifest) -> never touched
#   * a file the package removed         -> deleted
#
#   spe.sh             install or update (asks about the WGN database)
#   spe.sh --with-wgn  install/update and also install the WGN database, no prompt
#   spe.sh --no-wgn    install/update and skip the WGN database, no prompt
#   spe.sh --uninstall remove the app, launcher, desktop entry and icon
#
# Non-interactive WGN choice can also be set with SWATPLUS_WGN=yes|no.
#
set -euo pipefail

# Download mirrors, tried in order: install.swat.tamu.edu, then install.swatplus.link.
BASES=(
  "https://install.swat.tamu.edu/downloads"
  "https://install.swatplus.link/downloads"
)
RELFEED="speditor/releases"   # feed path under each base
APP_NAME="SWAT+ Editor"
MAIN_EXE="SWATPlusEditor"
INSTALL_DIR="$HOME/.local/share/SWATPlus/SWATPlusEditor"
MANIFEST="$INSTALL_DIR/.manifest.txt"
BIN_DIR="$HOME/.local/bin"
LAUNCHER="$BIN_DIR/swatplus-editor"
DESKTOP_DIR="$HOME/.local/share/applications"
DESKTOP_FILE="$DESKTOP_DIR/swatplus-editor.desktop"
ICON_DIR="$HOME/.local/share/icons/hicolor/256x256/apps"
ICON_FILE="$ICON_DIR/swatplus-editor.png"
ICON_REL="swatplus-editor.png"               # icon lives at the mirror root
HICOLOR="$HOME/.local/share/icons/hicolor"
# The WGN (weather generator) reference database. The editor looks for it in a
# few places; the canonical one is ~/SWATPlus/Databases (the Unix analog of the
# Windows C:/SWAT/SWATPlus/Databases layout) on macOS, and a "Databases" dir two
# levels above the app folder on Linux. Both match the editor's discovery.
WGN_ZIP="swatplus_wgn.zip"                    # lives at the mirror root
DATASETS_REL="4.0/swatplus_datasets.sqlite"   # reference datasets DB, under the mirror base
if [ "$(uname -s)" = "Darwin" ]; then
  WGN_DIR="$HOME/SWATPlus/Databases"
else
  WGN_DIR="$(dirname "$(dirname "$INSTALL_DIR")")/Databases"
fi
WGN_DB="$WGN_DIR/swatplus_wgn.sqlite"
DATASETS_DB="$WGN_DIR/swatplus_datasets.sqlite"

say()  { printf '\033[1;36m==>\033[0m %s\n' "$*"; }
warn() { printf '\033[1;33m!! \033[0m %s\n' "$*" >&2; }
die()  { printf '\033[1;31mxx \033[0m %s\n' "$*" >&2; exit 1; }

# --- options -----------------------------------------------------------------
UNINSTALL=no
WANT_WGN="${SWATPLUS_WGN:-ask}"
case "$WANT_WGN" in
  1|y|Y|yes|Yes|YES|true|TRUE)  WANT_WGN=yes ;;
  0|n|N|no|No|NO|false|FALSE)    WANT_WGN=no ;;
  *)                             WANT_WGN=ask ;;
esac
for arg in "$@"; do
  case "$arg" in
    --uninstall) UNINSTALL=yes ;;
    --with-wgn)  WANT_WGN=yes ;;
    --no-wgn)    WANT_WGN=no ;;
  esac
done

GOOD_BASE=""   # the mirror that last served a valid file; tried first next time

# Portable sha256 (Linux: sha256sum; macOS: shasum).
sha256() {
  if command -v sha256sum >/dev/null 2>&1; then sha256sum "$1" | cut -d' ' -f1
  else shasum -a 256 "$1" | cut -d' ' -f1; fi
}

# Human-readable byte count (e.g. 180M).
hbytes() {
  awk -v b="${1:-0}" 'BEGIN{
    if (b>=1073741824) printf "%.1fG", b/1073741824;
    else if (b>=1048576) printf "%.0fM", b/1048576;
    else if (b>=1024)    printf "%.0fK", b/1024;
    else                 printf "%dB", b;
  }'
}

# Portable file size in bytes (GNU stat, then BSD/macOS stat).
_fsize() { stat -c %s "$1" 2>/dev/null || stat -f %z "$1" 2>/dev/null || echo 0; }

# _looks_valid <kind> <file> [expected-sha]
# Some mirrors answer 200 with an HTML page for missing files (soft-404), so a
# 200 is not enough: validate the content matches what we asked for.
_looks_valid() {
  local kind="$1" f="$2" exp="${3:-}"
  [ -s "$f" ] || return 1
  case "$kind" in
    gz)       head -c2 "$f" | od -An -tx1 | tr -d ' \n' | grep -qi '^1f8b' ;;
    zip)      head -c2 "$f" | od -An -tx1 | tr -d ' \n' | grep -qi '^504b' ;;
    png)      head -c8 "$f" | od -An -tx1 | tr -d ' \n' | grep -qi '^89504e470d0a1a0a' ;;
    sqlite)   head -c16 "$f" | od -An -tx1 | tr -d ' \n' | grep -qi '^53514c69746520666f726d6174203300' ;;
    manifest) grep -qE '^[0-9a-f]{64}'$'\t' "$f" ;;
    version)  grep -qE '^v?[0-9]+\.[0-9]+' "$f" ;;
    sha)      [ "$(sha256 "$f")" = "$exp" ] ;;
    *)        return 0 ;;
  esac
}

# Compact, terminal-width-adaptive download progress. The bar is capped to a
# short width and only drawn when stderr is a TTY; piped/logged runs get a
# single quiet line instead of thousands of carriage-return updates.
_bar_width() {
  local cols w
  cols="$(tput cols 2>/dev/null || echo "${COLUMNS:-80}")"
  w=$(( cols - 22 ))
  [ "$w" -gt 28 ] && w=28
  [ "$w" -lt 8 ]  && w=8
  printf '%s' "$w"
}
_draw_bar() {  # <done> <total>
  local done="$1" total="$2" w pct fill bar
  w="$(_bar_width)"
  pct=$(( done * 100 / total )); [ "$pct" -gt 100 ] && pct=100
  fill=$(( done * w / total ));  [ "$fill" -gt "$w" ] && fill="$w"
  bar="$(printf '%*s' "$fill" '' | tr ' ' '#')"
  printf '\r  [%-*s] %3d%% %s   ' "$w" "$bar" "$pct" "$(hbytes "$done")" >&2
}
download_show() {  # <url> <out>
  local url="$1" out="$2" total len pid rc
  total="$(curl -fsSLI --connect-timeout 8 --proto '=https' --tlsv1.2 "$url" 2>/dev/null \
            | awk 'tolower($1)=="content-length:"{v=$2} END{gsub(/[\r ]/,"",v); print v}')"
  if { [ -t 2 ] || [ -t 1 ]; } && [ -n "$total" ] && [ "$total" -ge 1 ] 2>/dev/null; then
    curl -fSL -s --connect-timeout 8 --retry 3 --retry-delay 2 --proto '=https' --tlsv1.2 "$url" -o "$out" &
    pid=$!
    while kill -0 "$pid" 2>/dev/null; do
      _draw_bar "$(_fsize "$out")" "$total"
      sleep 0.2
    done
    if wait "$pid"; then rc=0; else rc=$?; fi
    [ "$rc" -eq 0 ] && _draw_bar "$total" "$total"
    printf '\n' >&2
    return "$rc"
  fi
  printf '    downloading %s (%s) ...\n' "$(basename "$url")" "$(hbytes "${total:-0}")" >&2
  curl -fSL -s --connect-timeout 8 --retry 3 --retry-delay 2 --proto '=https' --tlsv1.2 "$url" -o "$out"
}

# fetch <rel-under-base> <out> <kind> [expected-sha] [progress]
# Tries each mirror; accepts the first that serves a *valid* file.
fetch() {
  local rel="$1" out="$2" kind="$3" exp="${4:-}" prog="${5:-}" base order=()
  if [ -n "$GOOD_BASE" ]; then
    order=("$GOOD_BASE"); for base in "${BASES[@]}"; do [ "$base" != "$GOOD_BASE" ] && order+=("$base"); done
  else
    order=("${BASES[@]}")
  fi
  # Bundle paths contain spaces (e.g. "Electron Framework.framework"), and curl
  # rejects a URL with a raw space. Encode them; '/' and the rest stay as-is.
  local relenc="${rel// /%20}"
  for base in "${order[@]}"; do
    if [ "$prog" = progress ]; then
      download_show "$base/$relenc" "$out" || continue
    else
      curl -fsSL --connect-timeout 8 --retry 3 --retry-delay 2 --proto '=https' --tlsv1.2 "$base/$relenc" -o "$out" 2>/dev/null || continue
    fi
    if _looks_valid "$kind" "$out" "$exp"; then GOOD_BASE="$base"; return 0; fi
  done
  return 1
}

TMP="$(mktemp -d)"
trap 'rm -rf "$TMP"' EXIT
command -v curl >/dev/null 2>&1 || die "curl is required but not installed."
command -v tar  >/dev/null 2>&1 || die "tar is required but not installed."

# ---- Optional WGN reference database (shared by macOS and Linux) -------------
# Ask about the WGN database (unless a flag or SWATPLUS_WGN already decided, or
# there is no terminal to ask on).
ask_wgn() {
  [ "$WANT_WGN" = ask ] || return 0
  if { [ -t 2 ] || [ -t 1 ]; } && [ -e /dev/tty ]; then
    printf '\n' >&2
    say "Optional: the SWAT+ WGN (weather generator) reference database (~180 MB)."
    say "It lets weather generator import work out of the box."
    printf '    Download and install it now? [y/N] ' >&2
    local reply=""
    read -t 120 -r reply < /dev/tty || reply=""
    case "$reply" in [Yy]*) WANT_WGN=yes ;; *) WANT_WGN=no ;; esac
  else
    WANT_WGN=no
  fi
}

# Download + install the WGN database into $WGN_DIR (created if needed). If a copy
# is already there, leave it. Honours --with-wgn / --no-wgn / SWATPLUS_WGN.
install_wgn() {
  [ "$WANT_WGN" = yes ] || { say "Skipping WGN database (add it later with: spe.sh --with-wgn)."; return 0; }
  if [ -f "$WGN_DB" ]; then
    say "WGN database already present: $WGN_DB"
    return 0
  fi
  mkdir -p "$WGN_DIR"
  say "Downloading WGN database (~180 MB)"
  fetch "$WGN_ZIP" "$TMP/wgn.zip" zip "" progress || die "Could not download $WGN_ZIP from any mirror."
  say "Extracting WGN database"
  local tmpd="$TMP/wgn"; mkdir -p "$tmpd"
  if command -v unzip >/dev/null 2>&1; then
    unzip -oq "$TMP/wgn.zip" -d "$tmpd"
  elif command -v python3 >/dev/null 2>&1; then
    python3 -m zipfile -e "$TMP/wgn.zip" "$tmpd"
  elif command -v bsdtar >/dev/null 2>&1; then
    bsdtar -xf "$TMP/wgn.zip" -C "$tmpd"
  else
    die "Need 'unzip' or python3 to extract the WGN database. Install one, then run: spe.sh --with-wgn"
  fi
  local src
  src="$(find "$tmpd" -type f -iname '*.sqlite' | head -1)"
  [ -n "$src" ] || die "The WGN archive did not contain a .sqlite file."
  mv -f "$src" "$WGN_DB"
  say "WGN database installed: $WGN_DB"
}

# Download the reference datasets database into $WGN_DIR (the same "Databases" folder
# the editor discovers). Always refreshed so the schema stays current: every project
# is created from this file, so a stale copy breaks New Project.
install_datasets() {
  mkdir -p "$WGN_DIR"
  say "Downloading the SWAT+ datasets database"
  if fetch "$DATASETS_REL" "$TMP/datasets.sqlite" sqlite; then
    mv -f "$TMP/datasets.sqlite" "$DATASETS_DB"
    say "Datasets database installed: $DATASETS_DB"
  else
    warn "Could not download the datasets database from any mirror; left any existing copy in place."
  fi
}

# ============================ macOS (Apple Silicon) ============================
# Installs SWATPlusEditor.app into ~/Applications, ad-hoc signs it, strips the
# quarantine xattr, and registers it with Launch Services. Re-running updates in
# place via the manifest (same skip-unchanged / full-if-changed model as Linux).
if [ "$(uname -s)" = "Darwin" ]; then
  MAC_CHANNEL="osx-arm64"
  MAC_TARBALL="SWATPlusEditor-osx-arm64.tar.gz"
  APP="$HOME/Applications/SWATPlusEditor.app"
  MAC_MANIFEST="$HOME/Applications/.SWATPlusEditor.manifest.txt"   # sidecar, outside the bundle
  LSREGISTER="/System/Library/Frameworks/CoreServices.framework/Versions/A/Frameworks/LaunchServices.framework/Versions/A/Support/lsregister"

  mac_register() {
    # Extracting/updating files inside a bundle invalidates its signature, and a
    # downloaded app is quarantined, so re-sign (ad-hoc), drop quarantine, and
    # force Launch Services to index it.
    codesign --force --deep -s - "$APP" 2>/dev/null || warn "codesign failed (the app may still run)"
    xattr -dr com.apple.quarantine "$APP" 2>/dev/null || true
    [ -x "$LSREGISTER" ] && "$LSREGISTER" -f "$APP" 2>/dev/null || true
  }

  if [ "$UNINSTALL" = yes ]; then
    say "Removing $APP_NAME"
    [ -x "$LSREGISTER" ] && "$LSREGISTER" -u "$APP" 2>/dev/null || true
    rm -rf "$APP" "$MAC_MANIFEST"
    say "Done. $APP_NAME has been removed."
    exit 0
  fi

  [ "$(uname -m)" = "arm64" ] || die "Unsupported macOS architecture: $(uname -m). Apple Silicon only."
  mkdir -p "$HOME/Applications"

  if [ ! -d "$APP" ] || [ ! -f "$MAC_MANIFEST" ]; then
    say "Installing $APP_NAME (first install)"
    fetch "$MAC_TARBALL" "$TMP/app.tgz" gz "" progress || die "Could not download $MAC_TARBALL from any mirror."
    rm -rf "$APP"
    tar -xzf "$TMP/app.tgz" -C "$HOME/Applications"
    fetch "$RELFEED/$MAC_CHANNEL/manifest.txt" "$MAC_MANIFEST" manifest \
      || warn "manifest not recorded; the next run will re-download the full app"
  else
    say "Updating $APP_NAME (syncing changed files)"
    fetch "$RELFEED/$MAC_CHANNEL/manifest.txt" "$TMP/manifest.txt" manifest \
      || die "Could not fetch manifest from any mirror."
    changed=0; kept=0
    while IFS=$'\t' read -r sha mode rel; do
      [ -n "${rel:-}" ] || continue
      oldsha="$(awk -F'\t' -v r="$rel" '$3==r{print $1; exit}' "$MAC_MANIFEST")"
      if [ "$oldsha" = "$sha" ] && [ -f "$APP/$rel" ]; then kept=$((kept+1)); continue; fi
      mkdir -p "$APP/$(dirname "$rel")"
      fetch "$RELFEED/$MAC_CHANNEL/current/$rel" "$APP/$rel" sha "$sha" \
        || die "Could not download '$rel' from any mirror."
      chmod "$mode" "$APP/$rel" 2>/dev/null || true
      changed=$((changed+1))
    done < "$TMP/manifest.txt"
    # Remove files the package no longer ships (never user-added ones).
    while IFS=$'\t' read -r _sha _mode rel; do
      [ -n "${rel:-}" ] || continue
      awk -F'\t' -v r="$rel" '$3==r{f=1} END{exit !f}' "$TMP/manifest.txt" \
        || { [ -f "$APP/$rel" ] && rm -f "$APP/$rel"; }
    done < "$MAC_MANIFEST"
    cp -f "$TMP/manifest.txt" "$MAC_MANIFEST"
    say "Updated: $changed replaced, $kept unchanged"
  fi

  say "Signing and registering with Launch Services"
  mac_register

  # Optional WGN database -> ~/SWATPlus/Databases (where the editor looks on macOS).
  ask_wgn
  install_wgn
  install_datasets

  say "Installed: $APP"
  say "Find 'SWAT+ Editor' in Launchpad or Spotlight."
  exit 0
fi

# ============================ Linux (folder install) ===========================

uninstall() {
  say "Removing $APP_NAME"
  rm -rf "$INSTALL_DIR"
  rm -f "$LAUNCHER" "$DESKTOP_FILE" "$ICON_FILE"
  command -v update-desktop-database >/dev/null 2>&1 && update-desktop-database "$DESKTOP_DIR" 2>/dev/null || true
  command -v gtk-update-icon-cache   >/dev/null 2>&1 && gtk-update-icon-cache -f "$HICOLOR" 2>/dev/null || true
  say "Done. $APP_NAME has been removed."
  [ -f "$WGN_DB" ] && say "Left the shared WGN database in place: $WGN_DB"
  exit 0
}
[ "$UNINSTALL" = yes ] && uninstall

[ "$(uname -s)" = "Linux" ] || die "This installer is for Linux. See install.swatplus.link for other platforms."

# Architecture -> channel + tarball name (channels match the app's update channels).
case "$(uname -m)" in
  x86_64|amd64)  CHANNEL=linux-amd64; TARBALL=SWATPlusEditor-linux-x64.tar.gz ;;
  aarch64|arm64) CHANNEL=linux-arm64; TARBALL=SWATPlusEditor-linux-arm64.tar.gz ;;
  *) die "Unsupported architecture: $(uname -m). Supported: x86_64, aarch64." ;;
esac

mkdir -p "$INSTALL_DIR" "$BIN_DIR" "$DESKTOP_DIR" "$ICON_DIR"

# Ask about the WGN database up front, before the long app download.
ask_wgn

apply_modes_and_manifest() {
  # Re-apply exec bits from the manifest (tar preserves them on first install, but
  # per-file updates need them set explicitly) and record the new manifest.
  while IFS=$'\t' read -r _sha mode rel; do
    [ -n "${rel:-}" ] || continue
    [ -f "$INSTALL_DIR/$rel" ] && chmod "$mode" "$INSTALL_DIR/$rel" 2>/dev/null || true
  done < "$TMP/manifest.txt"
  cp -f "$TMP/manifest.txt" "$MANIFEST"
  chmod +x "$INSTALL_DIR/$MAIN_EXE" 2>/dev/null || true
}

say "Fetching manifest for $CHANNEL"
fetch "$RELFEED/$CHANNEL/manifest.txt" "$TMP/manifest.txt" manifest \
  || die "Could not fetch manifest from any mirror. Is $CHANNEL published yet?"

if [ ! -f "$MANIFEST" ]; then
  # ---- Fresh install: pull the whole tarball (fast), then record the manifest.
  say "Installing $APP_NAME (first install)"
  fetch "$TARBALL" "$TMP/app.tar.gz" gz "" progress || die "Could not download $TARBALL from any mirror."
  tar -xzf "$TMP/app.tar.gz" -C "$INSTALL_DIR"
  apply_modes_and_manifest
else
  # ---- Update: download only files whose package hash changed.
  say "Updating $APP_NAME (syncing changed files)"
  declare -A OLD=()
  while IFS=$'\t' read -r sha _mode rel; do
    [ -n "${rel:-}" ] && OLD["$rel"]="$sha"
  done < "$MANIFEST"

  changed=0; kept=0
  while IFS=$'\t' read -r sha mode rel; do
    [ -n "${rel:-}" ] || continue
    if [ "${OLD[$rel]:-}" = "$sha" ] && [ -f "$INSTALL_DIR/$rel" ]; then
      kept=$((kept+1)); continue        # unchanged in package -> keep (your edits survive)
    fi
    mkdir -p "$INSTALL_DIR/$(dirname "$rel")"
    fetch "$RELFEED/$CHANNEL/current/$rel" "$INSTALL_DIR/$rel" sha "$sha" \
      || die "Could not download '$rel' from any mirror."
    chmod "$mode" "$INSTALL_DIR/$rel" 2>/dev/null || true
    changed=$((changed+1))
  done < "$TMP/manifest.txt"

  # Remove files the package used to own but no longer ships (never user-added ones).
  declare -A NEW=()
  while IFS=$'\t' read -r _sha _mode rel; do
    [ -n "${rel:-}" ] && NEW["$rel"]=1
  done < "$TMP/manifest.txt"
  removed=0
  for rel in "${!OLD[@]}"; do
    if [ -z "${NEW[$rel]:-}" ] && [ -f "$INSTALL_DIR/$rel" ]; then
      rm -f "$INSTALL_DIR/$rel"; removed=$((removed+1))
    fi
  done

  apply_modes_and_manifest
  say "Updated: $changed replaced, $kept unchanged, $removed removed"
fi

say "Installing icon"
fetch "$ICON_REL" "$ICON_FILE" png || warn "Could not fetch the application icon (the launcher will still work)."

say "Creating launcher: $LAUNCHER"
cat > "$LAUNCHER" <<EOF
#!/usr/bin/env bash
exec "$INSTALL_DIR/$MAIN_EXE" "\$@"
EOF
chmod +x "$LAUNCHER"

say "Creating desktop entry"
cat > "$DESKTOP_FILE" <<EOF
[Desktop Entry]
Type=Application
Name=$APP_NAME
Comment=Edit, write and run SWAT+ model inputs
Exec="$INSTALL_DIR/$MAIN_EXE" %F
Icon=swatplus-editor
Categories=Science;Education;Geography;
Terminal=false
StartupNotify=true
StartupWMClass=SWATPlusEditor
EOF

command -v update-desktop-database >/dev/null 2>&1 && update-desktop-database "$DESKTOP_DIR" 2>/dev/null || true
command -v gtk-update-icon-cache   >/dev/null 2>&1 && gtk-update-icon-cache -f "$HICOLOR" 2>/dev/null || true

# ---- Optional WGN reference database -----------------------------------------
install_wgn

# ---- Reference datasets database (always installed) --------------------------
install_datasets

say "Installed to $INSTALL_DIR"
case ":$PATH:" in
  *":$BIN_DIR:"*) say "Launch with: swatplus-editor  (or from your application menu)" ;;
  *) say "Launch from your application menu, or run: $LAUNCHER"
     warn "$BIN_DIR is not on PATH; add it to use the 'swatplus-editor' command." ;;
esac