#!/usr/bin/env bash
# Retrieval evaluation for an embedding model, on German documents of the kind
# this homelab actually files.
#
# Companion to scripts/ai-eval.sh, but it cannot work the same way. A chat
# answer can be read and judged; an embedding is 1024 opaque numbers. What is
# measurable is whether the right document comes back first, so this builds a
# small labelled corpus, embeds it, and scores real questions against it.
#
# Metrics:
#   Recall@1  the correct document ranked first — what Paperless's RAG actually
#             feeds to the classifier
#   Recall@3  it at least made the shortlist
#   MRR       mean reciprocal rank, so near-misses are visible rather than
#             collapsing to pass/fail
#   margin    cosine gap between the correct document and the best wrong one.
#             A model that is right by 0.002 is guessing; one that is right by
#             0.15 has actually separated them.
#
# The corpus deliberately contains confusable pairs — two Allianz policies, four
# documents that are all "Rechnungen", a car insurance and a vehicle
# registration sharing a number plate. An easy corpus would rate everything 100%
# and tell you nothing.
#
# Usage:
#   scripts/embed-eval.sh                                    # cluster default
#   scripts/embed-eval.sh -m some-model
#   scripts/embed-eval.sh -e https://host/v1/embeddings -m model -k KEY
#
# Env equivalents: EMBED_EVAL_ENDPOINT, EMBED_EVAL_MODEL, EMBED_EVAL_KEY.
#
# To compare two models on our llama-embed, -m is not enough: swap --hf-repo in
# kubernetes/apps/llama-server/embed.yaml, let it roll out, then re-run. Note
# that changing the model also invalidates any vectors Paperless has stored, so
# comparisons are best made before an index exists.
#
# Exit code is the number of queries whose correct document did not rank first.
set -uo pipefail

ENDPOINT="${EMBED_EVAL_ENDPOINT:-http://llama-embed.ai.svc.cluster.local:8080/v1/embeddings}"
MODEL="${EMBED_EVAL_MODEL:-qwen3-embedding-0.6b}"
APIKEY="${EMBED_EVAL_KEY:-}"
# Run curl from inside this pod when the endpoint is a cluster-internal name.
# llama-embed has no ingress by design, so there is nothing to reach from a
# workstation.
VIA_POD="${EMBED_EVAL_VIA_POD:-}"

while [ $# -gt 0 ]; do
    case "$1" in
        -e|--endpoint) ENDPOINT="$2"; shift 2 ;;
        -m|--model)    MODEL="$2";    shift 2 ;;
        -k|--key)      APIKEY="$2";   shift 2 ;;
        -p|--via-pod)  VIA_POD="$2";  shift 2 ;;
        -h|--help)     sed -n '2,34p' "$0"; exit 0 ;;
        *) echo "unknown argument: $1 (try --help)" >&2; exit 2 ;;
    esac
done

HERE="$(cd "$(dirname "$0")" && pwd)"
CASES="$HERE/embed-eval"
WORK="$(mktemp -d)"
trap 'rm -rf "$WORK"' EXIT

for f in "$CASES/corpus.jsonl" "$CASES/queries.jsonl"; do
    [ -r "$f" ] || { echo "missing fixture: $f" >&2; exit 2; }
done

# Default to reaching the Service from a pod that is allowed to, unless the
# caller pointed us somewhere externally reachable.
if [ -z "$VIA_POD" ] && [ "${ENDPOINT#*svc.cluster.local}" != "$ENDPOINT" ]; then
    VIA_POD="deploy/llama-server -n ai"
fi

# embed <text> -> the vector as a JSON array on stdout
embed() {
    local payload out
    payload=$(jq -nc --arg m "$MODEL" --arg t "$1" '{model:$m,input:$t}')
    if [ -n "$VIA_POD" ]; then
        # stderr is kept and surfaced on failure rather than discarded: a
        # silently swallowed kubectl or curl error is indistinguishable from a
        # model that returned nothing.
        # shellcheck disable=SC2086
        out=$(kubectl exec $VIA_POD -- curl -sS -m 120 \
              -H 'Content-Type: application/json' \
              ${APIKEY:+-H "Authorization: Bearer $APIKEY"} \
              -d "$payload" "$ENDPOINT" 2>"$WORK/err.txt")
    else
        out=$(curl -sS -m 120 -H 'Content-Type: application/json' \
              ${APIKEY:+-H "Authorization: Bearer $APIKEY"} \
              -d "$payload" "$ENDPOINT")
    fi
    printf '%s' "$out" | jq -c '.data[0].embedding // empty'
}

echo "endpoint: $ENDPOINT"
echo "model:    $MODEL"
[ -n "$VIA_POD" ] && echo "via pod:  $VIA_POD"
echo ""

