#!/usr/bin/env bash
# Install the Rust collector (rust/collector, `ping-collector`) and move work
# from the Python app to it, or back.
#
#   sudo bash scripts/install_collector.sh                     # install binary + unit; checks only
#   sudo bash scripts/install_collector.sh --cutover           # Rust collects (COLLECTION_ENABLED=false)
#   sudo bash scripts/install_collector.sh --cutover-forecast  # Rust also forecasts and sends the
#                                                              # ntfy notifications (FORECAST_IN_COLLECTOR=true)
#   sudo bash scripts/install_collector.sh --rollback-forecast # forecasts + notifications back to Python
#   sudo bash scripts/install_collector.sh --rollback          # Python does everything again
#
# Every install also exports the trained models for the Rust forecaster and
# compares its forecast with Python's on the live data
# (scripts/forecast_parity.py); --cutover-forecast refuses to switch when they
# differ (FORCE=1 overrides).
#
# The binary comes from BIN=/path/to/ping-collector (e.g. the "ping-collector-
# linux-x86_64" CI artifact) or from rust/target/release/ping-collector, built
# beforehand as a normal user:
#   cargo build --release --locked --manifest-path rust/Cargo.toml
# (needs Rust >= 1.88 from rustup -- distro cargo is too old -- plus pkg-config
# and libssl-dev; see rust/README.md).
set -euo pipefail

APP_DIR="${APP_DIR:-/opt/thailand-water-monitor}"
SERVICE_USER="${SERVICE_USER:-water-monitor}" # as in scripts/ping-collector.service
ENV_FILE="${APP_DIR}/.env"
STATUS_FILE="${APP_DIR}/collector-status.json"
PYTHON="${APP_DIR}/.venv/bin/python"
REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
MODE="${1:-install}"

log() { printf '\033[1;32m==>\033[0m %s\n' "$*"; }
warn() { printf '\033[1;33m[warn]\033[0m %s\n' "$*" >&2; }
die() { printf '\033[1;31m[error]\033[0m %s\n' "$*" >&2; exit 1; }

[ "$(id -u)" -eq 0 ] || die "run as root (sudo)"
[ -f "$ENV_FILE" ] || die "$ENV_FILE not found"

set_env() {
    local key="$1" value="$2"
    if grep -q "^${key}=" "$ENV_FILE"; then
        sed -i "s|^${key}=.*|${key}=${value}|" "$ENV_FILE"
    else
        printf '%s=%s\n' "$key" "$value" >>"$ENV_FILE"
    fi
}

# Run a command the way the services run: same user, working directory and
# EnvironmentFile parsing (systemd's, not a .env library's).
as_service() {
    if systemd-run --version >/dev/null 2>&1; then
        systemd-run --quiet --pipe --wait --collect \
            -p User="$SERVICE_USER" -p Group="$(id -gn "$SERVICE_USER")" \
            -p WorkingDirectory="$APP_DIR" -p EnvironmentFile="$ENV_FILE" "$@"
    else
        (cd "$APP_DIR" && runuser -u "$SERVICE_USER" -- "$@")
    fi
}

show() {
    sleep "$1"
    systemctl --no-pager --lines=0 status ping-collector.service || true
    journalctl -u ping-collector -n 25 --no-pager || true
}

case "$MODE" in
    --rollback)
        log "Stopping ping-collector; the Python app collects, forecasts and notifies again"
        systemctl disable --now ping-collector.service 2>/dev/null || true
        set_env FORECAST_IN_COLLECTOR false
        set_env COLLECTION_ENABLED true
        systemctl restart water-monitor.service
        log "Done. Check: journalctl -u water-monitor -n 50"
        exit 0
        ;;
    --rollback-forecast)
        log "Forecasts and notifications back to the Python app; ping-collector keeps collecting"
        set_env FORECAST_IN_COLLECTOR false
        systemctl restart ping-collector.service water-monitor.service
        log "Done. Check: journalctl -u water-monitor -n 50"
        exit 0
        ;;
    install | --cutover | --cutover-forecast) ;;
    *) die "unknown option: $MODE" ;;
esac

BIN="${BIN:-${REPO_DIR}/rust/target/release/ping-collector}"
[ -x "$BIN" ] || die "no binary at $BIN: build it (see header) or pass BIN=/path/to/ping-collector"

log "Installing $BIN -> ${APP_DIR}/bin/ping-collector ($("$BIN" --version))"
install -D -m 0755 -o root -g root "$BIN" "${APP_DIR}/bin/ping-collector"
install -m 0644 "${REPO_DIR}/scripts/ping-collector.service" /etc/systemd/system/ping-collector.service
systemctl daemon-reload

# Both services read and write /scraping/status through this file (inside
# the app dir: the units use PrivateTmp).
set_env SCRAPING_STATUS_FILE "$STATUS_FILE"
touch "$STATUS_FILE"
chown "${SERVICE_USER}:$(id -gn "$SERVICE_USER")" "$STATUS_FILE"

log "Dry run against the configured database (writes nothing)"
as_service "${APP_DIR}/bin/ping-collector" once --dry-run

forecast_ok=1
if compgen -G "${APP_DIR}/models/flood_*.joblib" >/dev/null; then
    log "Exporting the trained models for the Rust forecaster (models/flood_*.json.gz)"
    as_service "$PYTHON" -m src.ml.export
    log "Forecast check: Python and ping-collector on the same data, moment and rain"
    as_service "$PYTHON" scripts/forecast_parity.py || forecast_ok=0
else
    warn "no trained models in ${APP_DIR}/models: forecast check skipped"
    forecast_ok=0
fi

case "$MODE" in
    --cutover)
        log "Cutover: ping-collector collects; the Python app forecasts and notifies"
        set_env COLLECTION_ENABLED false
        systemctl enable ping-collector.service
        systemctl restart ping-collector.service water-monitor.service
        show 20
        log "Next step: --cutover-forecast. Rollback: sudo bash scripts/install_collector.sh --rollback"
        ;;
    --cutover-forecast)
        grep -q '^COLLECTION_ENABLED=false' "$ENV_FILE" ||
            die "ping-collector must collect first: run --cutover"
        if [ "$forecast_ok" != 1 ] && [ "${FORCE:-0}" != 1 ]; then
            die "the forecast check did not pass; not switching (FORCE=1 overrides)"
        fi
        log "Cutover: ping-collector also forecasts and notifies; the Python app serves its forecast"
        set_env FORECAST_IN_COLLECTOR true
        systemctl enable ping-collector.service
        systemctl restart ping-collector.service water-monitor.service
        show 30
        log "Rollback if needed: sudo bash scripts/install_collector.sh --rollback-forecast"
        ;;
    *)
        if [ "$(systemctl is-active ping-collector.service 2>/dev/null)" = active ]; then
            systemctl restart ping-collector.service
            log "Installed; ping-collector restarted on the new binary"
        else
            log "Installed, not started. Switch over with: sudo bash scripts/install_collector.sh --cutover"
        fi
        ;;
esac