#!/usr/bin/env bash
#
# Retrain the flood forecast models safely. Run by water-monitor-retrain.timer
# (monthly) or by hand:  sudo systemctl start water-monitor-retrain.service
#
# Why a script rather than ExecStart=train_flood_model.py:
#   * train.py writes each station's bundle straight into models/ over ~12 min,
#     and the API's hourly precompute reloads bundles by mtime. Training into
#     a staging dir and mv-ing (atomic on one filesystem) means the API never
#     sees a half-written joblib file or a mixed old/new set.
#   * A run that produced gauge-only (v2) bundles, or trained too few stations,
#     must NOT replace the deployed models. train.py already aborts on a
#     missing rain series; this script re-checks the written metrics anyway.
#   * No API restart is needed: predict.py reloads changed bundles on the next
#     precompute (every scrape cycle, hourly), so the new models are live
#     within an hour. Restart manually if you want them live immediately.
#
# Exit codes: 0 ok, 2 training refused (see log), 3 verification failed.
set -euo pipefail

APP_DIR="${APP_DIR:-/opt/thailand-water-monitor}"
PYTHON="${PYTHON:-${APP_DIR}/.venv/bin/python}"
MODELS_DIR="${APP_DIR}/models"
STAGE_DIR="${MODELS_DIR}/.staging"
# P.4A is NOT_TRAINABLE by design (17% fill); 15 of 16 is the normal outcome.
MIN_TRAINED="${MIN_TRAINED:-14}"
EXPECT_VERSION_PREFIX="${EXPECT_VERSION_PREFIX:-hgb-v3+}"

log() { printf '%s retrain: %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" "$*"; }

cd "${APP_DIR}"
[ -x "${PYTHON}" ] || { log "no interpreter at ${PYTHON} (run uv sync)"; exit 3; }

rm -rf "${STAGE_DIR}"
mkdir -p "${STAGE_DIR}"
log "training into ${STAGE_DIR} (python=${PYTHON}, OMP_NUM_THREADS=${OMP_NUM_THREADS:-unset})"

# train_flood_model.py exits 2 on a missing rain series (RainUnavailableError)
# instead of silently writing v2 bundles -- propagate that unchanged.
set +e
"${PYTHON}" scripts/train_flood_model.py --stations all --models-dir "${STAGE_DIR}" "$@"
rc=$?
set -e
if [ "${rc}" -ne 0 ]; then
    log "training failed (exit ${rc}); deployed models untouched"
    rm -rf "${STAGE_DIR}"
    exit "${rc}"
fi

# Verify before promoting. Reads metrics.json from the stage dir.
VERSION="$("${PYTHON}" - "${STAGE_DIR}/metrics.json" <<'PY'
import json, sys
m = json.load(open(sys.argv[1]))
print(m["model_version"])
PY
)"
TRAINED="$("${PYTHON}" - "${STAGE_DIR}/metrics.json" <<'PY'
import json, sys
m = json.load(open(sys.argv[1]))
print(sum(1 for s in m["stations"].values() if s.get("status") == "trained"))
PY
)"
log "staged model_version=${VERSION} trained_stations=${TRAINED}"

case "${VERSION}" in
    "${EXPECT_VERSION_PREFIX}"*) ;;
    *)
        log "REFUSING to deploy: version '${VERSION}' does not start with '${EXPECT_VERSION_PREFIX}'"
        rm -rf "${STAGE_DIR}"
        exit 3
        ;;
esac
if [ "${TRAINED}" -lt "${MIN_TRAINED}" ]; then
    log "REFUSING to deploy: only ${TRAINED} stations trained (< ${MIN_TRAINED})"
    rm -rf "${STAGE_DIR}"
    exit 3
fi

# Promote: per-file rename is atomic; readers see either the old or the new
# bundle, never a partial one. Keep one previous generation for rollback.
mkdir -p "${MODELS_DIR}/.previous"
for f in "${STAGE_DIR}"/flood_*.joblib "${STAGE_DIR}/metrics.json"; do
    name="$(basename "${f}")"
    if [ -f "${MODELS_DIR}/${name}" ]; then
        mv -f "${MODELS_DIR}/${name}" "${MODELS_DIR}/.previous/${name}"
    fi
    mv -f "${f}" "${MODELS_DIR}/${name}"
done
rm -rf "${STAGE_DIR}"
log "deployed ${VERSION} (${TRAINED} stations); previous generation in models/.previous. The API picks it up on its next hourly precompute."