// Requests to the model, and webcam frames streamed to it. Each new camera frame
// is sent as it arrives, unless `inflight` frames are already waiting, so no
// stale frame queues up; an answer older than the last one delivered is dropped.
// Every answer is timed on this page's clock, for the HUD.

const token = new URLSearchParams(location.search).get("token") || "";
// Frames in flight at once. On this machine one: the model answers frames one after
// another, so a second only waits in line (round trip 64 ms with one, 125 ms with
// two, at the same 15 decisions/s on an M5 Pro). Across a network two, so a frame
// travels while the last one is answered.
const LOCAL = ["localhost", "127.0.0.1", "[::1]"].includes(location.hostname);
const INFLIGHT = LOCAL ? 1 : 2;
const recent = [];  // {at, latency, model} of the last two seconds, every request

// One request to the model: `{image}` (a JPEG data URL) or `{state}` (text or an
// object), and its questions. A question is `{text}` (yes/no), `{text, options}`
// (one of named options) or `{text, levels}` (an ordered score).
export async function decide(body) {
  const t0 = performance.now();
  const out = await (socket.usable() ? socket.ask(body) : post(body));
  const now = performance.now();
  recent.push({ at: now, latency: now - t0, model: out.ms });
  while (recent.length && now - recent[0].at > 2000) recent.shift();
  return out;
}

async function post(body) {
  const res = await fetch("decide", {
    method: "POST",
    headers: { "Content-Type": "application/json", "X-Token": token },
    body: JSON.stringify(body),
  });
  if (!res.ok) throw new Error(`model server: ${res.status}`);
  return res.json();
}

// One WebSocket for every request: no HTTP headers per frame, and behind a proxy
// (a Space) no per-request checks. Answers carry their request's id. A socket that
// never opens sends the page back to plain POSTs for good.
const socket = {
  ws: null, pending: new Map(), next: 0, broken: false, ready: null,

  usable() { return !this.broken; },

  open() {
    if (this.ready) return this.ready;
    const url = new URL(`ws${token ? `?token=${encodeURIComponent(token)}` : ""}`, location.href);
    url.protocol = url.protocol === "https:" ? "wss:" : "ws:";
    const ws = new WebSocket(url);
    this.ready = new Promise((resolve, reject) => {
      ws.onopen = () => resolve(ws);
      ws.onerror = () => reject(new Error("websocket failed"));
    });
    ws.onmessage = (e) => {
      const out = JSON.parse(e.data);
      const p = this.pending.get(out.id);
      if (!p) return;
      this.pending.delete(out.id);
      out.error ? p.reject(new Error(`model server: ${out.error}`)) : p.resolve(out);
    };
    ws.onclose = () => {
      for (const p of this.pending.values()) p.reject(new Error("websocket closed"));
      this.pending.clear();
      this.ready = null;  // the next request opens a new one
    };
    this.ws = ws;
    return this.ready;
  },

  async ask(body) {
    let ws;
    try {
      ws = await this.open();
    } catch (e) {
      this.broken = true;
      this.ready = null;
      return post(body);
    }
    const id = this.next++;
    return new Promise((resolve, reject) => {
      this.pending.set(id, { resolve, reject });
      ws.send(JSON.stringify({ ...body, id }));
    });
  },
};

// decisions per second, and the median round trip and model time, over two seconds
export function summary() {
  const now = performance.now();
  const live = recent.filter((s) => now - s.at < 2000);
  const median = (xs) => xs.sort((a, b) => a - b)[Math.floor(xs.length / 2)] ?? 0;
  return { rate: live.length / 2, latency: median(live.map((s) => s.latency)), model: median(live.map((s) => s.model)) };
}

// The expected level of a score's distribution, 0 to levels - 1.
export const expected = (probs) => probs.reduce((s, p, i) => s + i * p, 0);

export class Stream {
  constructor(video) {
    this.video = video;
    this.canvas = document.createElement("canvas");
  }

  // `questions()` is read at every frame, so a game can change them between phases.
  // `crop` keeps a horizontal slice of the camera frame, as fractions of its width.
  start({ questions, side = 384, inflight = INFLIGHT, crop = [0, 1] }, onAnswer) {
    this.stop();
    const run = { alive: true, sent: 0, shown: -1, inFlight: 0 };
    this.run = run;
    // a frame goes out at each camera frame while fewer than `inflight` wait, and at each answer: the model is
    // never left idle until the camera's next frame (up to 33 ms at 30 fps)
    const pump = () => {
      const qs = run.alive ? questions() : [];
      if (run.inFlight < inflight && this.video.videoWidth && qs.length) this.#send(run, qs, side, crop, onAnswer, pump);
    };
    const onFrame = () => {
      if (!run.alive) return;
      pump();
      this.video.requestVideoFrameCallback(onFrame);
    };
    this.video.requestVideoFrameCallback(onFrame);
  }

  stop() {
    if (this.run) this.run.alive = false;
  }

  async #send(run, questions, side, crop, onAnswer, next) {
    const id = run.sent++;
    const frame = this.snapshot(side, crop);
    const captured = performance.now();
    run.inFlight++;
    let answered = false;
    try {
      const out = await decide({ image: frame, questions });
      answered = true;
      if (!run.alive || id <= run.shown) return;
      run.shown = id;
      onAnswer({ probs: out.probs, model: out.ms, captured, latency: performance.now() - captured, frame, questions });
    } catch (e) {
      console.warn("decide failed", e);
    } finally {
      run.inFlight--;
      if (answered) next();  // after a failure, the camera's next frame: no loop against a server that is down
    }
  }

  // the current camera frame as a JPEG data URL, the full frame's longer side at most `side`
  snapshot(side = 384, [x0, x1] = [0, 1]) {
    const { videoWidth: vw, videoHeight: h } = this.video;
    const w = vw * (x1 - x0);
    const scale = Math.min(1, side / Math.max(vw, h));
    this.canvas.width = Math.round(w * scale);
    this.canvas.height = Math.round(h * scale);
    this.canvas.getContext("2d").drawImage(this.video, vw * x0, 0, w, h, 0, 0, this.canvas.width, this.canvas.height);
    return this.canvas.toDataURL("image/jpeg", 0.8);
  }
}