// @ts-check
/**
 * LFM3-Omni on WebGPU, off the main thread so the avatar keeps its frame rate.
 * One conversation lives here: the history, the prefix-cached session, and the
 * context budget. Everything the main thread needs comes back as a stream of events.
 *
 * main -> worker:
 *   { kind: "load", model, system }
 *   { kind: "respond", id, audio }   a user turn, 16 kHz Float32
 *   { kind: "respond", id, text }    a user turn typed in the chat box
 *   { kind: "abort", id }            barge-in: stop generating this turn
 *   { kind: "system", system }       new persona / instructions
 *   { kind: "reset" }                forget the conversation
 *
 * worker -> main:
 *   { kind: "progress", loaded?, total?, fraction?, status? }
 *   { kind: "ready", device, sampleRate }   { kind: "failed", name, message }
 *   { kind: "text", id, text }     { kind: "audio", id, samples }
 *   { kind: "finished", id, reason, usage, generated, error? }
 *
 * Voice: with Piper enabled (default, `?voice=` selects or disables with
 * "model") replies generate as text only and each sentence is spoken by a
 * native Spanish Piper voice, resampled into the model's audio rate.
 */

import {
  Lfm3OmniForConditionalGeneration as Model,
  ContextOverflowError,
  DeviceLostError,
} from "../vendor/Lfm3OmniForConditionalGeneration.js";

const LOAD_ATTEMPTS = 4; // Chromium occasionally returns no adapter on a cold GPU process, and networks drop requests
const LOAD_RETRY_MS = 1500;
// Her replies stay in history as text only: her speech costs ~16 tokens/s and left
// no room for what was said, so a 44-turn chat now fits and she still recalls it.
// Past the budget half of it goes, well under the session's 8192 tokens.
const SEARCH = new URL(location.href).searchParams;
const HISTORY_BUDGET_TOKENS = Number(SEARCH.get("historyBudget")) || 6000;

const PRE_SPEECH_SAMPLES = 1200; // keep 50 ms before her first sound
/** @type {any} */ let model = null;
/** @type {any} */ let session = null;
let system = "";
/** @type {any[]} */ let history = [];
let idle = Promise.resolve();
/** @type {AbortController | null} */ let rereading = null;
/** @type {Map<number, AbortController>} */
const running = new Map();
const simulateRtf = Number(SEARCH.get("simulateRtf")) || 0;
const holdLoadMs = Number(SEARCH.get("holdLoadMs")) || 0; // testing: stretch out the load

const post = (msg, transfer = []) => self.postMessage(msg, { transfer });

// Piper TTS: native Spanish voice. `?voice=model` keeps the model's own voice; any
// other value is a Piper voice id (e.g. es_ES-sharvard-medium) + `?speaker=` (F=1).
let ttsSession = null;
let ttsWarm = Promise.resolve();
const ttsVoice = SEARCH.get("voice") ?? "es_ES-sharvard-medium";
const ttsSpeaker = Number(SEARCH.get("speaker")) || 1;
if (ttsVoice !== "model") {
  post({ kind: "progress", status: "Preparando voz" });
  ttsWarm = (async () => {
    try {
      const lib = await import("../vendor/piper/piper-tts-web.js");
      ttsSession = await lib.TtsSession.create({ voiceId: ttsVoice, speakerId: ttsSpeaker, logger: (m) => console.log("tts:", m) });
      post({ kind: "progress", status: "Voz lista" });
    } catch (err) {
      console.warn("voice load failed:", err);
    }
  })();
  void ttsWarm;
}

/** Split off a sentence boundary at the front of `pendingText`. */
function takeSentence(buf) {
  const m = buf.match(/[\s\S]*?[.!?…](\s|$)/);
  return m ? m[0] : null;
}

function decodeWav(buf) {
  const dv = new DataView(buf);
  let off = 12, fmt = null, data = null;
  while (off < buf.byteLength - 8) {
    const id = String.fromCharCode(dv.getUint8(off), dv.getUint8(off + 1), dv.getUint8(off + 2), dv.getUint8(off + 3));
    const size = dv.getUint32(off + 4, true);
    if (id === "fmt ") fmt = { comp: dv.getUint16(off + 8, true), ch: dv.getUint16(off + 10, true), rate: dv.getUint32(off + 12, true), bits: dv.getUint16(off + 22, true) };
    else if (id === "data") data = { off: off + 8, size };
    off += 8 + size + (size % 2);
  }
  if (!fmt || !data) throw new Error("wav malformed");
  let pcm;
  if (fmt.bits === 16) {
    const n = Math.floor(data.size / 2), out = new Float32Array(n);
    for (let k = 0; k < n; k++) out[k] = dv.getInt16(data.off + k * 2, true) / 32768;
    pcm = out;
  } else if (fmt.bits === 32 && fmt.comp === 3) {
    const n = Math.floor(data.size / 4), out = new Float32Array(n);
    for (let k = 0; k < n; k++) out[k] = dv.getFloat32(data.off + k * 4, true);
    pcm = out;
  } else throw new Error(`unsupported wav bit depth ${fmt.bits}`);
  return { pcm, rate: fmt.rate };
}

