main.js

// @ts-check
/**
 * Minimal voice conversation app, talking to a Hugging Face speech-to-speech
 * backend over WebRTC.
 *
 * Click the orb -> we ask for the mic, POST a session on the LB, post our
 * SDP offer to the routed compute, set up the data channel and the audio
 * tracks. The orb visually reflects the live state (idle, connecting,
 * listening, user-speaking, processing, ai-speaking).
 *
 * @typedef {"idle" | "connecting" | "listening" | "user-speaking" | "processing" | "ai-speaking" | "error"} AppState
 */

import { S2sRealtimeClient } from "./s2s-client.js";

const DEFAULT_LB_URL = "https://mewklqcmavs5gkid.us-east-1.aws.endpoints.huggingface.cloud";
const DEFAULT_VOICE = "Aiden";
const DEFAULT_INSTRUCTIONS =
  "You are a friendly voice assistant. " +
  "Keep replies short, warm, and spoken. Avoid long monologues.";

const STORAGE_KEYS = {
  lbUrl: "s2s.lbUrl",
  voice: "s2s.voice",
  instructions: "s2s.instructions",
};

// ─── Settings (localStorage) ────────────────────────────────────────────

/**
 * @typedef {Object} Settings
 * @property {string} lbUrl
 * @property {string} voice
 * @property {string} instructions
 *
 * @returns {Settings}
 */
function loadSettings() {
  return {
    lbUrl: localStorage.getItem(STORAGE_KEYS.lbUrl) || DEFAULT_LB_URL,
    voice: localStorage.getItem(STORAGE_KEYS.voice) || DEFAULT_VOICE,
    instructions: localStorage.getItem(STORAGE_KEYS.instructions) || DEFAULT_INSTRUCTIONS,
  };
}

/** @param {Settings} s */
function saveSettings(s) {
  localStorage.setItem(STORAGE_KEYS.lbUrl, s.lbUrl);
  localStorage.setItem(STORAGE_KEYS.voice, s.voice);
  localStorage.setItem(STORAGE_KEYS.instructions, s.instructions);
}

// ─── State views ────────────────────────────────────────────────────────

/**
 * @typedef {Object} StateView
 * @property {string} caption
 * @property {boolean} disabled
 *
 * @type {Record<AppState, StateView>}
 */
const STATE_VIEWS = {
  idle:            { caption: "Tap to start",  disabled: false },
  connecting:      { caption: "Connecting",    disabled: true  },
  listening:       { caption: "",              disabled: false },
  "user-speaking": { caption: "",              disabled: false },
  processing:      { caption: "",              disabled: false },
  "ai-speaking":   { caption: "",              disabled: false },
  error:           { caption: "Tap to retry",  disabled: false },
};

/** @type {Record<AppState, string>} */
const STATE_CLASS = {
  idle: "state-idle",
  connecting: "state-connecting",
  listening: "state-listening",
  "user-speaking": "state-user-speaking",
  processing: "state-processing",
  "ai-speaking": "state-ai-speaking",
  error: "state-error",
};

/** @type {ReadonlySet<AppState>} */
const LIVE_STATES = new Set(["listening", "user-speaking", "processing", "ai-speaking"]);

// ─── DOM refs ───────────────────────────────────────────────────────────

/**
 * @template {HTMLElement} T
 * @param {string} selector
 * @returns {T}
 */
function $(selector) {
  const el = document.querySelector(selector);
  if (!el) throw new Error(`Missing element: ${selector}`);
  return /** @type {T} */ (el);
}

/** @type {HTMLButtonElement} */
const circleBtn = $("#main-circle");
/** @type {HTMLParagraphElement} */
const circleCaption = $("#circle-caption");
/** @type {HTMLElement} */
const orbWrap = $(".orb-wrap");
/** @type {HTMLButtonElement} */
const micBtn = $("#mic-btn");
/** @type {HTMLButtonElement} */
const stopBtn = $("#stop-btn");

