#!/usr/bin/env bash
# Watchdog for the TFLoffroad YouTube miner.
# Invoked every 5 minutes by cron.
#
# Behaviour:
#   - Crawler alive + heartbeat fresh  -> leave it alone.
#   - Crawler dead OR heartbeat older than STALE_SECONDS -> kill stragglers,
#     then start a fresh crawler in nohup so it survives this script exit.
#   - Crawl already complete (every video fetched or permanently failed) -> exit.

set -euo pipefail

BASE_DIR="/root/youtube/TFLoffroad"
SCRIPT="$BASE_DIR/TFLoffroad_youtube_miner.py"
PID_FILE="$BASE_DIR/crawler.pid"
HEARTBEAT_FILE="$BASE_DIR/crawler.heartbeat"
STATE_FILE="$BASE_DIR/crawler_state.json"
VIDEO_CSV="$BASE_DIR/TFLoffroad_youtube.csv"
LOG_FILE="$BASE_DIR/watchdog.log"
STDOUT_LOG="$BASE_DIR/crawler.stdout.log"
STDERR_LOG="$BASE_DIR/crawler.stderr.log"

# 30 minutes without a heartbeat = the crawler is wedged.  The miner pauses
# 3-6 minutes between work windows, which fits well inside this window.
STALE_SECONDS=1800

mkdir -p "$BASE_DIR"

ts() { date -Iseconds; }

pid_value() {
  if [[ -s "$PID_FILE" ]]; then
    tr -dc '0-9' < "$PID_FILE" || true
  fi
}

is_running() {
  local pid; pid="$(pid_value)"
  [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null
}

is_stale() {
  if [[ ! -s "$HEARTBEAT_FILE" ]]; then
    return 0
  fi
  local now heartbeat_age
  now="$(date +%s)"
  heartbeat_age=$((now - $(stat -c %Y "$HEARTBEAT_FILE")))
  (( heartbeat_age > STALE_SECONDS ))
}

start_crawler() {
  echo "$(ts) watchdog: starting crawler" >> "$LOG_FILE"
  cd "$BASE_DIR"
  nohup python3 "$SCRIPT" >> "$STDOUT_LOG" 2>> "$STDERR_LOG" &
  echo "$!" > "$PID_FILE"
  echo "$(ts) watchdog: crawler started pid=$!" >> "$LOG_FILE"
}

kill_crawler() {
  local pid; pid="$(pid_value)"
  if [[ -n "$pid" ]]; then
    echo "$(ts) watchdog: killing pid=$pid" >> "$LOG_FILE"
    kill "$pid" 2>/dev/null || true
    for _ in 1 2 3 4 5; do
      kill -0 "$pid" 2>/dev/null || break
      sleep 2
    done
    if kill -0 "$pid" 2>/dev/null; then
      echo "$(ts) watchdog: forcing pid=$pid" >> "$LOG_FILE"
      kill -9 "$pid" 2>/dev/null || true
    fi
  fi
  rm -f "$PID_FILE" "$BASE_DIR/crawler.lock"
}

# A run is considered complete when every video is either completed or
# permanently marked failed.
is_complete() {
  python3 - "$STATE_FILE" "$VIDEO_CSV" <<'PY'
import csv, json, sys
from pathlib import Path
state_path = Path(sys.argv[1]); video_csv = Path(sys.argv[2])
if not state_path.exists() or not video_csv.exists():
    raise SystemExit(1)
state = json.loads(state_path.read_text(encoding="utf-8"))
with video_csv.open("r", encoding="utf-8", newline="") as fh:
    total = sum(1 for _ in csv.DictReader(fh))
completed = len(set(state.get("completed", [])))
failed = state.get("failed") or {}
raise SystemExit(0 if total > 0 and completed + len(failed) >= total else 1)
PY
}

if is_running; then
  pid="$(pid_value)"
  if is_stale; then
    echo "$(ts) watchdog: crawler pid=$pid stale (>${STALE_SECONDS}s); restarting" >> "$LOG_FILE"
    kill_crawler
    sleep 5
    start_crawler
  else
    echo "$(ts) watchdog: crawler healthy pid=$pid" >> "$LOG_FILE"
  fi
else
  echo "$(ts) watchdog: crawler not running" >> "$LOG_FILE"
  if is_complete; then
    echo "$(ts) watchdog: crawl already complete; leaving alone" >> "$LOG_FILE"
    exit 0
  fi
  start_crawler
fi