#!/usr/bin/env bash
# End-to-end smoke tests for the Gemma 4 stack. Exit non-zero if anything fails.
#   ./scripts/e2e-test.sh
# Env overrides: LS=http://localhost:8080  EMB=http://localhost:8081
#                RRK=http://localhost:8082  OWUI=http://localhost:3000
#                MODEL=gemma-e4b  TEST_WAV=/path/clip.wav  TEST_IMG=/path/img.png
set -u

LS="${LS:-http://localhost:8080}"
EMB="${EMB:-http://localhost:8081}"
RRK="${RRK:-http://localhost:8082}"
OWUI="${OWUI:-http://localhost:3000}"
MODEL="${MODEL:-gemma-e4b}"
# Auto-use the bundled demo assets so `make test` needs no setup.
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
TEST_IMG="${TEST_IMG:-$ROOT/assets/test_image.png}"
TEST_WAV="${TEST_WAV:-$ROOT/assets/test_audio_en.wav}"

PASS=0; FAIL=0
ok(){ echo "  ✅ $1"; PASS=$((PASS+1)); }
no(){ echo "  ❌ $1"; FAIL=$((FAIL+1)); }
hdr(){ echo; echo "== $1 =="; }

jqget(){ python3 -c "import sys,json;d=json.load(sys.stdin);print(eval(sys.argv[1]))" "$1" 2>/dev/null; }

hdr "1) llama-swap lists model profiles"
MODELS=$(curl -fsS "$LS/v1/models" 2>/dev/null)
echo "$MODELS" | grep -q "$MODEL" && ok "model '$MODEL' present" || no "model '$MODEL' missing"

hdr "2) Gemma text chat (Persian)"
RESP=$(curl -fsS "$LS/v1/chat/completions" -H 'Content-Type: application/json' -d \
  "{\"model\":\"$MODEL\",\"messages\":[{\"role\":\"user\",\"content\":\"به فارسی یک جمله کوتاه بنویس.\"}],\"max_tokens\":512,\"temperature\":1.0}" 2>/dev/null)
TXT=$(echo "$RESP" | jqget "(d['choices'][0]['message'].get('content') or d['choices'][0]['message'].get('reasoning_content') or '')")
[ -n "$TXT" ] && ok "got reply: ${TXT:0:60}" || no "no text reply ($RESP)"

hdr "3) Gemma vision (image)"
if [ -n "$TEST_IMG" ] && [ -f "$TEST_IMG" ]; then
  python3 -c "import json,base64;b=base64.b64encode(open('$TEST_IMG','rb').read()).decode();json.dump({'model':'$MODEL','messages':[{'role':'user','content':[{'type':'text','text':'Describe this image in one short sentence.'},{'type':'image_url','image_url':{'url':'data:image/png;base64,'+b}}]}],'max_tokens':512},open('/tmp/e2e_img.json','w'))"
  R=$(curl -s -m 600 "$LS/v1/chat/completions" -H 'Content-Type: application/json' -d @/tmp/e2e_img.json 2>/dev/null)
  V=$(echo "$R" | jqget "(d['choices'][0]['message'].get('content') or d['choices'][0]['message'].get('reasoning_content') or '')")
  [ -n "$V" ] && ok "vision reply: ${V:0:60}" || no "no vision reply ($R)"
else
  echo "  ⏭  skipped (set TEST_IMG=/path/img.png)"
fi

hdr "4) Gemma native audio (input_audio)"
if [ -n "$TEST_WAV" ] && [ -f "$TEST_WAV" ]; then
  python3 -c "import json,base64;b=base64.b64encode(open('$TEST_WAV','rb').read()).decode();json.dump({'model':'$MODEL','messages':[{'role':'user','content':[{'type':'text','text':'Transcribe this audio exactly.'},{'type':'input_audio','input_audio':{'data':b,'format':'wav'}}]}],'max_tokens':512,'temperature':1.0},open('/tmp/e2e_aud.json','w'))"
  R=$(curl -s -m 1500 "$LS/v1/chat/completions" -H 'Content-Type: application/json' -d @/tmp/e2e_aud.json 2>/dev/null)
  A=$(echo "$R" | jqget "(d['choices'][0]['message'].get('content') or d['choices'][0]['message'].get('reasoning_content') or '')")
  [ -n "$A" ] && ok "audio reply: ${A:0:80}" || no "no audio reply ($R)"
else
  echo "  ⏭  skipped (set TEST_WAV=/path/clip.wav)"
fi

hdr "5) TEI embeddings (Persian)"
R=$(curl -fsS "$EMB/v1/embeddings" -H 'Content-Type: application/json' -d \
  '{"model":"x","input":["پایتخت ایران تهران است"]}' 2>/dev/null)
DIM=$(echo "$R" | jqget "len(d['data'][0]['embedding'])")
[ -n "$DIM" ] && [ "$DIM" -gt 0 ] 2>/dev/null && ok "embedding dim=$DIM" || no "no embedding ($R)"

hdr "6) TEI rerank (Persian relevance ordering)"
R=$(curl -fsS "$RRK/rerank" -H 'Content-Type: application/json' -d \
  '{"query":"پایتخت ایران کجاست؟","texts":["تهران پایتخت ایران است.","موز یک میوه است."]}' 2>/dev/null)
TOP=$(echo "$R" | python3 -c "import sys,json;d=json.load(sys.stdin);d=sorted(d,key=lambda x:-x['score']);print(d[0]['index'])" 2>/dev/null)
[ "$TOP" = "0" ] && ok "relevant text ranked first" || no "rerank wrong/failed (top=$TOP, $R)"

hdr "7) Open WebUI reachable"
CODE=$(curl -s -o /dev/null -w "%{http_code}" "$OWUI/health" 2>/dev/null)
[ "$CODE" = "200" ] && ok "Open WebUI /health 200" || no "Open WebUI not healthy (HTTP $CODE)"

echo; echo "=================================================="
echo "  RESULT: $PASS passed, $FAIL failed"
echo "=================================================="
[ "$FAIL" -eq 0 ]