/** @type {HTMLButtonElement} */
const settingsBtn = $("#settings-btn");
/** @type {HTMLDialogElement} */
const settingsModal = $("#settings-modal");
/** @type {HTMLInputElement} */
const inputLbUrl = $("#lb-url");
/** @type {HTMLSelectElement} */
const inputVoice = $("#voice");
/** @type {HTMLTextAreaElement} */
const inputInstructions = $("#instructions");
/** @type {HTMLButtonElement} */
const restartBtn = $("#restart-conversation");
/** @type {HTMLElement} */
const restartHint = $("#restart-hint");
const settingsForm = /** @type {HTMLFormElement} */ (settingsModal.querySelector("form"));

// ─── Runtime state ──────────────────────────────────────────────────────

/** @type {AppState} */
let currentState = "idle";
let settings = loadSettings();

/** @type {S2sRealtimeClient | null} */
let client = null;
/** @type {MediaStream | null} */
let micStream = null;
/** @type {HTMLAudioElement | null} */
let outputSink = null;
let micMuted = false;

// ─── UI plumbing ────────────────────────────────────────────────────────

/** @param {AppState} next */
function setState(next) {
  currentState = next;
  const view = STATE_VIEWS[next];
  circleBtn.disabled = view.disabled;
  circleBtn.className = `circle ${STATE_CLASS[next]}`;
  if (next !== "error") setCaption(view.caption);

  const live = LIVE_STATES.has(next);
  orbWrap.classList.toggle("live", live);
  micBtn.setAttribute("aria-hidden", live ? "false" : "true");
  stopBtn.setAttribute("aria-hidden", live ? "false" : "true");
  micBtn.tabIndex = live ? 0 : -1;
  stopBtn.tabIndex = live ? 0 : -1;

  updateRestartAvailability();
}

function updateRestartAvailability() {
  const live = LIVE_STATES.has(currentState);
  restartBtn.disabled = !live;
  restartHint.hidden = live;
}

/**
 * @param {string} text
 * @param {"" | "error" | "muted"} [kind]
 */
function setCaption(text, kind = "") {
  const trimmed = text.trim();
  circleCaption.textContent = trimmed;
  circleCaption.className = `circle-caption${kind ? ` ${kind}` : ""}${trimmed ? "" : " empty"}`;
}

// ─── Settings modal ─────────────────────────────────────────────────────

function openSettings() {
  inputLbUrl.value = settings.lbUrl;
  inputVoice.value = settings.voice;
  inputInstructions.value = settings.instructions;
  updateRestartAvailability();
  settingsModal.showModal();
}

settingsBtn.addEventListener("click", openSettings);

settingsForm.addEventListener("submit", (event) => {
  const submitter = /** @type {HTMLButtonElement | null} */ ((/** @type {SubmitEvent} */ (event)).submitter);
  if (submitter?.value !== "save") return;

  settings = {
    lbUrl: inputLbUrl.value.trim() || DEFAULT_LB_URL,
    voice: inputVoice.value || DEFAULT_VOICE,
    instructions: inputInstructions.value.trim() || DEFAULT_INSTRUCTIONS,
  };
  saveSettings(settings);

  // If we have a live session, push the new voice/instructions on the wire
  // without tearing the peer connection down.
  if (client && LIVE_STATES.has(currentState)) {
    client.updateSession({ voice: settings.voice, instructions: settings.instructions });
  }
});

restartBtn.addEventListener("click", async () => {
  if (!LIVE_STATES.has(currentState)) return;
  settings = {
    lbUrl: inputLbUrl.value.trim() || DEFAULT_LB_URL,
    voice: inputVoice.value || DEFAULT_VOICE,
    instructions: inputInstructions.value.trim() || DEFAULT_INSTRUCTIONS,
  };
  saveSettings(settings);
  settingsModal.close();
  try {
    await teardown();
    await doStart();
  } catch (err) {
    onFatalError(err);
  }
});

// ─── Click the orb ──────────────────────────────────────────────────────

circleBtn.addEventListener("click", async () => {
  try {
    if (currentState === "idle" || currentState === "error") {
      await doStart();
    }
  } catch (err) {
    onFatalError(err);
  }
});

micBtn.addEventListener("click", () => {
  if (!micStream) return;
  micMuted = !micMuted;
  for (const track of micStream.getAudioTracks()) {
    track.enabled = !micMuted;
  }
  micBtn.classList.toggle("muted", micMuted);
  micBtn.setAttribute("aria-label", micMuted ? "Unmute" : "Mute");
  micBtn.title = micMuted ? "Unmute" : "Mute";
});

