"""Memory a FRESH seat-config instance (A') needs per input length: ascending lengths, one request each,
GPU memory of its PID read 2 s after each response. ORT's arena never shrinks, so the reading after a
request is the high-water mark so far. Prefixes of the public SCOTUS audio (16 kHz mono).
usage: len_sweep.py ARM PORT SECONDS..."""
import io
import json
import subprocess
import sys
import time

import soundfile as sf

sys.path.insert(0, "/tank/spikes/parakeet-ab/code")
from bench import post  # noqa: E402

arm, port = sys.argv[1], sys.argv[2]
lens = [float(x) for x in sys.argv[3:]]
pids = set(subprocess.run(["docker", "top", arm, "-eo", "pid"], capture_output=True, text=True).stdout.split()[1:])
a, sr = sf.read("/tank/spikes/scriberr-slicer/public/scotus.wav", dtype="float32")


def mem():
    out = subprocess.run(["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits"],
                         capture_output=True, text=True).stdout
    return sum(int(l.split(",")[1]) for l in out.splitlines() if l.split(",")[0].strip() in pids)


print(json.dumps(dict(arm=arm, seconds=0, status=None, mib=mem())), flush=True)
for s in lens:
    bio = io.BytesIO()
    sf.write(bio, a[: int(s * sr)], sr, format="WAV", subtype="PCM_16")
    r = post(f"http://127.0.0.1:{port}/v1/audio/transcriptions", bio.getvalue(), timeout=1800)
    time.sleep(2)
    print(json.dumps(dict(arm=arm, seconds=s, status=r["status"], e2e_ms=r["e2e_ms"], mib=mem())), flush=True)