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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
CODE_NEEDLE_REPO_URL="https://github.com/alexziskind1/codeneedle.git"

usage() {
  cat <<'EOF'
Usage: ./scripts/benchmark-codeneedle-suite.sh [options]

Vergleicht mehrere lokale Ollama-Modelle mit CodeNeedle gegen ein OpenAI-
kompatibles Endpoint-Schema. Das Skript klont CodeNeedle zur Laufzeit unter
.cache/benchmarks, baut das Upstream-Dockerfile mit Podman oder Docker und
fuehrt die Standard-Corpora http_server und jquery gegen den lokalen Ollama-
Server aus.

Optionen:
  --model <name>         Modellname in Ollama; mehrfach nutzbar
                         (default: qwen3-coding:latest, qwen25-coding:latest,
                         deepseek-coding:latest)
  --corpus <name>        CodeNeedle-Corpus; mehrfach nutzbar
                         (default: http_server, jquery)
  --base-url <url>       OpenAI-kompatible Basis-URL ohne /v1-Suffix
                         (default: http://127.0.0.1:11434)
  --api-key <value>      API-Key fuer den Endpoint (default: not-needed)
  --temperature <n>      Sampling-Temperatur (default: 0)
  --max-tokens <n>       Antwortbudget je Query (default: 8000)
  --timeout-sec <n>      Timeout je Query in Sekunden (default: 600)
  --sample-k <n>         Anzahl der getesteten Funktionen je Corpus (default: 16)
  --sample-seed <n>      Seed fuer die Stichprobe (default: 42)
  --skip-preflight       Ueberspringt CodeNeedle-Preflight-Probe pro Lauf
  --runtime-root <path>  Zielordner fuer temporaere Laufdaten (default: auto)
  --results-root <path>  Zielordner fuer kuratierte Resultate (default: auto)
  --artifact-root <path> Legacy-Alias; nutzt denselben Pfad fuer Runtime und Resultate
  --repo-dir <path>      Ablageort fuer den Upstream-Clone (default: auto)
  --image-tag <tag>      Name des lokal gebauten Runner-Images (default: ghostnet-codeneedle:latest)
  --container-engine <e> podman oder docker (default: aus .env/.env.example oder auto)
  --refresh-upstream     Upstream-Clone vor dem Lauf mit git pull aktualisieren
  --rebuild-image        Image vor dem Lauf neu bauen
  --keep-runtime         Runtime-Dateien nach dem Export nicht loeschen
  -h, --help             Hilfe anzeigen
EOF
}

need_cmd() {
  if ! command -v "$1" >/dev/null 2>&1; then
    echo "ERROR: Missing required command '$1'." >&2
    exit 1
  fi
}

read_env_var() {
  local key="$1"
  local file line

  for file in "$PROJECT_DIR/.env" "$PROJECT_DIR/.env.example"; do
    if [ ! -f "$file" ]; then
      continue
    fi

    line="$(grep -E "^${key}=" "$file" | tail -n 1 || true)"
    if [ -n "$line" ]; then
      printf '%s\n' "${line#*=}"
      return 0
    fi
  done

  return 1
}

slugify() {
  printf '%s' "$1" | tr '/:@ ' '----' | tr -cd '[:alnum:]_.-'
}

MODELS=(
  "qwen3-coding:latest"
  "qwen25-coding:latest"
  "deepseek-coding:latest"
)
CORPORA=(
  "http_server"
  "jquery"
)
BASE_URL="http://127.0.0.1:11434"
API_KEY="not-needed"
TEMPERATURE="0"
MAX_TOKENS=8000
TIMEOUT_SEC=600
SAMPLE_K=16
SAMPLE_SEED=42
SKIP_PREFLIGHT=0
RUNTIME_ROOT=""
RESULTS_ROOT=""
LEGACY_ARTIFACT_ROOT=""
REPO_DIR=""
IMAGE_TAG="ghostnet-codeneedle:latest"
CONTAINER_ENGINE="${CONTAINER_ENGINE:-}"
REFRESH_UPSTREAM=0
REBUILD_IMAGE=0
KEEP_RUNTIME=0

while [ $# -gt 0 ]; do
  case "$1" in
    --model)
      if [ "${MODELS[0]:-}" = "qwen3-coding:latest" ] && [ "${#MODELS[@]}" -eq 3 ]; then
        MODELS=()
      fi
      MODELS+=("$2")
      shift 2
      ;;
    --corpus)
      if [ "${CORPORA[0]:-}" = "http_server" ] && [ "${#CORPORA[@]}" -eq 2 ]; then
        CORPORA=()
      fi
      CORPORA+=("$2")
      shift 2
      ;;
    --base-url)
      BASE_URL="$2"
      shift 2
      ;;
    --api-key)
      API_KEY="$2"
      shift 2
      ;;
    --temperature)
      TEMPERATURE="$2"
      shift 2
      ;;
    --max-tokens)
      MAX_TOKENS="$2"
      shift 2
      ;;
    --timeout-sec)
      TIMEOUT_SEC="$2"
      shift 2
      ;;
    --sample-k)
      SAMPLE_K="$2"
      shift 2
      ;;
    --sample-seed)
      SAMPLE_SEED="$2"
      shift 2
      ;;
    --skip-preflight)
      SKIP_PREFLIGHT=1
      shift
      ;;
    --runtime-root)
      RUNTIME_ROOT="$2"
      shift 2
      ;;
    --results-root)
      RESULTS_ROOT="$2"
      shift 2
      ;;
    --artifact-root)
      LEGACY_ARTIFACT_ROOT="$2"
      shift 2
      ;;
    --repo-dir)
      REPO_DIR="$2"
      shift 2
      ;;
    --image-tag)
      IMAGE_TAG="$2"
      shift 2
      ;;
    --container-engine)
      CONTAINER_ENGINE="$2"
      shift 2
      ;;
    --refresh-upstream)
      REFRESH_UPSTREAM=1
      shift
      ;;
    --rebuild-image)
      REBUILD_IMAGE=1
      shift
      ;;
    --keep-runtime)
      KEEP_RUNTIME=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "ERROR: Unknown option '$1'." >&2
      usage >&2
      exit 1
      ;;
  esac