stopBtn.addEventListener("click", async () => {
  await teardown();
});

// ─── High-level flow ────────────────────────────────────────────────────

async function doStart() {
  setState("connecting");
  setCaption("Asking for mic…", "muted");
  try {
    micStream = await navigator.mediaDevices.getUserMedia({
      audio: {
        echoCancellation: true,
        noiseSuppression: true,
        autoGainControl: true,
      },
    });
  } catch (err) {
    throw new Error(
      `Microphone access denied${err instanceof Error ? `: ${err.message}` : ""}`,
    );
  }

  const inputTrack = micStream.getAudioTracks()[0];
  if (!inputTrack) throw new Error("No microphone track available");

  const c = new S2sRealtimeClient({
    loadBalancerUrl: settings.lbUrl,
    voice: settings.voice,
    instructions: settings.instructions,
    inputTrack,
  });
  client = c;

  c.addEventListener("status", (e) => {
    const detail = /** @type {CustomEvent<{ status: string }>} */ (e).detail;
    onClientStatus(detail.status);
  });
  c.addEventListener("outputTrack", (e) => {
    const detail = /** @type {CustomEvent<{ track: MediaStreamTrack }>} */ (e).detail;
    attachOutputTrack(detail.track);
  });
  c.addEventListener("transcript", (e) => {
    const detail = /** @type {CustomEvent<{ role: "user" | "assistant"; text: string; partial: boolean }>} */ (
      e
    ).detail;
    console.log(`[transcript:${detail.role}${detail.partial ? "·delta" : ""}] ${detail.text}`);
  });
  c.addEventListener("error", (e) => {
    const detail = /** @type {CustomEvent<{ error: unknown }>} */ (e).detail;
    onFatalError(detail.error);
  });
  c.addEventListener("session", (e) => {
    const detail = /** @type {CustomEvent<{ info: { sessionId: string } }>} */ (e).detail;
    console.log("[s2s] session created:", detail.info.sessionId);
  });

  await c.connect();
}

/** @param {string} status */
function onClientStatus(status) {
  switch (status) {
    case "creating-session":
    case "connecting":
      setState("connecting");
      break;
    case "connected":
      setState("listening");
      break;
    case "user-speaking":
      setState("user-speaking");
      break;
    case "processing":
      setState("processing");
      break;
    case "ai-speaking":
      setState("ai-speaking");
      break;
    case "closed":
      // teardown() will move us to idle
      break;
    case "error":
      setState("error");
      break;
  }
}

/** @param {MediaStreamTrack} track */
function attachOutputTrack(track) {
  if (!outputSink) {
    outputSink = document.createElement("audio");
    outputSink.autoplay = true;
    outputSink.setAttribute("playsinline", "true");
    document.body.appendChild(outputSink);
  }
  const stream = new MediaStream([track]);
  outputSink.srcObject = stream;
  outputSink.play().catch((err) => {
    console.warn("[main] audio playback blocked:", err);
  });
}

async function teardown() {
  if (client) {
    try {
      await client.close();
    } catch (err) {
      console.warn("[main] error closing client:", err);
    }
    client = null;
  }
  if (micStream) {
    for (const track of micStream.getTracks()) track.stop();
    micStream = null;
  }
  if (outputSink) {
    outputSink.srcObject = null;
    outputSink.remove();
    outputSink = null;
  }
  micMuted = false;
  micBtn.classList.remove("muted");
  setState("idle");
}

/** @param {unknown} err */
function onFatalError(err) {
  console.error("[main] fatal:", err);
  setState("error");
  const message = err instanceof Error ? err.message : String(err);
  setCaption(truncateError(message), "error");
  // Best-effort cleanup; ignore errors thrown during cleanup itself.
  void teardown().catch(() => {
    setState("error");
    setCaption(truncateError(message), "error");
  });
}

/** @param {string} text */
function truncateError(text) {
  if (text.length <= 90) return text;
  return text.slice(0, 87) + "…";
}

// ─── Bootstrap ──────────────────────────────────────────────────────────

setState("idle");

// Strip the .booting class after one rAF so the orb's first paint is
// committed without a fade/scale flash.
requestAnimationFrame(() => {
  document.body.classList.remove("booting");
});