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

LLAMA_DIR="~/code/llama.cpp"
CHORUS_DIR="~/code/chorus"
TESTDATA_DIR="$CHORUS_DIR/testdata"
OUT_DIR="$CHORUS_DIR/runs/$(date +%Y%m%d_%H%M%S)"

HOST="127.0.0.1"
PORT="8001"
BASE_URL="http://${HOST}:${PORT}"

mkdir -p "$OUT_DIR"

MODELS=(
  "qwen3-1.7b-q8_0|Qwen/Qwen3-1.7B-GGUF:Q8_0"
  "qwen3-4b-q4_k_m|Qwen/Qwen3-4B-GGUF:Q4_K_M"
  "qwen3-8b-q4_k_m|Qwen/Qwen3-8B-GGUF:Q4_K_M"
)

TEST_FILES=(
  "$TESTDATA_DIR"/*.json
)

wait_for_server() {
  local pid="$1"

  echo "Waiting for llama-server..."
  for i in {1..180}; do
    if ! kill -0 "$pid" 2>/dev/null; then
      echo "ERROR: llama-server exited while starting." >&2
      return 1
    fi

    if curl -fsS "$BASE_URL/v1/models" >/dev/null 2>&1; then
      echo "llama-server is ready."
      return 0
    fi

    sleep 1
  done

  echo "ERROR: llama-server did not become ready." >&2
  return 1
}

stop_server() {
  local pid="${1:-}"

  if [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null; then
    echo "Stopping llama-server pid=$pid"
    kill "$pid" 2>/dev/null || true

    for i in {1..20}; do
      if ! kill -0 "$pid" 2>/dev/null; then
        return 0
      fi
      sleep 1
    done

    echo "Force killing llama-server pid=$pid"
    kill -9 "$pid" 2>/dev/null || true
  fi
}

cd "$CHORUS_DIR"

SUMMARY="$OUT_DIR/summary.tsv"
echo -e "model\tfile\tseconds\texit_code\toutput\tlog" > "$SUMMARY"

for model_entry in "${MODELS[@]}"; do
  model_name="${model_entry%%|*}"
  hf_model="${model_entry#*|}"

  model_out_dir="$OUT_DIR/$model_name"
  mkdir -p "$model_out_dir"

  llama_log="$model_out_dir/llama-server.log"

  echo
  echo "============================================================"
  echo "Starting model: $model_name"
  echo "HF repo/file: $hf_model"
  echo "============================================================"

  cd "$LLAMA_DIR"

  stdbuf -oL -eL ./build/bin/llama-server \
    -hf "$hf_model" \
    --host "$HOST" \
    --port "$PORT" \
    --ctx-size 4096 \
    --alias scorer \
    --threads 12 \
    --threads-batch 24 \
    --batch-size 2048 \
    --ubatch-size 512 \
    --flash-attn on \
    > "$llama_log" 2>&1 &

  LLAMA_PID=$!

  trap 'stop_server "$LLAMA_PID"' EXIT

  if ! wait_for_server "$LLAMA_PID"; then
    echo "llama-server failed. See log: $llama_log" >&2
    stop_server "$LLAMA_PID"
    exit 1
  fi

  cd "$CHORUS_DIR"

  for input_file in "${TEST_FILES[@]}"; do
    base="$(basename "$input_file" .json)"
    out_file="$model_out_dir/out_${base}.json"
    err_file="$model_out_dir/log_${base}.txt"

    echo "Running $model_name on $base.json"

    start_ns="$(date +%s%N)"
    set +e
    python filter_comments.py "$input_file" > "$out_file" 2> "$err_file"
    exit_code=$?
    set -e
    end_ns="$(date +%s%N)"

    seconds="$(awk "BEGIN { printf \"%.2f\", ($end_ns - $start_ns) / 1000000000 }")"

    echo -e "${model_name}\t${base}.json\t${seconds}\t${exit_code}\t${out_file}\t${err_file}" >> "$SUMMARY"

    echo "  finished in ${seconds}s, exit=$exit_code"
  done

  stop_server "$LLAMA_PID"
  trap - EXIT

  sleep 2
done

echo
echo "Done."
echo "Results: $OUT_DIR"
echo "Summary: $SUMMARY"