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

LLAMA_CPP_COMMIT="${TIMMY_LLAMA_CPP_COMMIT:-6d05498314db1b57f81c271080018aa2d0b89be9}"
MODEL_SHA256="${TIMMY_MODEL_SHA256:-0cf76814555b8665149075b74ab6b5c1d428ea1d3d01c1918c12012e8d7c9f58}"
MMPROJ_SHA256="${TIMMY_MMPROJ_SHA256:-ae07ea1facd07dd3230c4483b63e8cda96c6944ad2481f33d531f79e892dd024}"
HOST="${TIMMY_MODEL_HOST:-127.0.0.1}"
PORT="${TIMMY_MODEL_PORT:-8080}"
ROOT="${TIMMY_SELFHOST_ROOT:-${XDG_DATA_HOME:-$HOME/.local/share}/timmy-selfhost}"
MODEL_DIR="$ROOT/models"
MODEL="$MODEL_DIR/SmolVLM2-2.2B-Instruct-Q4_K_M.gguf"
MMPROJ="$MODEL_DIR/mmproj-SmolVLM2-2.2B-Instruct-Q8_0.gguf"
SOURCE="$ROOT/llama.cpp"
BUILD="$SOURCE/build"
HF_BASE="https://huggingface.co/ggml-org/SmolVLM2-2.2B-Instruct-GGUF/resolve/main"
SERVER="$BUILD/bin/llama-server"
RUNTIME="$ROOT/run"
PIDFILE="$RUNTIME/llama-server.pid"
LOGFILE="$RUNTIME/llama-server.log"

verify_files() {
  for file in "$MODEL" "$MMPROJ"; do
    if [[ ! -f "$file" ]]; then
      printf 'Missing required file: %s. Run %s install.\n' "$file" "$0" >&2
      return 1
    fi
  done
  printf '%s  %s\n%s  %s\n' "$MODEL_SHA256" "$MODEL" "$MMPROJ_SHA256" "$MMPROJ" | sha256sum --check --status
}

download_verified() {
  local url="$1" target="$2" expected="$3" temporary="$2.download"
  rm -f "$temporary"
  curl --fail --location --retry 3 --output "$temporary" "$url"
  printf '%s  %s\n' "$expected" "$temporary" | sha256sum --check --status || {
    rm -f "$temporary"
    printf 'Downloaded file failed pinned SHA-256 verification: %s\n' "$target" >&2
    return 1
  }
  mv "$temporary" "$target"
}

install_worker() {
  command -v git >/dev/null
  command -v cmake >/dev/null
  command -v curl >/dev/null
  mkdir -p "$ROOT" "$MODEL_DIR"
  if [[ ! -d "$SOURCE/.git" ]]; then
    git clone https://github.com/ggml-org/llama.cpp "$SOURCE"
  fi
  git -C "$SOURCE" fetch origin "$LLAMA_CPP_COMMIT"
  git -C "$SOURCE" checkout --detach "$LLAMA_CPP_COMMIT"
  cmake -S "$SOURCE" -B "$BUILD" -DCMAKE_BUILD_TYPE=Release -DLLAMA_CURL=OFF
  cmake --build "$BUILD" --parallel "${TIMMY_BUILD_JOBS:-4}" --target llama-server
  download_verified "$HF_BASE/SmolVLM2-2.2B-Instruct-Q4_K_M.gguf" "$MODEL" "$MODEL_SHA256"
  download_verified "$HF_BASE/mmproj-SmolVLM2-2.2B-Instruct-Q8_0.gguf" "$MMPROJ" "$MMPROJ_SHA256"
  verify_files
  printf 'Pinned model files verified.\n'
}

health_worker() {
  curl --fail --silent --show-error --retry 5 --retry-connrefused --retry-delay 1 --max-time 3 "http://$HOST:$PORT/v1/models" | grep -q 'SmolVLM2-2.2B-Instruct'
  printf 'Self-hosted worker healthy at http://%s:%s/v1\n' "$HOST" "$PORT"
}

start_worker() {
  verify_files
  if [[ ! -x "$SERVER" ]]; then
    printf 'Missing llama-server executable: %s. Run %s install.\n' "$SERVER" "$0" >&2
    return 1
  fi
  mkdir -p "$RUNTIME"
  if [[ -f "$PIDFILE" ]] && kill -0 "$(cat "$PIDFILE")" 2>/dev/null; then
    health_worker
    return
  fi
  rm -f "$PIDFILE"
  nohup "$SERVER" \
    --model "$MODEL" \
    --mmproj "$MMPROJ" \
    --alias SmolVLM2-2.2B-Instruct \
    --host "$HOST" \
    --port "$PORT" \
    --ctx-size 4096 \
    --parallel 1 \
    --threads "${TIMMY_MODEL_THREADS:-4}" \
    --no-webui >"$LOGFILE" 2>&1 &
  printf '%s\n' "$!" > "$PIDFILE"
  for _ in {1..30}; do
    if health_worker >/dev/null 2>&1; then
      printf 'Self-hosted worker started privately at http://%s:%s/v1\n' "$HOST" "$PORT"
      return
    fi
    sleep 0.2
  done
  kill "$(cat "$PIDFILE")" 2>/dev/null || true
  rm -f "$PIDFILE"
  printf 'Worker did not become healthy; inspect %s\n' "$LOGFILE" >&2
  return 1
}

stop_worker() {
  if [[ ! -f "$PIDFILE" ]]; then
    printf 'Self-hosted worker is not running.\n'
    return
  fi
  local pid
  pid="$(cat "$PIDFILE")"
  kill "$pid" 2>/dev/null || true
  for _ in {1..25}; do
    if ! kill -0 "$pid" 2>/dev/null; then
      rm -f "$PIDFILE"
      printf 'Self-hosted worker stopped.\n'
      return
    fi
    sleep 0.1
  done
  kill -9 "$pid" 2>/dev/null || true
  rm -f "$PIDFILE"
  printf 'Self-hosted worker stopped.\n'
}

case "${1:-}" in
  receipt)
    printf 'llama.cpp commit: %s\n' "$LLAMA_CPP_COMMIT"
    printf 'SmolVLM2-2.2B-Instruct-Q4_K_M.gguf sha256: %s\n' "$MODEL_SHA256"
    printf 'mmproj-SmolVLM2-2.2B-Instruct-Q8_0.gguf sha256: %s\n' "$MMPROJ_SHA256"
    printf 'bind address: %s:%s\n' "$HOST" "$PORT"
    ;;
  verify)
    verify_files
    printf 'Pinned model files verified.\n'
    ;;
  install)
    install_worker
    ;;
  start)
    start_worker
    ;;
  health)
    health_worker
    ;;
  stop)
    stop_worker
    ;;
  *)
    printf 'Usage: %s {install|start|health|stop|receipt|verify}\n' "$0" >&2
    exit 64
    ;;
esac