// A custom Gradio HTML stage. The DOM stays mounted while native Gradio events
// deliver PCM packets; each packet is scheduled once on one Web Audio clock.
const root = element.closest('#the-stage') || element;
const $ = name => element.querySelector(`[data-stage="${name}"]`);
const state = {run: null, seen: 0, sequence: 0, take: null, current: null, takes: [], busy: false, keeping: false, cancelled: new Set()};
const button = id => { const el = document.getElementById(id); return el?.matches('button') ? el : el?.querySelector('button'); };
const message = (text, error = false) => {
  $('message').hidden = !text; $('message').textContent = text || '';
  $('message').classList.toggle('error', error);
};
const player = {
  ctx: null, analyser: null, nodes: new Set(), nextTime: 0, firstTime: null, epoch: 0, finished: false,
  ensure() {
    if (!this.ctx) {
      const AC = window.AudioContext || window.webkitAudioContext;
      this.ctx = new AC({sampleRate: 44100});
      this.analyser = this.ctx.createAnalyser();
      this.analyser.fftSize = 1024; this.analyser.smoothingTimeConstant = .7;
      this.analyser.connect(this.ctx.destination);
      // The completed replay uses the same analyser as the live stream.
      this.ctx.createMediaElementSource($('audio')).connect(this.analyser);
    }
    this.ctx.resume().catch(() => message('Press play to enable audio in this browser.'));
  },
  reset() {
    this.epoch++; this.finished = false;
    for (const node of this.nodes) { try { node.stop(); } catch {} node.disconnect(); }
    this.nodes.clear(); this.nextTime = 0; this.firstTime = null;
    $('audio').pause();
  },
  push(pcm, rate) {
    if (!pcm.length) return;
    this.ensure();
    const buffer = this.ctx.createBuffer(1, pcm.length, rate), data = buffer.getChannelData(0);
    for (let i = 0; i < pcm.length; i++) data[i] = pcm[i] / 32768;
    const node = this.ctx.createBufferSource(); node.buffer = buffer; node.connect(this.analyser);
    const at = Math.max(this.nextTime, this.ctx.currentTime + .025);
    if (this.firstTime === null) this.firstTime = at;
    this.nextTime = at + buffer.duration;
    const epoch = this.epoch;
    this.nodes.add(node);
    node.onended = () => {
      this.nodes.delete(node); node.disconnect();
      if (epoch === this.epoch && this.finished && !this.nodes.size) this.revealReplay();
    };
    node.start(at);
  },
  played() { return this.firstTime === null ? 0 : Math.max(0, Math.min(this.ctx.currentTime, this.nextTime) - this.firstTime); },
  buffered() { return this.ctx ? Math.max(0, this.nextTime - this.ctx.currentTime) : 0; },
  revealReplay() {
    if (state.current?.blob && !state.busy) {
      $('audio').hidden = false;
      // Do not autoplay or replace the source when generation completes: all
      // scheduled PCM drains first, then the finished WAV is ready for replay.
      $('audio').currentTime = 0;
      metrics(state.current, false);
    }
    syncControls();
  },
};
function syncControls() {
  const studio = document.getElementById('studio');
  if (studio) {
    studio.dataset.stageBusy = String(state.busy);
    studio.dataset.stagePlaying = String(player.nodes.size > 0 || (!$('audio').paused && !$('audio').ended));
  }
  const perform = button('perform'), retake = button('retake');
  if (state.busy && perform && !perform.disabled) perform.disabled = true;
  const noRetake = state.busy || !state.current?.id || !!state.current?.partial;
  if (retake && retake.disabled !== noRetake) retake.disabled = noRetake;
  const keepDisabled = state.busy || state.keeping || !state.current?.id || !state.current?.blob || !!state.current?.partial;
  if ($('keep').disabled !== keepDisabled) $('keep').disabled = keepDisabled;
  const keepLabel = state.keeping ? 'Keeping…' : 'Keep this voice';
  if ($('keep').textContent !== keepLabel) $('keep').textContent = keepLabel;
}
function busy(value) {
  state.busy = value;
  if (!value && button('perform')) button('perform').disabled = false;
  syncControls();
}
function decode(b64, take) {
  const bin = atob(b64), bytes = new Uint8Array(take.carry.length + bin.length);
  bytes.set(take.carry);
  for (let i = 0; i < bin.length; i++) bytes[take.carry.length + i] = bin.charCodeAt(i);
  const even = bytes.length - bytes.length % 2;
  take.carry = bytes.slice(even);
  return new Int16Array(bytes.buffer, 0, even / 2);
}
function wav(parts, rate) {
  const bytes = parts.reduce((n, p) => n + p.byteLength, 0), head = new ArrayBuffer(44), v = new DataView(head);
  const text = (o, s) => { for (let i = 0; i < s.length; i++) v.setUint8(o + i, s.charCodeAt(i)); };
  text(0, 'RIFF'); v.setUint32(4, 36 + bytes, true); text(8, 'WAVE'); text(12, 'fmt ');
  v.setUint32(16, 16, true); v.setUint16(20, 1, true); v.setUint16(22, 1, true);
  v.setUint32(24, rate, true); v.setUint32(28, rate * 2, true); v.setUint16(32, 2, true); v.setUint16(34, 16, true);
  text(36, 'data'); v.setUint32(40, bytes, true);
  return new Blob([head, ...parts], {type: 'audio/wav'});
}
function metrics(take, live = false) {
  const box = $('metrics'); box.replaceChildren();
  const add = (label, value) => {
    const span = document.createElement('span'), bold = document.createElement('b');
    span.append(label + ' '); bold.textContent = value; span.append(bold); box.append(span);
  };
  const m = take.metrics;
  if (m.ttfa_ms != null) add('first audio', `${m.ttfa_ms} ms`);
  if (m.client_ms != null) add('you heard it', `${m.client_ms} ms`);
  if (m.speed != null) add('speed', `${m.speed}× real-time`);
  if (m.audio_s != null) add('length', `${m.audio_s.toFixed(1)} s`);
  if (live && take.parts.length) add('buffer', `${player.buffered().toFixed(1)} s`);
  if (m.trimmed_s >= .5) add('silence trimmed', `${m.trimmed_s.toFixed(1)} s`);
  if (m.finish_reason === 'length') add('note', 'hit the length cap');
  if (m.finish_reason === 'stall') add('note', 'stalled, cut short');
  if (take.note) add('note', take.note);
  if (!box.childNodes.length) box.textContent = state.busy ? 'Preparing the actor…' : 'Your performance will appear here.';
}
function show(take) {
  $('title').textContent = `Take ${take.n}` + (take.note ? ' · retake' : '');
  $('voice').textContent = take.voice;
  $('audio').hidden = !take.blob || player.nodes.size > 0;
  $('download').disabled = !take.blob;
  $('keep').hidden = take.mode === 'clone';
  root.classList.toggle('has-take', !!take.blob);
  metrics(take, !take.blob || player.nodes.size > 0);
  syncControls();
}
function finish(take, partial = false) {
  if (!take.parts.length) return;
  take.partial = partial;
  if (partial) take.id = null;
  take.blob = wav(take.parts, take.rate); take.url = URL.createObjectURL(take.blob);
  if (!state.takes.includes(take)) state.takes.unshift(take);
  state.current = take;
  $('audio').src = take.url; // Exactly one source assignment per completed take.
  player.finished = true;
  show(take); slate();
  if (!player.nodes.size) player.revealReplay();
}
function slate() {
  const box = $('slate'); box.replaceChildren();
  for (const take of state.takes) {
    const item = document.createElement('div'); item.className = 'stage-item' + (take === state.current ? ' current' : '');
    const number = document.createElement('span'); number.className = 'stage-number'; number.textContent = `#${take.n}`;
    const body = document.createElement('div'), title = document.createElement('div'), sub = document.createElement('div');
    title.className = 'stage-item-title'; title.textContent = take.note ? `“${take.note}”` : take.text.replace(/\s+/g, ' ');
    sub.className = 'stage-item-sub'; sub.textContent = `${take.voice} · ${(take.metrics.audio_s || 0).toFixed(1)} s` + (take.partial ? ' · stopped' : '');
    body.append(title, sub);
    const play = document.createElement('button'); play.type = 'button'; play.className = 'stage-mini'; play.textContent = '▶ Play';
    play.addEventListener('click', () => {
      if (state.busy) return;
      player.reset(); state.current = take; state.take = take;
      $('audio').src = take.url; $('audio').hidden = false; show(take); slate();
      player.ensure(); $('audio').play().catch(() => {});
      if (take.id && !take.partial) trigger('click', {action: 'select', take_id: take.id});
    });
    item.append(number, body, play); box.append(item);
  }
}
function consume(event) {
  if (!event || event.type === 'idle' || state.cancelled.has(event.run_id)) return;
  if (event.type === 'notice') { state.keeping = false; message(event.message, event.error); syncControls(); return; }
  if (event.type === 'selected') return;
  if (event.run_id && event.run_id !== state.run) {
    player.reset(); state.run = event.run_id; state.seen = 0; state.sequence = 0;
    state.take = {id: null, n: state.takes.length + 1, text: event.text || '', mode: event.mode, voice: event.voice,
                  note: event.note, parts: [], carry: new Uint8Array(0), rate: 44100, metrics: {}, started: performance.now()};
    message(''); $('audio').hidden = true; $('download').disabled = true;
    busy(true); show(state.take);
  }
  if (event.sequence && event.sequence <= state.sequence) return;
  state.sequence = event.sequence || state.sequence;
  const take = state.take;
  if (!take && event.type !== 'error') return;
  if (take) {
    if (event.sample_rate) take.rate = event.sample_rate;
    if (event.take_id) take.id = event.take_id;
    if (event.ttfa_ms != null) take.metrics.ttfa_ms = event.ttfa_ms;
    // A cumulative packet ledger also handles coalesced final browser updates.
    const packets = event.packets || [];
    while (state.seen < packets.length) {
      const pcm = decode(packets[state.seen++], take);
      if (pcm.length) { take.parts.push(pcm); player.push(pcm, take.rate); }
    }
    take.metrics.audio_s = take.parts.reduce((n, p) => n + p.length, 0) / take.rate;
  }
  if (event.type === 'queued') message(event.position <= 1 ? 'Next in line…' : `In line · ${event.position - 1} ahead`);
  else if (event.type === 'retrying') message(event.message);
  else if (event.type === 'audio') { message(''); metrics(take, true); }
  else if (event.type === 'done') {
    for (const key of ['ttfa_ms', 'audio_s', 'speed', 'total_ms', 'queue_ms', 'trimmed_s', 'finish_reason']) {
      if (event[key] != null) take.metrics[key] = event[key];
    }
    busy(false); finish(take);
    if (event.warning) message(event.warning);
  } else if (event.type === 'error') {
    busy(false); message(event.message, true);
    if (take?.parts.length && !take.blob) finish(take, true);
    else if (state.current) show(state.current);
  }
  syncControls();
}
$('download').addEventListener('click', () => {
  const take = state.current; if (!take?.url) return;
  const a = document.createElement('a'); a.href = take.url; a.download = `drama3-take-${take.n}.wav`;
  document.body.append(a); a.click(); a.remove();
});
$('keep').addEventListener('click', () => {
  if (state.busy || !state.current?.id || state.current.partial) return;
  state.keeping = true; syncControls(); trigger('click', {action: 'keep', take_id: state.current.id});
});
$('audio').addEventListener('play', () => { player.ensure(); syncControls(); });
$('audio').addEventListener('pause', syncControls);
$('audio').addEventListener('ended', syncControls);
function installNote() {
  const note = document.querySelector('#stage-note textarea');
  if (note && !note.dataset.stageEnter) {
    note.dataset.stageEnter = 'true';
    note.addEventListener('keydown', e => {
      if (e.key === 'Enter' && !e.shiftKey && !e.metaKey && !e.ctrlKey) { e.preventDefault(); button('retake')?.click(); }
    });
  }
  syncControls();
}
const observer = new MutationObserver(installNote);
observer.observe(document.getElementById('studio') || element, {childList: true, subtree: true, attributes: true, attributeFilter: ['disabled']});
installNote();
function draw() {
  if (!element.isConnected) { observer.disconnect(); player.reset(); player.ctx?.close(); return; }
  const canvas = $('viz'), w = canvas.clientWidth, h = canvas.clientHeight, dpr = window.devicePixelRatio || 1;
  if (w && h) {
    if (canvas.width !== Math.round(w*dpr) || canvas.height !== Math.round(h*dpr)) { canvas.width = Math.round(w*dpr); canvas.height = Math.round(h*dpr); }
    const g = canvas.getContext('2d'); g.setTransform(dpr,0,0,dpr,0,0); g.clearRect(0,0,w,h);
    const style = getComputedStyle(element), ink = style.getPropertyValue('--ink').trim(), border = style.getPropertyValue('--border').trim();
    const data = new Uint8Array(player.analyser?.frequencyBinCount || 0);
    player.analyser?.getByteFrequencyData(data);
    const live = player.nodes.size > 0 || (!$('audio').paused && !$('audio').ended);
    const bars = Math.max(24, Math.floor(w/9));
    for (let b=0;b<bars;b++) {
      const idx = Math.floor(Math.pow(b/bars,1.4)*data.length*.35), v = data.length ? data[idx]/255 : 0;
      const bh = Math.max(2,v*(h-10)); g.fillStyle = live && v>.02 ? ink : border;
      g.beginPath(); g.roundRect(b*w/bars+2,(h-bh)/2,w/bars-4,bh,2); g.fill();
    }
    const played = $('audio').hidden ? player.played() : $('audio').currentTime;
    canvas.setAttribute('aria-label', `Audio waveform · ${played.toFixed(1)} seconds played`);
    if (state.take && player.firstTime !== null && player.ctx.state === 'running' && player.ctx.currentTime >= player.firstTime && state.take.metrics.client_ms == null) {
      state.take.metrics.client_ms = Math.round(performance.now()-state.take.started); metrics(state.take, true);
    }
  }
  syncControls(); requestAnimationFrame(draw);
}
window.dramaStage = {
  unlock: () => player.ensure(),
  stop: () => {
    if (state.run) state.cancelled.add(state.run);
    player.reset(); busy(false);
    if (state.take?.parts.length && !state.take.blob) finish(state.take, true);
    else if (state.current) { $('audio').hidden = false; show(state.current); }
    message('Cut! Take stopped.');
  },
};
watch('value', () => consume(props.value));
consume(props.value);
requestAnimationFrame(draw);