function trimSilence(x, rate, padSec) {
  const th = 0.012;
  let a = 0, b = x.length - 1;
  while (a < x.length && Math.abs(x[a]) < th) a++;
  while (b > a && Math.abs(x[b]) < th) b--;
  const pad = Math.round(rate * padSec);
  return x.slice(Math.max(0, a - pad), Math.min(x.length, b + 1 + pad));
}

function resample(x, rateIn, rateOut) {
  if (!rateOut || rateIn === rateOut) return x;
  const ratio = rateIn / rateOut, n = Math.max(1, Math.round(x.length / ratio)), out = new Float32Array(n);
  for (let k = 0; k < n; k++) {
    const p = k * ratio, j = Math.floor(p);
    out[k] = x[j] + (x[Math.min(j + 1, x.length - 1)] - x[j]) * (p - j);
  }
  return out;
}

/** Synthesize one sentence with Piper and hand it to the main thread like model audio. */
async function speakAndPost(id, text) {
  if (!text.trim()) return 0;
  try {
    if (!ttsSession) {
      console.warn("no voice yet, staying silent for:", text);
      return 0;
    }
    const { pcm, rate } = decodeWav(await (await ttsSession.predict(text)).arrayBuffer());
    let samples = trimSilence(pcm, rate, 0.12);
    samples = resample(samples, rate, model.audio_sampling_rate);
    if (samples.length > 16) post({ kind: "audio", id, samples }, [samples.buffer]);
    return samples.length;
  } catch (err) {
    console.warn("tts synthesis failed:", err);
    return 0;
  }
}

function messages() {
  return [{ role: "system", content: system }, ...history];
}

/** Forget the oldest user turn and everything answering it. */
function dropOldestExchange() {
  const next = history.findIndex((m, i) => i > 0 && m.role === "user");
  if (next < 0) return false;
  history.splice(0, next);
  return true;
}

/** Past the budget half the history goes at once, so the next turn's one re-prefill buys many cheap turns. */
function compactHistory() {
  const exchanges = history.filter((m) => m.role === "user").length;
  for (let i = 0; i < Math.floor(exchanges / 2); i++) dropOldestExchange();
}

/** The reply as history keeps it: what she said, not the sound of it. */
function textOnly(message) {
  if (!Array.isArray(message.content)) return message;
  return { ...message, content: message.content.filter((c) => c.type === "text").map((c) => c.text).join("") };
}

/**
 * Settle the cache while she's idle, re-reading a trimmed history and rewinding
 * past her audio, so nobody waits for it on their next turn. Stops short of the last reply: the cache can only rewind to its end,
 * and this prompt's end must still be a prefix of the next one.
 */
async function reread(signal) {
  const stream = session.generate(messages().slice(0, -1), { generation_mode: "text", max_new_tokens: 1, signal });
  for await (const _ of stream);
  await stream.result;
}

async function load({ model: modelId, system: sys }) {
  system = sys;
  let lastErr;
  for (let attempt = 1; attempt <= LOAD_ATTEMPTS; attempt++) {
    try {
      if (!navigator.gpu) {
        throw Object.assign(new Error("This browser doesn't support WebGPU."), { name: "WebGPUUnavailable" });
      }
      post({ kind: "progress", status: "Checking your GPU" });
      const check = await Model.checkAvailability(modelId);
      if (!check.ok && !/No WebGPU adapter/i.test(check.reason ?? "")) {
        throw Object.assign(new Error(check.reason), { name: "WebGPUUnavailable" });
      }
      model = await Model.from_pretrained(modelId, {
        onProgress: (p) => {
          if (p.kind === "bytes" && p.total > 0) post({ kind: "progress", loaded: p.loaded, total: p.total });
          else if (p.status === "ready") post({ kind: "progress", fraction: 1 });
          else post({ kind: "progress", status: p.message ?? p.status });
        },
      });
      session = model.create_session();
      post({ kind: "progress", status: "Getting ready" });
      await session.prime(messages()).result;
      if (holdLoadMs) await new Promise((r) => setTimeout(r, holdLoadMs));
      post({ kind: "ready", device: model.deviceInfo(), sampleRate: model.audio_sampling_rate });
      return;
    } catch (err) {
      lastErr = err;
      console.warn(`model load attempt ${attempt} failed:`, err);
      const adapterGlitch = /No WebGPU adapter/i.test(String(err?.message));
      if (adapterGlitch && attempt === LOAD_ATTEMPTS) lastErr = Object.assign(new Error(err.message), { name: "WebGPUUnavailable" });
      // A multi-gigabyte download over a real connection will sometimes drop a request. That's worth another go, after a growing pause; what's already cached isn't fetched again.
      const networkGlitch = /failed to fetch|networkerror|network error|load failed|err_/i.test(String(err?.message));
      if ((!adapterGlitch && !networkGlitch) || attempt === LOAD_ATTEMPTS) break;
      await new Promise((r) => setTimeout(r, LOAD_RETRY_MS * 2 ** (attempt - 1)));
    }
  }
  post({ kind: "failed", name: lastErr?.name ?? "Error", message: lastErr?.message ?? String(lastErr) });
}