done

need_cmd git
need_cmd jq

if [ -z "$CONTAINER_ENGINE" ]; then
  CONTAINER_ENGINE="$(read_env_var CONTAINER_ENGINE || true)"
fi

if [ -z "$CONTAINER_ENGINE" ]; then
  if command -v podman >/dev/null 2>&1; then
    CONTAINER_ENGINE="podman"
  elif command -v docker >/dev/null 2>&1; then
    CONTAINER_ENGINE="docker"
  else
    echo "ERROR: Neither podman nor docker is available." >&2
    exit 1
  fi
fi

case "$CONTAINER_ENGINE" in
  podman)
    ENGINE_CMD=(podman)
    MOUNT_SUFFIX=':Z'
    ;;
  docker)
    ENGINE_CMD=(docker)
    MOUNT_SUFFIX=''
    ;;
  *)
    echo "ERROR: Invalid container engine '$CONTAINER_ENGINE'. Use 'podman' or 'docker'." >&2
    exit 1
    ;;
esac

need_cmd "${ENGINE_CMD[0]}"

RUN_ID="codeneedle-suite-$(date +%Y%m%d-%H%M%S)"

if [ -n "$LEGACY_ARTIFACT_ROOT" ]; then
  RUNTIME_ROOT="$LEGACY_ARTIFACT_ROOT"
  if [ -z "$RESULTS_ROOT" ]; then
    RESULTS_ROOT="$LEGACY_ARTIFACT_ROOT"
    KEEP_RUNTIME=1
  fi
fi

if [ -z "$RUNTIME_ROOT" ]; then
  RUNTIME_ROOT="$PROJECT_DIR/.cache/benchmarks/$RUN_ID"
fi

if [ -z "$RESULTS_ROOT" ]; then
  RESULTS_ROOT="$PROJECT_DIR/docs/benchmarks/codeneedle/$RUN_ID"
fi

if [ -z "$REPO_DIR" ]; then
  REPO_DIR="$PROJECT_DIR/.cache/benchmarks/codeneedle/upstream"
fi

mkdir -p "$RUNTIME_ROOT" "$RESULTS_ROOT"

if [ ! -d "$REPO_DIR/.git" ]; then
  mkdir -p "$(dirname "$REPO_DIR")"
  git clone --depth 1 "$CODE_NEEDLE_REPO_URL" "$REPO_DIR"
elif [ "$REFRESH_UPSTREAM" -eq 1 ]; then
  git -C "$REPO_DIR" pull --ff-only
fi

UPSTREAM_REF="$(git -C "$REPO_DIR" rev-parse HEAD)"

if [ "$REBUILD_IMAGE" -eq 1 ] || ! "${ENGINE_CMD[@]}" image inspect "$IMAGE_TAG" >/dev/null 2>&1; then
  "${ENGINE_CMD[@]}" build -t "$IMAGE_TAG" -f "$REPO_DIR/Dockerfile" "$REPO_DIR"
fi

SUMMARY_TABLE="$RESULTS_ROOT/summary.tsv"
RUN_METADATA_FILE="$RESULTS_ROOT/run-metadata.txt"
printf 'corpus\tmodel\trun_rc\tpassed\ttotal\terrored\tprimary_matched\tprimary_total\thallucinated\tbonus\n' > "$SUMMARY_TABLE"

