#!/bin/bash
# server boot gates + smoke + 50 timed requests with the tt-model.yaml serve env (chip from CHIP, port from PORT)
# M defaults to the repo this script lives in (so a verify worktree can run its own copy unmodified)
M=${M:-$(cd "$(dirname "$0")/../.." && pwd)}
VENV=${VENV:-/home/ttuser/experiments/tt-models/models/hamer-p150/.venv}
W=${W:-$M/.work}; mkdir -p $W
cd $M
source /home/ttuser/experiments/tt-models/tools/chipenv.sh ${CHIP:?} $VENV >/dev/null 2>&1 || exit 1
eval "$(python3 - <<'PY'
import yaml
d = yaml.safe_load(open("tt-model.yaml"))
for k, v in d["serve"]["env"].items():
    if k == "TT_METAL_VISIBLE_DEVICES":
        continue  # chipenv pins the chip by BDF
    print(f"export {k}={str(v)!r}")
PY
)"
[ -n "$EXTRA" ] && export $EXTRA
# uvicorn flags: tt-model.yaml serve.args (what the image launches) unless UVI is set (UVI="--loop asyncio --http h11"
# restores uvicorn's pre-integration stack for an A/B)
UVI=${UVI-$(python3 -c "
import yaml
a = yaml.safe_load(open('tt-model.yaml'))['serve'].get('args', [])
print(' '.join(' '.join(map(str, x)) if isinstance(x, list) else str(x) for x in a))")}
echo "uvicorn flags: $UVI"
export PYTHONPATH=$M/code TT_MESH_SHAPE=1x1 TT_DEVICE_ID=0 HAMER_WARMUP_ITERS=3 HF_MODEL=changh95/hamer-weights
env | grep -E "^(TT_FUSED|HAMER_)" | sort > $W/serve_env.txt
PORT=${PORT:-20117}
rm -f $W/serve.log
timeout -s INT 900 python -m uvicorn --host 127.0.0.1 --port $PORT --lifespan on $UVI hamer.server.app:app > $W/serve.log 2>&1 &
SP=$!
for i in $(seq 1 300); do
  grep -q "Application startup complete" $W/serve.log && break
  kill -0 $SP 2>/dev/null || break
  sleep 1
done
grep -E "warm-up PCC|Warmup complete|Application startup complete|Error|Traceback" $W/serve.log | head -8
if grep -q "Application startup complete" $W/serve.log; then
  python code/hamer/server/smoke_test.py --url http://127.0.0.1:$PORT --image all 2>&1 | tail -10
  python3 - $PORT <<'PY'
import base64, json, sys, time, statistics, urllib.request
from pathlib import Path
port = sys.argv[1]
media = Path("media"); boxes = json.loads((media / "bboxes.json").read_text())
reqs = []
for name, ents in sorted(boxes.items()):
    img = base64.b64encode((media / name).read_bytes()).decode()
    for e in ents:
        reqs.append({"image": img, "bbox": e["bbox"], "is_right": bool(e.get("is_right", 1)), "return_vertices": False})
mt, tot, pre, rtt = [], [], [], []
N = int(__import__('os').environ.get('NREQ', '60'))
for i in range(N):
    body = json.dumps(reqs[i % len(reqs)]).encode()
    r = urllib.request.Request(f"http://127.0.0.1:{port}/predict", data=body, headers={"Content-Type": "application/json"})
    t0 = time.perf_counter()
    raw = urllib.request.urlopen(r, timeout=60).read()
    dt = (time.perf_counter() - t0) * 1e3
    out = json.loads(raw)
    if i >= 10:
        rtt.append(dt)
        t = out.get("timing_ms", {})
        mt.append(t.get("model", t.get("device", 0))); tot.append(t.get("total", 0)); pre.append(t.get("preprocess", 0))
print("timing keys", list(out.get("timing_ms", {}).keys()))
print(f"TIMED {len(mt)} requests: model median {statistics.median(mt):.3f} min {min(mt):.3f}  total median {statistics.median(tot):.3f}")
# round 11: the server total is host-load-bound (shared host): report its spread and the load average with it
q = lambda v, p: sorted(v)[min(len(v) - 1, int(p * len(v)))]
import os as _os
print(f"TIMED spread: model p10 / p50 / p90 {q(mt, .1):.3f} / {q(mt, .5):.3f} / {q(mt, .9):.3f}  "
      f"total p10 / p50 / p90 {q(tot, .1):.3f} / {q(tot, .5):.3f} / {q(tot, .9):.3f}  "
      f"preprocess p50 {q(pre, .5):.3f}  client RTT p10 / p50 / p90 {q(rtt, .1):.3f} / {q(rtt, .5):.3f} / {q(rtt, .9):.3f}  host load average {' / '.join(f'{v:.1f}' for v in _os.getloadavg())}")
PY
fi
pkill -INT -P $SP 2>/dev/null; wait $SP 2>/dev/null; echo "server exit $?"