/** One user turn: spoken (audio) or typed (text). */
async function respond({ id, audio, text }) {
  if (!audio && !text) {
    post({ kind: "finished", id, reason: "noop" });
    return;
  }
  const controller = new AbortController();
  running.set(id, controller);
  const { signal } = controller;
  rereading?.abort();
  const prior = idle;
  /** @type {() => void} */ let release = () => {};
  idle = new Promise((r) => (release = r));
  let reason = "stop";
  let usage = null;
  let generated = 0;
  let error;
  try {
    await prior;
    signal.throwIfAborted();
    history.push(audio ? { role: "user", content: [{ type: "audio", audio }] } : { role: "user", content: text });

    for (;;) {
      const textMode = ttsVoice !== "model" && !audio;
      const stream = session.generate(
        messages(),
        textMode ? { generation_mode: "text", signal } : { generation_mode: "interleaved", signal },
      );
      let emitted = false;
      let pendingText = "";
      const started = performance.now();
      try {
        for await (const ev of stream) {
          if (ev.type === "text") {
            emitted = true;
            if (textMode) {
              pendingText += ev.text;
              for (let s; (s = takeSentence(pendingText)); ) {
                pendingText = pendingText.slice(s.length);
                const clean = s.trim();
                if (clean) {
                  post({ kind: "text", id, text: clean });
                  generated += (await speakAndPost(id, clean)) ?? 0;
                }
              }
            } else {
              post({ kind: "text", id, text: ev.text });
            }
          } else if (ev.type === "audio") {
            emitted = true;
            let samples = ev.audio.slice();
            if (generated === 0) {
              // Her replies open with ~0.8 s of silence. Don't make anyone wait for it.
              const start = samples.findIndex((v) => Math.abs(v) > 0.01);
              if (start < 0) continue;
              samples = samples.slice(Math.max(0, start - PRE_SPEECH_SAMPLES));
            }
            generated += samples.length;
            if (simulateRtf) {
              const due = started + (generated / model.audio_sampling_rate / simulateRtf) * 1000;
              await new Promise((r) => setTimeout(r, Math.max(0, due - performance.now())));
            }
            post({ kind: "audio", id, samples }, [samples.buffer]);
          }
        }
        if (textMode && pendingText.trim()) {
          const clean = pendingText.trim();
          post({ kind: "text", id, text: clean });
          generated += (await speakAndPost(id, clean)) ?? 0;
          pendingText = "";
        }
      } catch (err) {
        if (!emitted && err instanceof ContextOverflowError && dropOldestExchange()) continue;
        throw err;
      }
      const result = await stream.result;
      usage = result.usage;
      reason = result.finish_reason;
      history.push(textOnly(result.message));
      break;
    }
  } catch (err) {
    if (err?.name === "AbortError" || signal.aborted) {
      reason = "cancelled";
    } else {
      console.error(err);
      reason = "error";
      error = { name: err?.name, message: err?.message ?? String(err), fatal: err instanceof DeviceLostError };
    }
  } finally {
    running.delete(id);
    post({ kind: "finished", id, reason, usage, generated, error });
    if (usage && !signal.aborted) {
      if (session.count_tokens(messages()).prompt_tokens > HISTORY_BUDGET_TOKENS) compactHistory();
      rereading = new AbortController();
      await reread(rereading.signal).catch((err) => {
        if (err?.name !== "AbortError" && !rereading?.signal.aborted) console.warn("re-read failed:", err);
      });
      rereading = null;
    }
    release();
  }
}

self.onmessage = ({ data }) => {
  switch (data.kind) {
    case "load":
      void load(data);
      break;
    case "respond":
      void respond(data);
      break;
    case "abort":
      running.get(data.id)?.abort();
      break;
    case "system":
      system = data.system;
      break;
    case "reset":
      idle = idle.then(() => {
        history = [];
      });
      break;
  }
};