app.js

const DEFAULT_BASE = 'https://huggingface.co/DownFlow/webasmr-f5tts/resolve/main';
const PRESETS = [
  ['助眠', '闭上眼睛，慢慢地呼吸，肩膀放松下来。'],
  ['耳语', '我就在你的耳边，轻轻说话，听得见吗？'],
  ['调情', '你呀，真是个笨蛋，让我看看你的表情。'],
  ['寸止', '再忍一下，数到三才允许，乖。'],
];

const $ = (id) => document.getElementById(id);
const els = {
  chip: $('cap-chip'), capText: $('cap-text'),
  text: $('text'), presets: $('presets'), go: $('go'), dl: $('dl'), hint: $('hint'),
  gain: $('gain'), gainOut: $('gain-out'),
  progress: $('progress'), bar: $('bar'), barFill: $('bar-fill'), stageLabel: $('stage-label'), stageDetail: $('stage-detail'),
  result: $('result'), player: $('player'), meta: $('meta'), error: $('error'),
  studioState: $('studio-state'), statSteps: $('stat-steps'),
  statLive: $('stat-live'), statTotal: $('stat-total'), statRt: $('stat-rt'), statRealtime: $('stat-realtime'),
};

const base = new URLSearchParams(location.search).get('base') || DEFAULT_BASE;
const worker = new Worker('/tts.worker.js', { type: 'module' });

let ready = false;
let readyPromise = null;
let resolveReady = null;
let synthesizing = false;
let pcm = null;
let pcmStats = null;
let wavUrl = null;
let tuning = null;
let nonce = 0;

for (const [name, phrase] of PRESETS) {
  const btn = document.createElement('button');
  btn.className = 'chip-btn';
  btn.textContent = name;
  btn.addEventListener('click', () => {
    els.text.value = phrase;
    for (const el of els.presets.children) el.classList.toggle('active', el === btn);
  });
  els.presets.appendChild(btn);
}

els.text.addEventListener('input', () => {
  for (const el of els.presets.children) el.classList.remove('active');
});

els.gain.addEventListener('input', () => {
  els.gainOut.textContent = `${els.gain.value}×`;
  if (pcm) refreshAudio();
});

function showError(message) {
  els.error.textContent = message;
  els.error.classList.add('on');
}

function clearError() {
  els.error.classList.remove('on');
}

function setState(label) {
  els.studioState.textContent = label;
}

function setBusy(busy) {
  els.go.disabled = busy;
}

function fmtMb(bytes) {
  return `${(bytes / 1e6).toFixed(0)} MB`;
}

function fmtMs(ms) {
  return ms >= 1000 ? `${(ms / 1000).toFixed(1)}s` : `${Math.round(ms)}ms`;
}

function statsFor(samples) {
  let sum = 0;
  let peak = 0;
  for (let i = 0; i < samples.length; i++) {
    const s = samples[i];
    sum += s * s;
    const a = Math.abs(s);
    if (a > peak) peak = a;
  }
  return { rms: Math.sqrt(sum / samples.length), peak };
}

function encodeWav(samples, gain) {
  const buffer = new ArrayBuffer(44 + samples.length * 2);
  const view = new DataView(buffer);
  const str = (offset, text) => { for (let i = 0; i < text.length; i++) view.setUint8(offset + i, text.charCodeAt(i)); };
  str(0, 'RIFF');
  view.setUint32(4, 36 + samples.length * 2, true);
  str(8, 'WAVE');
  str(12, 'fmt ');
  view.setUint32(16, 16, true);
  view.setUint16(20, 1, true);
  view.setUint16(22, 1, true);
  view.setUint32(24, 24000, true);
  view.setUint32(28, 48000, true);
  view.setUint16(32, 2, true);
  view.setUint16(34, 16, true);
  str(36, 'data');
  view.setUint32(40, samples.length * 2, true);
  for (let i = 0; i < samples.length; i++) {
    const v = Math.max(-1, Math.min(1, samples[i] * gain));
    view.setInt16(44 + i * 2, v * 32767, true);
  }
  return new Blob([buffer], { type: 'audio/wav' });
}

function refreshAudio() {
  if (!pcm) return;
  const gain = Number(els.gain.value);
  if (wavUrl) URL.revokeObjectURL(wavUrl);
  wavUrl = URL.createObjectURL(encodeWav(pcm, gain));
  els.player.src = wavUrl;
  els.dl.disabled = false;
  els.result.classList.add('on');
}

function setProgress(percent, { indeterminate = false } = {}) {
  els.progress.classList.add('on');
  els.bar.classList.toggle('indeterminate', indeterminate);
  els.barFill.style.width = indeterminate ? '30%' : `${Math.max(2, Math.min(100, percent))}%`;
}

function updateBar(stage) {
  if (stage.type === 'download' && stage.total) {
    setProgress((stage.loaded / stage.total) * 100);
    els.stageLabel.textContent = `下载模型 ${stage.label || ''}`;
    els.stageDetail.textContent = `${fmtMb(stage.loaded)} / ${fmtMb(stage.total)}`;
  } else if (stage.type === 'progress' && stage.stage === 'synth') {
    const step = stage.step ?? 0;
    const steps = stage.steps || tuning?.tts?.flowSteps || 32;
    if (step >= steps) {
      setProgress(100);
      els.stageLabel.textContent = '声码器';
      els.stageDetail.textContent = 'vocoding';
    } else {
      setProgress((step / steps) * 100);
      els.stageLabel.textContent = '流匹配合成';
      els.stageDetail.textContent = `step ${step}/${steps}`;
    }
  } else if (stage.type === 'progress' && stage.stage === 'warmup') {
    setProgress(100, { indeterminate: true });
    els.stageLabel.textContent = '预热内核';
    els.stageDetail.textContent = '';
  } else if (stage.type === 'progress' && stage.stage === 'session') {
    setProgress(100, { indeterminate: true });
    els.stageLabel.textContent = stage.done ? '会话就绪' : '初始化会话';
    els.stageDetail.textContent = stage.message || '';
  }
}

