static/app.js

import { ReachyMini } from "https://cdn.jsdelivr.net/gh/pollen-robotics/reachy_mini@v1.7.0/js/reachy-mini.js";

const $ = id => document.getElementById(id);
let robot = null;
let hfToken = null;
let frameBuffer = [];  // circular buffer of recent frames for sharp-pick

// ─── Init ───
document.addEventListener("DOMContentLoaded", async () => {
  let clientId = null;
  try {
    const cfg = await (await fetch("/api/oauth-config")).json();
    clientId = cfg.clientId || null;
  } catch (e) { console.warn("OAuth config fetch failed:", e); }

  robot = new ReachyMini({ enableMicrophone: false, clientId });
  attachRobotListListener();

  if (clientId && await robot.authenticate()) {
    hfToken = getStoredToken();
    showRobotSelect();
  }
  initEvents();
});

function getStoredToken() {
  return sessionStorage.getItem("hf_token");
}

function showRobotSelect() {
  $("loginScreen").hidden = true;
  $("robotScreen").hidden = false;
  $("username").textContent = robot.username;
  connectSignaling();
  refreshQuota();
}

function showApp(robotInfo) {
  $("robotScreen").hidden = true;
  $("appScreen").hidden = false;
  $("robotName").textContent = robotInfo.meta?.name || robotInfo.id;
}

// ─── Events ───
function initEvents() {
  $("loginBtn").addEventListener("click", () => robot.login());
  $("logoutBtn").addEventListener("click", () => {
    robot.logout();
    location.reload();
  });
  $("disconnectBtn").addEventListener("click", async () => {
    await robot.stopSession();
    $("appScreen").hidden = true;
    $("robotScreen").hidden = false;
  });

  $("makeBtn").addEventListener("click", runMakeMe);
  $("againBtn").addEventListener("click", () => {
    $("revealSection").hidden = true;
    $("promptInput").focus();
  });

  for (const chip of document.querySelectorAll(".chip")) {
    chip.addEventListener("click", () => {
      $("promptInput").value = chip.dataset.prompt;
      $("promptInput").focus();
    });
  }

  // Mic button — push-to-talk
  const micBtn = $("micBtn");
  micBtn.addEventListener("mousedown", startRecording);
  micBtn.addEventListener("touchstart", (e) => { e.preventDefault(); startRecording(); });
  micBtn.addEventListener("mouseup", stopRecording);
  micBtn.addEventListener("mouseleave", stopRecording);
  micBtn.addEventListener("touchend", stopRecording);

  // Spacebar shortcut
  document.addEventListener("keydown", (e) => {
    if (e.code === "Space" && !e.repeat && !e.target.matches("input,textarea,button")) {
      e.preventDefault();
      startRecording();
    }
  });
  document.addEventListener("keyup", (e) => {
    if (e.code === "Space" && !e.target.matches("input,textarea")) {
      e.preventDefault();
      if (mediaRec?.state === "recording") stopRecording();
    }
  });

  // Enter in prompt = Make me
  $("promptInput").addEventListener("keydown", (e) => {
    if (e.key === "Enter") runMakeMe();
  });
}

// ─── Signaling / robot selection ───
async function connectSignaling() {
  try { await robot.connect(); }
  catch (e) {
    $("robotList").innerHTML = '<p class="error">Connection failed: ' + e.message + '</p>';
  }
}

function attachRobotListListener() {
  robot.addEventListener("robotsChanged", (e) => {
    const list = e.detail.robots;
    const container = $("robotList");
    if (list.length === 0) {
      container.innerHTML = '<p class="muted">No robots online</p>';
      return;
    }
    container.innerHTML = "";
    for (const r of list) {
      const btn = document.createElement("button");
      btn.className = "robot-btn";
      btn.textContent = r.meta?.name || r.id;
      btn.addEventListener("click", () => selectRobot(r));
      container.appendChild(btn);
    }
  });
}

async function selectRobot(r) {
  showApp(r);
  robot.attachVideo($("videoEl"));
  try {
    await robot.startSession(r.id);
    $("videoEl").play().catch(() => {});
    startFrameBuffer();
  } catch (e) {
    alert("Session failed: " + e.message);
    $("appScreen").hidden = true;
    $("robotScreen").hidden = false;
  }
}

// ─── Frame buffer (circular, for sharp-pick on capture) ───
const BUFFER_SIZE = 10;
let bufferStopFlag = false;

