#!/usr/bin/env bash
# A fake colab CLI for testing scripts/colab/smoke_window.sh without Colab: one "VM" directory under $FAKE_ST, uploads
# and downloads that really copy files, `status` reading BUSY while a job runs, and `exec` that runs the real job file
# (mock_job.py) with its --env settings. FAKE_MODE picks what goes wrong:
#   normal          the exec runs the job and returns
#   neighbourclose  a neighbour session closes between the two rate reads (15.90 -> 10.60), so OWN_RATE falls back to GPU_RATE
#   lostfirst       the first exec loses its connection and hangs with nothing run; the second runs the job
#   late            the first exec's job arrives FAKE_LATE_S seconds late (and survives the exec being killed, as a
#                   remote kernel would); the second runs the job at once
#   neverstarts     every exec hangs with nothing run
# FAKE_GATE (any mode) sets what each VM health gate exec reports: good (default), bad or dead.
ST=${FAKE_ST:?FAKE_ST}; VM=$ST/vm; mkdir -p "$VM/content/pin"
while [ $# -gt 0 ]; do
  case "$1" in
    --config) shift 2;;
    *) break;;
  esac
done
cmd=$1; shift
maybe_banner() {
  local emit=0
  if [ "${FAKE_BANNER:-}" = always ]; then
    emit=1
  elif [ "${FAKE_BANNER:-}" = from_settings ]; then
    local sf="${COLAB_SETTINGS_FILE:-$HOME/.config/colab-cli/settings.json}"
    if ! grep -q '"enable_update_check"[[:space:]]*:[[:space:]]*false' "$sf" 2>/dev/null; then
      emit=1
    fi
  fi
  if [ "$emit" = 1 ]; then
    echo "[colab] A new version of google-colab-cli is available: 0.7.4 -> 0.8.0"
    echo "[colab] Run \`uv tool upgrade google-colab-cli\` to update."
    echo "[colab] To silence this check, set \"enable_update_check\": false in ~/.config/colab-cli/settings.json"
  fi
}
maybe_banner
case "$cmd" in
  usage)
    if [ "${FAKE_MODE:-normal}" = neighbourclose ]; then
      if [ -f "$ST/open" ]; then r=10.60; else r=15.90; fi
    elif [ -f "$ST/open" ]; then r=5.30; else r=0.00; fi
    echo "Current balance: 100.00 compute units"; echo "Usage rate: $r/hr";;
  new)
    if [ -n "${PIN_OPEN_LOCK:-}" ] && ! flock -n "$PIN_OPEN_LOCK" true 2>/dev/null; then
      echo locked >> "$ST/new-lock.log"
    else
      echo unlocked >> "$ST/new-lock.log"
    fi
    if [ "${FAKE_NEW_FAIL:-0}" = 1 ]; then
      echo "[colab] Requesting runtime..."
      exit 1
    fi
    rm -f "$ST/sessions_gone" "$ST/unassigned"
    echo "$*" >> "$ST/new-args.log"; touch "$ST/open"; echo $(( $(cat "$ST/news" 2>/dev/null || echo 0) + 1 )) > "$ST/news"; echo "[colab] Session READY.";;
  sessions)
    if [ -n "${FAKE_DROP_SESSION_AT_CALLS:-}" ] && [ -f "$ST/open" ]; then
      vm_calls=$(wc -l "$VM"/content/pin/toolcount/*/calls.jsonl 2>/dev/null | awk 'END{print $1+0}')
      if [ "${vm_calls:-0}" -ge "$FAKE_DROP_SESSION_AT_CALLS" ]; then
        touch "$ST/sessions_gone"
      fi
    fi
    if [ -f "$ST/open" ] && [ ! -f "$ST/sessions_gone" ]; then
      n_vm=$(cat "$ST/news" 2>/dev/null || echo 1)
      vm_id=$(echo "${FAKE_VM:-}" | cut -d, -f"$n_vm")
      [ -n "$vm_id" ] || vm_id="gpu-a100-s-kkb-usc1b2-fake${n_vm}"
      hw="${PIN_GPU:-A100}"
      shape="Standard"; [ "${PIN_HIGH_MEM:-0}" = 1 ] && shape="High-RAM"
      variant="GPU"; [ "$hw" = CPU ] && variant="CPU"
      echo "[${PIN_SESSION:-fake}] $vm_id | Hardware: $hw | Shape: $shape | Variant: $variant"
    else
      if [ -f "$ST/unassigned" ] || [ -f "$ST/sessions_gone" ]; then
        echo "[colab] Pruned 1 stale local session(s)."
      fi
      echo "[colab] No active sessions found on server."
    fi
    true;;
  stop)
    if [ -n "${PIN_OPEN_LOCK:-}" ] && ! flock -n "$PIN_OPEN_LOCK" true 2>/dev/null; then
      touch "$ST/stopped_under_lock"
      echo locked >> "$ST/stop-lock.log"
    else
      echo unlocked >> "$ST/stop-lock.log"
    fi
    rm -f "$ST/open" "$ST/sessions_gone" "$ST/unassigned"; echo "[colab] Session terminated.";;
  status)
    [ -n "${FAKE_LEAK_URL:-}" ] && echo "$FAKE_LEAK_URL"
    if [ -n "${FAKE_BANNER:-}" ]; then
      echo "Session: [${PIN_SESSION:-fake}] gpu-a100-s-kkb-usc1b2-1a2b3c4d5e | Hardware: A100 | Runtime: Python 3 (GPU) | Kernel: python3"
    fi
    if ls "$ST"/busy.* >/dev/null 2>&1; then echo "[fake] Status: BUSY (exec(job))"; else echo "[fake] Status: IDLE"; fi;;
  upload)
    [ "$1" = -s ] && shift 2
    u_try=$(( $(cat "$ST/upload-tries" 2>/dev/null || echo 0) + 1 )); echo "$u_try" > "$ST/upload-tries"
    if [ "${FAKE_UPLOAD_MODE:-}" = timeout_silent ]; then
      echo "Uploading chunk 1/1 ($(basename "$1"))..."
      exit 124
    fi
    if [ "${FAKE_UPLOAD_FAIL_ONCE:-0}" = 1 ] && [ "$u_try" = 1 ]; then
      echo "${FAKE_LEAK_URL:-upload failed}"
      exit 1
    fi
    mkdir -p "$VM$(dirname "$2")"; cp "$1" "$VM$2"
    echo "vm$(cat "$ST/news" 2>/dev/null || echo 0) $2" >> "$ST/uploads.log"; echo "[colab] Uploaded";;
  download)
    sess_name="${PIN_SESSION:-fake}"
    if [ "$1" = -s ]; then sess_name="$2"; shift 2; fi
    d_try=$(( $(cat "$ST/download-tries" 2>/dev/null || echo 0) + 1 )); echo "$d_try" > "$ST/download-tries"
    if [ "${FAKE_DOWNLOAD_LEAK_ONCE:-0}" = 1 ] && [ "$d_try" = 1 ]; then
      echo "[colab] Download failed: ${FAKE_LEAK_URL:-download failed}"
      exit 1
    fi
    if [ ! -f "$ST/open" ]; then
      echo "[colab] Session '$sess_name' not found."
      exit 1
    fi
    case "$1" in
      /content/pin/job.log)
        if [ -s "$VM$1" ] && [ -n "${FAKE_HB_LINE:-}" ] && ! grep -Fq "$FAKE_HB_LINE" "$VM$1" 2>/dev/null; then
          printf '%s\n' "$FAKE_HB_LINE" >> "$VM$1"
        fi
        ;;
      */toolcount/*/calls.jsonl)
        if [ ! -s "$VM$1" ] && { [ -n "${FAKE_PARTIAL_LAST_LINE_ON:-}" ] || [ -n "${FAKE_UNASSIGN_MID_DOWNLOAD_ON:-}" ] || [ -n "${FAKE_CORRUPT_COPYBACK_ON:-}" ] || [ -n "${FAKE_TRUNCATE_COPYBACK_ON:-}" ] || [ -n "${FAKE_DROP_AFTER_COPYBACKS:-}" ]; }; then
          for _ in $(seq 1 30); do
            [ -s "$VM$1" ] && break
            sleep 0.2
          done
        fi
        if [ -s "$VM$1" ]; then
          cb_try=$(( $(cat "$ST/copyback-calls-tries" 2>/dev/null || echo 0) + 1 ))
          echo "$cb_try" > "$ST/copyback-calls-tries"
          if [ -n "${FAKE_UNASSIGN_MID_DOWNLOAD_ON:-}" ] && [ "$cb_try" = "$FAKE_UNASSIGN_MID_DOWNLOAD_ON" ]; then
            head -c 32 "$VM$1" > "$2"
            if [ "${FAKE_UNASSIGN_DROP_SESSION:-0}" = 1 ]; then
              rm -f "$ST/open"; touch "$ST/unassigned"
            fi
            echo "[colab] Download failed: 502 Server Error for url: ${FAKE_COPYBACK_LEAK_URL:-https://8080-m-s-kkb-usc1b2.prod.colab.dev/api/contents$1}"
            exit 1
          fi
          if [ -n "${FAKE_PARTIAL_LAST_LINE_ON:-}" ] && [ "$cb_try" = "$FAKE_PARTIAL_LAST_LINE_ON" ]; then
            prev_l=$(cat "$ST/copyback-calls-lines" 2>/dev/null || echo 0)
            for _ in $(seq 1 30); do
              [ "$(wc -l < "$VM$1" 2>/dev/null || echo 0)" -gt "$prev_l" ] && break
              sleep 0.2
            done
            cp "$VM$1" "$2"
            wc -l < "$VM$1" > "$ST/copyback-calls-lines"
            printf '{"call_id": "partial_mid_append' >> "$2"
            echo "[colab] Downloaded '$1' to '$2'${FAKE_COPYBACK_LEAK_URL:+ $FAKE_COPYBACK_LEAK_URL}"
            exit 0
          fi
          if [ -n "${FAKE_CORRUPT_COPYBACK_ON:-}" ] && [ "$cb_try" = "$FAKE_CORRUPT_COPYBACK_ON" ]; then
            printf '{"call_id": ""}\n' > "$2"
            cat "$VM$1" >> "$2"
            echo "[colab] Downloaded '$1' to '$2'${FAKE_COPYBACK_LEAK_URL:+ $FAKE_COPYBACK_LEAK_URL}"
            exit 0
          fi
          if [ -n "${FAKE_TRUNCATE_COPYBACK_ON:-}" ] && [ "$cb_try" = "$FAKE_TRUNCATE_COPYBACK_ON" ]; then
            printf '{"call_id": "truncated_pull"}\n' > "$2"
            echo "[colab] Downloaded '$1' to '$2'${FAKE_COPYBACK_LEAK_URL:+ $FAKE_COPYBACK_LEAK_URL}"
            exit 0
          fi
          wc -l < "$VM$1" > "$ST/copyback-calls-lines"
        fi
        ;;
      */toolcount/*/generations.jsonl)
        if [ ! -s "$VM$1" ] && { [ -n "${FAKE_PARTIAL_LAST_LINE_ON:-}" ] || [ -n "${FAKE_UNASSIGN_MID_DOWNLOAD_ON:-}" ] || [ -n "${FAKE_CORRUPT_COPYBACK_ON:-}" ] || [ -n "${FAKE_TRUNCATE_COPYBACK_ON:-}" ] || [ -n "${FAKE_DROP_AFTER_COPYBACKS:-}" ]; }; then
          for _ in $(seq 1 30); do
            [ -s "$VM$1" ] && break
            sleep 0.2
          done
        fi
        if [ -s "$VM$1" ]; then
          cb_g_try=$(( $(cat "$ST/copyback-gen-tries" 2>/dev/null || echo 0) + 1 ))
          echo "$cb_g_try" > "$ST/copyback-gen-tries"
          if [ -n "${FAKE_DROP_AFTER_COPYBACKS:-}" ] && [ "$cb_g_try" -ge "$FAKE_DROP_AFTER_COPYBACKS" ]; then
            cp "$VM$1" "$2"
            rm -f "$ST/open"; touch "$ST/unassigned"
            echo "[colab] Downloaded '$1' to '$2'${FAKE_COPYBACK_LEAK_URL:+ $FAKE_COPYBACK_LEAK_URL}"
            exit 0
          fi
        fi
        ;;
    esac
    if [ -f "$VM$1" ]; then
      cp "$VM$1" "$2"
      echo "[colab] Downloaded '$1' to '$2'${FAKE_COPYBACK_LEAK_URL:+ $FAKE_COPYBACK_LEAK_URL}"
    else
      echo "[colab] Download failed: $1"
      exit 1
    fi;;
  exec)
    file=""; envs=()
    while [ $# -gt 0 ]; do case "$1" in -f) file=$2; shift 2;; --env) envs+=("$2"); shift 2;; -s|--timeout) shift 2;; *) shift;; esac; done
    # The VM health gate's exec (vm_gate.py): FAKE_GATE is a comma list, one entry per gate exec: good, bad, or dead (the
    # exec dies with nothing printed, as windows 11c and 12e did). Gate execs don't count as job execs.
    if [ "$(basename "$file")" = vm_gate.py ]; then
      if [ -n "${PIN_OPEN_LOCK:-}" ] && ! flock -n "$PIN_OPEN_LOCK" true 2>/dev/null; then
        echo locked >> "$ST/gate-lock.log"
      else
        echo unlocked >> "$ST/gate-lock.log"
      fi
      g=$(( $(cat "$ST/gates" 2>/dev/null || echo 0) + 1 )); echo $g > "$ST/gates"
      printf '%s\n' "${envs[@]}" > "$ST/gate-env-$g"
      what=$(echo "${FAKE_GATE:-good}" | cut -d, -f"$g"); what=${what:-good}
      case "$what" in
        good) echo 'VM_GATE {"ok": true, "why": [], "gpu_name": "NVIDIA A100-SXM4-40GB", "disk_free_gb": 204.2}';;
        bad) echo 'VM_GATE {"ok": false, "why": ["card in use: 1200 MiB used, 1 processes"], "disk_free_gb": 204.2}';;
        dead) echo "RuntimeError: Connection was lost."; exec sleep 3600;;  # exec: the timeout kills the sleep itself
      esac
      exit 0
    fi
    [ -n "${FAKE_LEAK_URL:-}" ] && echo "$FAKE_LEAK_URL"
    n=$(( $(cat "$ST/execs" 2>/dev/null || echo 0) + 1 )); echo $n > "$ST/execs"
    job="$ST/job-$n.sh"
    { echo "#!/usr/bin/env bash"; echo "touch '$ST/busy.$n'"
      if [ "${FAKE_EXEC_IGNORE_TERM:-0}" = 1 ]; then
        echo "trap 'rm -f \"$ST/busy.$n\"' EXIT"
        echo "trap '' TERM INT"
      else
        echo "trap 'rm -f \"$ST/busy.$n\"; [ -n \"\${JPID:-}\" ] && kill \"\$JPID\" 2>/dev/null || true' TERM INT EXIT"
      fi
      for e in "${envs[@]}"; do printf 'export %q\n' "$e"; done
      echo "export PIN_CONTENT='$VM/content'"; echo "export PIN_WORK='$VM/content/pin'"; echo "python3 '$file' & JPID=\$!; wait \$JPID"
      if [ -n "${FAKE_EXEC_HANG_AFTER_DONE_S:-}" ]; then
        echo "sleep '${FAKE_EXEC_HANG_AFTER_DONE_S}' & HPID=\$!; wait \$HPID"
      fi
    } > "$job"
    chmod +x "$job"
    lost() { echo "RuntimeError: Connection was lost."; exec sleep 3600; }
    case "$FAKE_MODE:$n" in
      normal:*|neighbourclose:*) exec "$job";;
      lostfirst:1|neverstarts:*) lost;;
      lostfirst:*) exec "$job";;
      late:1) setsid bash -c "sleep ${FAKE_LATE_S:-4}; '$job'" > /dev/null 2>&1 < /dev/null & lost;;
      late:*) exec "$job";;
      # latefirstwins: the first exec's job arrives just after the watcher's IDLE read and takes the lock; the retry's
      # job arrives later, waits for it, and stops on this window's result.
      latefirstwins:1) setsid bash -c "sleep ${FAKE_LATE_S:-1.5}; '$job'" > /dev/null 2>&1 < /dev/null & lost;;
      latefirstwins:*) sleep "${FAKE_SECOND_DELAY_S:-2}"; exec "$job";;
    esac;;
esac