function applyCaps(caps) {
  if (caps?.webgpu) {
    const bits = [caps.vendor, caps.architecture].filter(Boolean).join(' · ') || 'webgpu';
    els.capText.textContent = `WebGPU · ${bits} · fp16`;
    els.chip.classList.add('ready');
  } else {
    els.capText.textContent = `WASM · int8 · ${caps?.numThreads ?? 1} threads`;
    els.chip.classList.add('ready');
  }
}

function setReady() {
  ready = true;
  setBusy(false);
  setState('READY');
  els.hint.textContent = '模型已缓存，直接输入文字即可生成。';
  setProgress(100);
  els.stageLabel.textContent = '就绪';
  els.stageDetail.textContent = '';
  setTimeout(() => els.progress.classList.remove('on'), 900);
}

function handleMessage(event) {
  const msg = event.data || {};
  switch (msg.type) {
    case 'progress':
      updateBar(msg);
      break;
    case 'fallback':
      console.warn('[tts]', msg.message);
      els.hint.textContent = msg.message;
      break;
    case 'caps':
      applyCaps(msg.caps);
      break;
    case 'ready':
      resolveReady?.(true);
      setReady();
      break;
    case 'audio': {
      pcm = msg.pcm instanceof Float32Array ? msg.pcm : new Float32Array(msg.pcm);
      pcmStats = statsFor(pcm);
      refreshAudio();
      break;
    }
    case 'done': {
      synthesizing = false;
      setBusy(false);
      setState('READY');
      const synthMs = msg.msTotal;
      const seconds = pcm ? pcm.length / 24000 : 0;
      els.statSteps.textContent = String(msg.steps ?? tuning?.tts?.flowSteps ?? 32);
      els.statLive.hidden = false;
      els.statRt.hidden = false;
      els.statTotal.textContent = fmtMs(synthMs);
      els.statRealtime.textContent = synthMs > 0 ? `${(seconds / (synthMs / 1000)).toFixed(2)}×` : '–';
      els.meta.innerHTML = [
        `时长 <b>${seconds.toFixed(2)}s</b>`,
        `合成 <b>${fmtMs(synthMs)}</b>`,
        `每步 <b>${fmtMs(synthMs / (msg.steps || 32))}</b>`,
        `RMS <b>${pcmStats ? pcmStats.rms.toFixed(4) : '–'}</b>`,
        `PEAK <b>${pcmStats ? pcmStats.peak.toFixed(3) : '–'}</b>`,
      ].map((t) => `<span>${t}</span>`).join('');
      if (msg.unknownChars?.length) els.hint.textContent = `未知字：${msg.unknownChars.join(' ')}`;
      els.player.play().catch(() => {});
      break;
    }
    case 'error':
      resolveReady?.(false);
      synthesizing = false;
      setBusy(false);
      setState('ERROR');
      els.chip.classList.add('error');
      showError(msg.message || 'unknown error');
      break;
    default:
      break;
  }
}

worker.onmessage = handleMessage;
worker.onerror = (err) => {
  resolveReady?.(false);
  setBusy(false);
  setState('ERROR');
  els.chip.classList.add('error');
  showError(err.message || 'worker failed to start');
};

async function fetchTuning() {
  try {
    const res = await fetch('/tuning.json');
    if (res.ok) {
      tuning = await res.json();
      const steps = tuning?.tts?.flowSteps;
      if (steps) els.statSteps.textContent = String(steps);
    }
  } catch {
    tuning = null;
  }
}

function ensureReady() {
  if (ready) return Promise.resolve(true);
  if (!readyPromise) {
    readyPromise = new Promise((resolve) => { resolveReady = resolve; });
    (async () => {
      setState('LOADING MODELS');
      setBusy(true);
      clearError();
      els.hint.textContent = '首次使用需下载模型，请在进度条中查看下载状态。';
      await fetchTuning();
      worker.postMessage({ type: 'init', base, __nonce: ++nonce });
    })();
  }
  return readyPromise;
}

els.go.addEventListener('click', async () => {
  if (synthesizing) return;
  const text = els.text.value.trim();
  if (!text) { showError('请输入要合成的文字。'); return; }
  clearError();
  const ok = await ensureReady();
  if (!ok) return;
  synthesizing = true;
  setBusy(true);
  setState('SYNTHESIZING');
  els.hint.textContent = '首次合成会预热内核，稍等片刻…';
  els.statLive.hidden = true;
  els.statRt.hidden = true;
  worker.postMessage({ type: 'synthesize', text, voiceId: 'xiaochun', __nonce: ++nonce });
});

els.dl.addEventListener('click', () => {
  if (!wavUrl) return;
  const a = document.createElement('a');
  a.href = wavUrl;
  a.download = `webasmr-${Date.now()}.wav`;
  a.click();
});

setState('IDLE');