# ---- embed the corpus ------------------------------------------------------
DOCS=$(wc -l < "$CASES/corpus.jsonl" | tr -d ' ')
echo "embedding $DOCS documents..."
: > "$WORK/vecs.jsonl"
start=$(date +%s)
while IFS= read -r line; do
    id=$(printf '%s' "$line" | jq -r '.id')
    text=$(printf '%s' "$line" | jq -r '.text')
    vec=$(embed "$text")
    if [ -z "$vec" ]; then
        echo "FAILED to embed '$id' — endpoint returned no vector." >&2
        [ -s "$WORK/err.txt" ] && { echo "stderr:" >&2; cat "$WORK/err.txt" >&2; }
        exit 2
    fi
    jq -nc --arg id "$id" --argjson v "$vec" '{id:$id, vec:$v}' >> "$WORK/vecs.jsonl"
done < "$CASES/corpus.jsonl"
elapsed=$(( $(date +%s) - start ))
DIM=$(head -1 "$WORK/vecs.jsonl" | jq '.vec | length')
echo "  $DOCS documents, ${DIM} dimensions, ${elapsed}s"
echo ""

# ---- score the queries -----------------------------------------------------
# Cosine similarity in jq. The vectors llama.cpp returns are normalised, but the
# full form costs nothing and keeps this correct against servers that do not.
rank() { # queryvec -> "id score" lines, best first
    jq -s -r --argjson q "$1" '
        def dot($a; $b): reduce range(0; $a|length) as $i (0; . + ($a[$i] * $b[$i]));
        def norm($a): (reduce ($a[]) as $x (0; . + ($x * $x))) | sqrt;
        map({id: .id, s: (dot(.vec; $q) / ((norm(.vec) * norm($q)) + 1e-12))})
        | sort_by(-.s)[]
        | "\(.id) \(.s)"
    ' "$WORK/vecs.jsonl"
}

MISS=0; TOTAL=0; HIT1=0; HIT3=0; MRRSUM=0
printf '%-46s %-22s %5s %8s\n' QUERY EXPECTED RANK MARGIN
printf '%s\n' "------------------------------------------------------------------------------------"
while IFS= read -r line; do
    q=$(printf '%s' "$line" | jq -r '.q')
    want=$(printf '%s' "$line" | jq -r '.want')
    qvec=$(embed "$q")
    if [ -z "$qvec" ]; then
        echo "FAILED to embed query: $q" >&2
        [ -s "$WORK/err.txt" ] && { echo "stderr:" >&2; cat "$WORK/err.txt" >&2; }
        exit 2
    fi
    rank "$qvec" > "$WORK/ranked.txt"

    pos=$(grep -n "^$want " "$WORK/ranked.txt" | cut -d: -f1)
    pos=${pos:-999}
    TOTAL=$((TOTAL + 1))
    [ "$pos" -eq 1 ] && HIT1=$((HIT1 + 1)) || MISS=$((MISS + 1))
    [ "$pos" -le 3 ] && HIT3=$((HIT3 + 1))
    MRRSUM=$(awk -v s="$MRRSUM" -v p="$pos" 'BEGIN{printf "%.6f", s + 1/p}')

    # Margin: the correct document's score minus the best incorrect one. On a
    # miss this is negative, which says how badly it lost.
    want_s=$(awk -v w="$want" '$1==w{print $2; exit}' "$WORK/ranked.txt")
    best_other=$(awk -v w="$want" '$1!=w{print $2; exit}' "$WORK/ranked.txt")
    margin=$(awk -v a="${want_s:-0}" -v b="${best_other:-0}" 'BEGIN{printf "%+.4f", a-b}')

    if [ "$pos" -eq 1 ]; then mark=" "; else mark="!"; fi
    printf '%s%-45.45s %-22s %5s %8s\n' "$mark" "$q" "$want" "$pos" "$margin"
    if [ "$pos" -ne 1 ]; then
        printf '   got instead: %s\n' "$(head -1 "$WORK/ranked.txt" | cut -d' ' -f1)"
    fi
done < "$CASES/queries.jsonl"

R1=$(awk -v h="$HIT1" -v t="$TOTAL" 'BEGIN{printf "%.0f", 100*h/t}')
R3=$(awk -v h="$HIT3" -v t="$TOTAL" 'BEGIN{printf "%.0f", 100*h/t}')
MRR=$(awk -v s="$MRRSUM" -v t="$TOTAL" 'BEGIN{printf "%.3f", s/t}')

echo ""
echo "========================================"
printf '  RESULT model=%s dim=%s recall@1=%s%% recall@3=%s%% mrr=%s\n' \
    "$MODEL" "$DIM" "$R1" "$R3" "$MRR"
echo "========================================"
exit "$MISS"