cat > "$RUN_METADATA_FILE" <<EOF
run_id=$RUN_ID
container_engine=$CONTAINER_ENGINE
image_tag=$IMAGE_TAG
upstream_repo=$CODE_NEEDLE_REPO_URL
upstream_ref=$UPSTREAM_REF
base_url=$BASE_URL
temperature=$TEMPERATURE
max_tokens=$MAX_TOKENS
timeout_sec=$TIMEOUT_SEC
sample_k=$SAMPLE_K
sample_seed=$SAMPLE_SEED
skip_preflight=$SKIP_PREFLIGHT
repo_dir=$REPO_DIR
runtime_root=$RUNTIME_ROOT
results_root=$RESULTS_ROOT
EOF

echo "CodeNeedle-Repo: $REPO_DIR @ $UPSTREAM_REF"
echo "Container-Engine: $CONTAINER_ENGINE"
echo "Image: $IMAGE_TAG"
echo "Base URL: $BASE_URL"
echo "Runtime-Wurzel: $RUNTIME_ROOT"
echo "Kuratierte Resultate: $RESULTS_ROOT"
echo

for corpus in "${CORPORA[@]}"; do
  for model in "${MODELS[@]}"; do
    corpus_slug="$(slugify "$corpus")"
    model_slug="$(slugify "$model")"
    runtime_dir="$RUNTIME_ROOT/$corpus_slug/$model_slug"
    results_dir="$RESULTS_ROOT/$corpus_slug/$model_slug"
    dump_file="$runtime_dir/results.json"
    stdout_file="$runtime_dir/stdout.log"
    stderr_file="$runtime_dir/stderr.log"
    run_cmd=(
      "${ENGINE_CMD[@]}" run --rm --network host
      -v "$REPO_DIR:/app$MOUNT_SUFFIX"
      -v "$runtime_dir:/runtime$MOUNT_SUFFIX"
      -w /app
      "$IMAGE_TAG"
      python bench.py run
      --corpus "$corpus"
      --model "$model"
      --base-url "$BASE_URL"
      --api-key "$API_KEY"
      --temperature "$TEMPERATURE"
      --max-tokens "$MAX_TOKENS"
      --timeout "$TIMEOUT_SEC"
      -k "$SAMPLE_K"
      --seed "$SAMPLE_SEED"
      --dump /runtime/results.json
    )

    if [ "$SKIP_PREFLIGHT" -eq 1 ]; then
      run_cmd+=(--skip-preflight)
    fi

    mkdir -p "$runtime_dir" "$results_dir"

    echo "=== CodeNeedle fuer $model auf $corpus ==="
    if "${run_cmd[@]}" >"$stdout_file" 2>"$stderr_file"; then
      run_rc=0
    else
      run_rc=$?
    fi

    cp "$stdout_file" "$results_dir/stdout.log"
    cp "$stderr_file" "$results_dir/stderr.log"

    passed=""
    total=""
    errored=""
    primary_matched=""
    primary_total=""
    hallucinated=""
    bonus=""

    if [ -f "$dump_file" ]; then
      cp "$dump_file" "$results_dir/results.json"
      passed="$(jq '[.results[] | select(.passed == true and (.error == null or .error == ""))] | length' "$dump_file")"
      total="$(jq '.results | length' "$dump_file")"
      errored="$(jq '[.results[] | select(.error != null and .error != "")] | length' "$dump_file")"
      primary_matched="$(jq '[.results[].primary_matched] | add // 0' "$dump_file")"
      primary_total="$(jq '[.results[].primary_total] | add // 0' "$dump_file")"
      hallucinated="$(jq '[.results[].hallucinated] | add // 0' "$dump_file")"
      bonus="$(jq '[.results[].bonus_matched] | add // 0' "$dump_file")"
    fi

    cat > "$results_dir/summary.txt" <<EOF
corpus=$corpus
model=$model
run_rc=$run_rc
passed=$passed
total=$total
errored=$errored
primary_matched=$primary_matched
primary_total=$primary_total
hallucinated=$hallucinated
bonus=$bonus
EOF

    printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n' \
      "$corpus" "$model" "$run_rc" "$passed" "$total" "$errored" \
      "$primary_matched" "$primary_total" "$hallucinated" "$bonus" >> "$SUMMARY_TABLE"

    if [ "$KEEP_RUNTIME" -eq 0 ] && [ "$runtime_dir" != "$results_dir" ]; then
      rm -rf "$runtime_dir"
    fi
  done
done

echo
echo "CodeNeedle-Vergleich gespeichert unter: $RESULTS_ROOT"
if command -v column >/dev/null 2>&1; then
  column -ts $'\t' "$SUMMARY_TABLE"
else
  cat "$SUMMARY_TABLE"
fi