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

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
source "${ROOT}/.env"

ts(){ date -Is; }

log_dir="${ROOT}/logs"
mkdir -p "$log_dir"
log_file="${log_dir}/update-$(date -Iseconds).log"
exec > >(tee -a "$log_file") 2>&1

echo "[$(ts)] update: starting (ROOT=$ROOT)"

echo "[$(ts)] update: pulling docker image(s)"
docker compose -f "${ROOT}/docker-compose.yml" pull

echo "[$(ts)] update: restarting services"
docker compose -f "${ROOT}/docker-compose.yml" up -d

echo "[$(ts)] update: waiting for Ollama API at ${OLLAMA_URL} ..."
for i in {1..120}; do
  if curl -sS "${OLLAMA_URL}/api/tags" >/dev/null 2>&1; then
    echo "[$(ts)] update: Ollama API is up."
    break
  fi
  if [[ $i -eq 120 ]]; then
    echo "[$(ts)] update: ERROR: API did not come up in time."
    exit 1
  fi
  sleep 1
done

echo "[$(ts)] update: pulling base model: ${BASE_MODEL}"
docker exec -it ollama ollama pull "${BASE_MODEL}"

if [[ -n "${EXTRA_MODELS:-}" ]]; then
  for m in ${EXTRA_MODELS}; do
    echo "[$(ts)] update: pulling extra model: $m"
    docker exec -it ollama ollama pull "$m"
  done
fi

echo "[$(ts)] update: rebuilding expert model: ${EXPERT_MODEL}"

tmp="$(mktemp)"
sed "s/\${BASE_MODEL}/${BASE_MODEL}/g" "${ROOT}/Modelfile" > "$tmp"

docker cp "$tmp" ollama:/tmp/Modelfile.jr-sql-expert
docker exec -it ollama ollama create "${EXPERT_MODEL}" -f /tmp/Modelfile.jr-sql-expert

rm -f "$tmp"

echo "[$(ts)] update: verifying model exists..."
docker exec -it ollama ollama list | grep -F "${EXPERT_MODEL}" >/dev/null && \
  echo "[$(ts)] update: OK: ${EXPERT_MODEL} is available."

# Warmup (loads model + GPU kernels, reduces first real query latency)
if [[ -x "${ROOT}/bin/sqlai" ]]; then
  echo "[$(ts)] update: warmup: sending a short request (no-metrics)..."
  "${ROOT}/bin/sqlai" ask --text "Warmup: reply with exactly 'OK'." --no-metrics || {
    echo "[$(ts)] update: WARN: warmup failed (continuing)."
  }
  echo "[$(ts)] update: warmup: done"
else
  echo "[$(ts)] update: WARN: ${ROOT}/bin/sqlai not executable; skipping warmup."
fi

echo "[$(ts)] update: complete"