function startFrameBuffer() {
  const video = $("videoEl");
  frameBuffer = [];
  bufferStopFlag = false;

  const push = () => {
    if (bufferStopFlag) return;
    if (video.videoWidth && video.readyState >= 2) {
      // Store a low-res canvas for sharpness scoring only
      const c = document.createElement("canvas");
      const scale = 0.15;
      c.width = video.videoWidth * scale;
      c.height = video.videoHeight * scale;
      c.getContext("2d").drawImage(video, 0, 0, c.width, c.height);
      // Also store timestamp for picking "frame at T"
      frameBuffer.push({ time: performance.now(), thumb: c });
      if (frameBuffer.length > BUFFER_SIZE) frameBuffer.shift();
    }
    if ("requestVideoFrameCallback" in video) {
      video.requestVideoFrameCallback(push);
    } else {
      requestAnimationFrame(push);
    }
  };
  if ("requestVideoFrameCallback" in video) {
    video.requestVideoFrameCallback(push);
  } else {
    requestAnimationFrame(push);
  }
}

function sharpnessScore(canvas) {
  // Laplacian variance approximation — higher = sharper
  const ctx = canvas.getContext("2d");
  const { data, width, height } = ctx.getImageData(0, 0, canvas.width, canvas.height);
  let sum = 0, sumSq = 0, count = 0;
  // 3x3 Laplacian kernel applied to luma, sampled every 2 px for speed
  for (let y = 1; y < height - 1; y += 2) {
    for (let x = 1; x < width - 1; x += 2) {
      const i = (y * width + x) * 4;
      const l = 0.3 * data[i] + 0.6 * data[i + 1] + 0.1 * data[i + 2];
      const t = (((y - 1) * width + x) * 4);
      const b = (((y + 1) * width + x) * 4);
      const le = ((y * width + x - 1) * 4);
      const ri = ((y * width + x + 1) * 4);
      const lT = 0.3 * data[t] + 0.6 * data[t + 1] + 0.1 * data[t + 2];
      const lB = 0.3 * data[b] + 0.6 * data[b + 1] + 0.1 * data[b + 2];
      const lL = 0.3 * data[le] + 0.6 * data[le + 1] + 0.1 * data[le + 2];
      const lR = 0.3 * data[ri] + 0.6 * data[ri + 1] + 0.1 * data[ri + 2];
      const lap = 4 * l - lT - lB - lL - lR;
      sum += lap; sumSq += lap * lap; count++;
    }
  }
  const mean = sum / count;
  return sumSq / count - mean * mean;
}

// ─── Voice prompting ───
let mediaRec = null;
let chunks = [];
let micStream = null;
let recStartTime = 0;

async function startRecording() {
  if (mediaRec?.state === "recording") return;
  try {
    micStream = await navigator.mediaDevices.getUserMedia({
      audio: { echoCancellation: true, noiseSuppression: true, autoGainControl: true },
    });
  } catch (e) {
    setStatus("Mic access denied", true);
    return;
  }

  const mime = MediaRecorder.isTypeSupported("audio/webm;codecs=opus")
    ? "audio/webm;codecs=opus"
    : (MediaRecorder.isTypeSupported("audio/mp4") ? "audio/mp4" : "");
  chunks = [];
  mediaRec = new MediaRecorder(micStream, mime ? { mimeType: mime } : {});
  mediaRec.ondataavailable = (e) => { if (e.data.size > 0) chunks.push(e.data); };
  mediaRec.start();
  recStartTime = performance.now();

  $("micBtn").classList.add("recording");
  $("micBtn").querySelector(".mic-label").textContent = "● REC";
  setStatus("Listening...");

  // Robot: perk up antennas
  if (robot?.state === "streaming") {
    robot.setAntennas(30, -30);
    robot.setAudioMuted(true);  // suppress robot speaker to avoid echo
  }
}

async function stopRecording() {
  if (!mediaRec || mediaRec.state !== "recording") return;
  const duration = performance.now() - recStartTime;
  const rec = mediaRec;
  mediaRec = null;

  await new Promise((resolve) => {
    rec.onstop = resolve;
    rec.stop();
  });
  micStream?.getTracks().forEach(t => t.stop());
  micStream = null;

  $("micBtn").classList.remove("recording");
  $("micBtn").querySelector(".mic-label").textContent = "Hold";

  // Robot: nod + reset antennas
  if (robot?.state === "streaming") {
    robot.setAntennas(0, 0);
    robot.setAudioMuted(false);
    robot.setHeadPose(0, 5, 0);
    setTimeout(() => robot.setHeadPose(0, 0, 0), 250);
  }

  if (duration < 400) {
    setStatus("Too short — hold the mic while speaking");
    return;
  }

  const blob = new Blob(chunks, { type: rec.mimeType || "audio/webm" });
  setStatus("Transcribing...");

  try {
    const fd = new FormData();
    fd.append("audio", blob, "prompt.webm");
    const res = await fetch("/api/transcribe", {
      method: "POST",
      body: fd,
      headers: { "Authorization": `Bearer ${hfToken}` },
    });
    if (!res.ok) {
      const txt = await res.text();
      setStatus(`Transcribe failed: ${txt.slice(0, 200)}`, true);
      return;
    }
    const { text } = await res.json();
    if (text) {
      $("promptInput").value = text;
      setStatus(`Got it: "${text}"`);
    } else {
      setStatus("Didn't catch anything — try again");
    }
  } catch (e) {
    setStatus("Transcribe failed: " + e.message, true);
  }
}

