#!/usr/bin/env bash
# Design-partner runbook, step 3: judge the trained model against the protocol.
# Run after the partner dataset was built (strictly after the protocol was
# declared) and a model was trained on it. Set TESSERA_MODEL_ID.
#
# The Idempotency-Key is derived from the model and the protocol: repeating the
# same judgement is a safe retry, and a different protocol is a new request.
set -euo pipefail

: "${TESSERA_API_BASE:?e.g. https://api.staging.tessera.zone/v1}"
: "${TESSERA_API_KEY:?the tsk_ key of the partner project}"
: "${TESSERA_MODEL_ID:?the trained model, mdl_...}"
: "${TESSERA_STATE_DIR:?set by the runner; ids are handed between steps}"
AUTH=(-H "Authorization: Bearer $TESSERA_API_KEY")
PROTOCOL_ID="$(cat "$TESSERA_STATE_DIR/protocol_id")"
MAX_POLLS="${TESSERA_MAX_POLLS:-720}"   # 5 s apart: one hour by default

fail() { echo "$1" >&2; exit 1; }

KEY="docs-partner-evaluation-$(printf '%s:%s' "$TESSERA_MODEL_ID" "$PROTOCOL_ID" | python3 -c '
import hashlib, sys
print(hashlib.sha256(sys.stdin.buffer.read()).hexdigest()[:24])
')"
EVALUATION="$(curl -sS --fail-with-body -X POST "$TESSERA_API_BASE/evaluations" "${AUTH[@]}" \
  -H "Content-Type: application/json" \
  -H "Idempotency-Key: $KEY" \
  -d '{"model_id": "'"$TESSERA_MODEL_ID"'", "protocol_id": "'"$PROTOCOL_ID"'"}')" \
  || fail "starting the evaluation failed: $EVALUATION"
REPORT_ID="$(printf '%s' "$EVALUATION" | python3 -c 'import json,sys; print(json.load(sys.stdin)["report_id"])')"

STATUS=""
for _ in $(seq 1 "$MAX_POLLS"); do
  DETAIL="$(curl -sS --fail-with-body "$TESSERA_API_BASE/evaluations/$REPORT_ID" "${AUTH[@]}")" \
    || fail "reading evaluation $REPORT_ID failed: $DETAIL"
  STATUS="$(printf '%s' "$DETAIL" | python3 -c 'import json,sys; print(json.load(sys.stdin)["evaluation"]["status"])')"
  if [ "$STATUS" = "succeeded" ] || [ "$STATUS" = "failed" ]; then break; fi
  sleep 5
done
echo "evaluation: $STATUS"
if [ "$STATUS" != "succeeded" ]; then
  printf '%s' "$DETAIL" | python3 -c '
import json, sys
row = json.load(sys.stdin)["evaluation"]
print("reason:", row.get("reason", "still running after the polling bound"), file=sys.stderr)
'
  exit 1
fi
printf '%s' "$DETAIL" | python3 -c '
import json, sys
verdict = json.load(sys.stdin)["report"]["verdict"]
print("verdict:", verdict["overall"])
print("provenance:", verdict["provenance"])
for criterion in verdict["criteria"]:
    print("  %s %s vs %s: %s" % (criterion["label"], criterion["metric"], criterion["arm"], criterion["status"]))
'
echo "$REPORT_ID" > "$TESSERA_STATE_DIR/report_id"