#!/usr/bin/env bats
# Seam: autoresearch measure client behavior.
#
# The measure client is the whole yardstick of the frank-archon speed
# session, so its METRIC arithmetic gets a behavior test: against a
# deterministic mock endpoint that emits a KNOWN number of completion
# tokens at a KNOWN per-token delay, the reported decode rate must land
# near the nominal rate, the usage accounting must come back exact, and an
# unhealthy endpoint must fail loudly (no half-measured numbers).
#
# The mock implements only what the client calls: POST
# /v1/messages/count_tokens (word count as the token count, which makes the
# client's words/token calibration come out to exactly 1.0) and a streaming
# POST /v1/chat/completions (one content chunk per completion token, then a
# usage chunk, then [DONE]).

load helpers

MEASURE="$REPO_ROOT/scripts/autoresearch/frank-archon-measure.py"
MOCK="$BATS_TEST_TMPDIR/mock.py"
PORT=""
MOCK_PID=""

setup() {
    # Deterministic mock; argv: <port> <delay-seconds>; env MOCK_FAIL=1
    # makes the chat endpoint answer 500.
    cat >"$MOCK" <<'PYEOF'
import json, os, sys, time
from http.server import BaseHTTPRequestHandler, HTTPServer

PORT = int(sys.argv[1])
DELAY = float(sys.argv[2])
FAIL = os.environ.get("MOCK_FAIL", "0") == "1"


class Handler(BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass

    def _send(self, obj, code=200):
        b = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(b)))
        self.end_headers()
        self.wfile.write(b)

    def do_POST(self):
        n = int(self.headers.get("Content-Length", "0"))
        body = json.loads(self.rfile.read(n) or b"{}")
        if self.path == "/v1/messages/count_tokens":
            text = body["messages"][0]["content"]
            self._send({"input_tokens": max(1, len(text.split()))})
        elif self.path == "/v1/chat/completions":
            if FAIL:
                self._send({"error": "engine down"}, 500)
                return
            max_tokens = int(body.get("max_tokens", 64))
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.end_headers()
            for i in range(max_tokens):
                ev = {"choices": [{"delta": {"content": "tok%d " % i}}]}
                self.wfile.write(("data: %s\n\n" % json.dumps(ev)).encode())
                self.wfile.flush()
                time.sleep(DELAY)
            prompt_words = max(1, len(body["messages"][0]["content"].split()))
            ev = {"choices": [{"delta": {}}],
                  "usage": {"prompt_tokens": prompt_words,
                            "completion_tokens": max_tokens}}
            self.wfile.write(("data: %s\n\n" % json.dumps(ev)).encode())
            self.wfile.write(b"data: [DONE]\n\n")
            self.wfile.flush()
        else:
            self._send({"error": "unknown path"}, 404)


HTTPServer(("127.0.0.1", PORT), Handler).serve_forever()
PYEOF
}
free_port() {
    python3 -c 'import socket; s = socket.socket(); s.bind(("127.0.0.1", 0)); print(s.getsockname()[1]); s.close()'
}

start_mock() { # $1 = ok|fail
    PORT="$(free_port)"
    if [ "$1" = "fail" ]; then
        MOCK_FAIL=1 python3 "$MOCK" "$PORT" 0.005 &
    else
        python3 "$MOCK" "$PORT" 0.005 &
    fi
    MOCK_PID=$!
    local i
    for i in $(seq 1 100); do
        if python3 -c "import socket; socket.create_connection(('127.0.0.1', $PORT), 0.2).close()" 2>/dev/null; then
            return 0
        fi
        sleep 0.1
    done
    echo "mock did not start" >&2
    return 1
}

stop_mock() {
    if [ -n "$MOCK_PID" ]; then
        kill "$MOCK_PID" 2>/dev/null || true
        wait "$MOCK_PID" 2>/dev/null || true
        MOCK_PID=""
    fi
}

@test "measure exists and compiles" {
    [ -f "$MEASURE" ]
    python3 -m py_compile "$MEASURE"
}

@test "measure computes decode rate and exact usage from a controlled endpoint" {
    start_mock ok || return 1
    run env AR_BASE="http://127.0.0.1:$PORT" AR_MODEL=mock python3 "$MEASURE"
    stop_mock
    [ "$status" -eq 0 ] || { echo "$output"; return 1; }
    # run() merges stderr in; the contract is: exactly ONE json result line
    # (the window script parses the same way), progress on stderr.
    printf '%s\n' "$output" | python3 -c '
import json, sys
lines = [l for l in sys.stdin.read().splitlines() if l.startswith("{")]
assert len(lines) == 1, "expected exactly one JSON line, got %d" % len(lines)
d = json.loads(lines[0])
# nominal 200 t/s (0.005 s/chunk) minus small per-chunk overhead
for k in ("decode_tps", "decode_tps_115", "decode_tps_1056", "decode_tps_4096"):
    v = d[k]
    assert 60.0 <= v <= 320.0, "%s=%s outside [60, 320]" % (k, v)
# the mock emits exactly max_tokens completion tokens; the accounting must be exact
assert d["completion_tokens_4096"] == 256, d["completion_tokens_4096"]
assert d["completion_tokens_115"] == 256
# prompt sizing: calibration is exactly 1.0 words/token in the mock, so the
# measured prompt lengths equal the targets
assert d["prompt_tokens_115"] == 115
assert d["prompt_tokens_4096"] == 4096
print("measure claims verified")
'
}

@test "measure fails loudly on an unhealthy endpoint instead of half-measuring" {
    start_mock fail || return 1
    run env AR_BASE="http://127.0.0.1:$PORT" AR_MODEL=mock python3 "$MEASURE"
    stop_mock
    [ "$status" -ne 0 ]
    ! printf '%s\n' "$output" | grep -qE '^\{'
}