// ─── Make Me flow ───
async function runMakeMe() {
  const prompt = $("promptInput").value.trim();
  if (!prompt) {
    setStatus("Type or speak a prompt first");
    $("promptInput").focus();
    return;
  }
  if (!robot || robot.state !== "streaming") {
    setStatus("Robot not connected");
    return;
  }

  $("makeBtn").disabled = true;

  try {
    await countdown(3);
    const blob = await captureSharpFrame();
    const beforeUrl = URL.createObjectURL(blob);
    $("beforeImg").src = beforeUrl;
    $("afterImg").src = beforeUrl;  // placeholder until transform completes
    $("downloadLink").hidden = true;
    $("revealSection").hidden = false;

    // Robot: thinking animation during render
    const thinkingTimer = startThinkingLoop();
    setStatus("Transforming...");

    const fd = new FormData();
    fd.append("image", blob, "capture.jpg");
    fd.append("prompt", prompt);
    const t0 = performance.now();
    const res = await fetch("/api/transform", {
      method: "POST",
      body: fd,
      headers: { "Authorization": `Bearer ${hfToken}` },
    });
    clearInterval(thinkingTimer);

    if (!res.ok) {
      const txt = await res.text();
      setStatus(`Transform failed: ${txt.slice(0, 200)}`, true);
      if (res.status === 429) setStatus("Daily limit reached. Come back tomorrow.", true);
      return;
    }
    const afterBlob = await res.blob();
    const afterUrl = URL.createObjectURL(afterBlob);
    $("afterImg").src = afterUrl;
    $("downloadLink").href = afterUrl;
    $("downloadLink").hidden = false;

    const elapsed = (performance.now() - t0) / 1000;
    $("timingLabel").textContent = `Rendered in ${elapsed.toFixed(1)}s`;
    setStatus("");

    // Robot: "ta-da" nod
    robot.setHeadPose(0, -5, 0);
    setTimeout(() => robot.setHeadPose(0, 0, 0), 350);

    refreshQuota();
  } finally {
    $("makeBtn").disabled = false;
  }
}

async function countdown(seconds) {
  const overlay = $("countdownOverlay");
  const num = $("countdownNumber");
  overlay.hidden = false;
  // Robot: center head, perk up for the photo
  robot.setHeadPose(0, 0, 0);
  robot.setAntennas(20, -20);
  for (let i = seconds; i >= 1; i--) {
    num.textContent = i;
    num.classList.remove("pulse");
    void num.offsetWidth;
    num.classList.add("pulse");
    await sleep(900);
  }
  overlay.hidden = true;
  // Flash
  const flash = $("flashOverlay");
  flash.classList.add("flash-active");
  setTimeout(() => flash.classList.remove("flash-active"), 120);
  robot.setAntennas(0, 0);
}

async function captureSharpFrame() {
  const video = $("videoEl");
  // Wait one more frame tick so the "flash moment" buffer is freshest
  await new Promise(r => setTimeout(r, 50));
  // Pick sharpest from buffer
  let bestThumb = null;
  let bestScore = -Infinity;
  for (const entry of frameBuffer) {
    const s = sharpnessScore(entry.thumb);
    if (s > bestScore) { bestScore = s; bestThumb = entry.thumb; }
  }
  // Draw full-res version of current frame (we can't retro-grab full-res, but sharpest thumb
  // was likely adjacent in time to current frame — so current full-res is close enough)
  const canvas = $("captureCanvas");
  canvas.width = video.videoWidth;
  canvas.height = video.videoHeight;
  canvas.getContext("2d").drawImage(video, 0, 0);
  return new Promise(r => canvas.toBlob(r, "image/jpeg", 0.92));
}

function startThinkingLoop() {
  let yaw = 0;
  let dir = 1;
  return setInterval(() => {
    yaw += dir * 8;
    if (yaw > 15 || yaw < -15) dir *= -1;
    if (robot?.state === "streaming") robot.setHeadPose(0, 0, yaw);
  }, 300);
}

// ─── Quota / status ───
async function refreshQuota() {
  try {
    const r = await fetch("/api/quota", {
      headers: { "Authorization": `Bearer ${hfToken}` },
    });
    if (!r.ok) return;
    const q = await r.json();
    const rem = q.transform.remaining;
    $("quotaBadge").textContent = `${rem}/${q.transform.limit} transforms left today`;
    $("quotaBadge").classList.toggle("warn", rem <= 2);
  } catch (_) {}
}

function setStatus(msg, isError = false) {
  const el = $("statusLine");
  el.textContent = msg;
  el.classList.toggle("error", isError);
}

function sleep(ms) { return new Promise(r => setTimeout(r, ms)); }