#!/usr/bin/env bash
# Smoke-test the running orchestrator using slices of Project Gutenberg #84
# (Frankenstein). Exercises:
#   1. short only             -> Swift
#   2. mixed short texts      -> Swift
#   3. long only              -> spawns Rust backend
#   4. mixed short + long     -> Swift in parallel with Rust
#   5. > 32k tokens           -> HTTP 400
#
# Assumes the orchestrator is already listening on :8080.
# Run scripts/run-orchestrator.sh in another terminal first.

set -euo pipefail

ORCH="${ORCH:-http://127.0.0.1:8080}"
TEXT="${TEXT:-/tmp/pg84.txt}"
GUTENBERG_URL="https://www.gutenberg.org/cache/epub/84/pg84.txt"

# Fetch the book if we don't have it yet.
if [ ! -f "$TEXT" ]; then
    echo "==> downloading $GUTENBERG_URL -> $TEXT"
    curl -sSL -o "$TEXT" "$GUTENBERG_URL"
fi

# Sanity: orchestrator reachable?
if ! curl -sf -o /dev/null "$ORCH/healthz"; then
    echo "orchestrator not reachable at $ORCH/healthz"
    echo "start it first: scripts/run-orchestrator.sh"
    exit 1
fi

# Drive the tests through a Python helper so we can slice the text and
# measure wall-clock per request cleanly.
python3 - "$ORCH" "$TEXT" <<'PY'
import json, sys, time, urllib.request, urllib.error

orch, text_path = sys.argv[1], sys.argv[2]
book = open(text_path, encoding="utf-8").read()

# Character slices. Token counts are rough: ~1 token per 4 chars for English.
slices = {
    "short  (~125 tokens)"  : book[1000:1500],
    "medium (~1k tokens)"   : book[2000:6000],
    "long   (~15k tokens)"  : book[5000:65000],
    "overflow (~50k tokens)": book[:200000],
}

def call(label, texts, expect_error=False):
    body = json.dumps({"texts": texts, "input_type": "document"}).encode()
    req = urllib.request.Request(
        f"{orch}/embed",
        data=body,
        headers={"Content-Type": "application/json"},
    )
    t0 = time.time()
    try:
        r = json.loads(urllib.request.urlopen(req, timeout=600).read())
        dt = time.time() - t0
        print(f"  {label:40s} ok   n={len(r['vectors'])} dim={r['dim']} "
              f"t={dt:.2f}s")
        return r
    except urllib.error.HTTPError as e:
        dt = time.time() - t0
        msg = e.read().decode()
        tag = "PASS" if expect_error else "FAIL"
        print(f"  {label:40s} {e.code}  t={dt:.2f}s {msg}  [{tag}]")
        if not expect_error:
            sys.exit(1)
        return None

print("== short only (Swift path) ==")
a = call("short only", [slices["short  (~125 tokens)"]])

print("== mixed short texts (Swift path) ==")
call("short + medium", [
    slices["short  (~125 tokens)"],
    slices["medium (~1k tokens)"],
])

print("== long only (spawns Rust) ==")
b = call("long only", [slices["long   (~15k tokens)"]])

print("== mixed short + long (Swift || Rust) ==")
c = call("short + long", [
    slices["short  (~125 tokens)"],
    slices["long   (~15k tokens)"],
])

print("== overflow rejection ==")
call("40k+ tokens", [slices["overflow (~50k tokens)"]], expect_error=True)

# Sanity: the short vector in the mixed call should match the short-only call
# bit-for-bit (shorts went through Swift unchanged; merge order preserved).
if a and c:
    v1 = a["vectors"][0][:5]
    v2 = c["vectors"][0][:5]
    match = v1 == v2
    print(f"\n  short-vec parity short-only vs mixed: "
          f"{'OK' if match else 'MISMATCH'}")
    if not match:
        print(f"    short-only : {v1}")
        print(f"    mixed      : {v2}")
        sys.exit(1)

print("\nall smoke tests passed")
PY