import re

with open('static/js/audiobook.js', 'r') as f:
    content = f.read()

# 1. Change listener
content = content.replace(
    "$('reh-tb-audiobook')?.addEventListener('click', audiobookExport);",
    "$('reh-tb-audiobook')?.addEventListener('click', audiobookExportDialog);"
)

# 2. Rename audiobookExport to accept opts
content = content.replace(
    "async function audiobookExport() {",
    "async function audiobookExport(opts = {}) {"
)

# 3. Handle activeBuckets inside audiobookExport
content = content.replace(
    "const allIdx = buckets.flatMap(b => b.idx);",
    """let activeBuckets = buckets;
  if (opts.selectedChapterIndices) {
    activeBuckets = buckets.filter((_, i) => opts.selectedChapterIndices.includes(i));
  }
  const allIdx = activeBuckets.flatMap(b => b.idx);"""
)

# 4. Modify the loop that creates files
loop_find = """  for (let c = 0; c < buckets.length; c++) {
    const blobs = buckets[c].idx.map(i => wavClips.get(i)).filter(Boolean);"""

loop_replace = """  if (opts.format === 'single') {
    let masterBlobs = [];
    for (let c = 0; c < activeBuckets.length; c++) {
      const blobs = activeBuckets[c].idx.map(i => wavClips.get(i)).filter(Boolean);
      if (!blobs.length) continue;
      if (encMsg) encMsg.textContent = `Merging chapter ${c + 1} / ${activeBuckets.length}…`;
      let mergedWav = await mergeWavBlobs(blobs, pauses.paragraphMs);
      if (c > 0 && activeBuckets[c].title && (pauses.chapterMs > 0 || _abChapterSfxDataUrl())) {
        const ref = _parseWavBytes(new Uint8Array(await mergedWav.arrayBuffer())).fmt;
        const lead = [];
        const sfx = await _abChapterSfxWavBlob(ref);
        if (sfx) lead.push(sfx);
        if (pauses.chapterMs > 0) lead.push(_buildWavBlob(_silencePcmBytes(pauses.chapterMs, ref), ref));
        if (lead.length) mergedWav = await mergeWavBlobs([...lead, mergedWav]);
      }
      masterBlobs.push(mergedWav);
    }
    if (masterBlobs.length) {
      if (encMsg) encMsg.textContent = `Encoding full book to MP3…`;
      let finalWav = await mergeWavBlobs(masterBlobs, 0);
      let blob = finalWav;
      try {
        const encResp = await fetch('/api/audio/encode-mp3', { method: 'POST', body: finalWav });
        if (encResp.ok) blob = await encResp.blob();
      } catch (e) {}
      const ext = blob === finalWav ? 'wav' : 'mp3';
      const name = `${title}.${ext}`;
      if (typeof readerDownload === 'function') readerDownload(blob, name);
      try {
        let exportUrl = `/api/audiobook-export/${encodeURIComponent(bookForExport)}/${encodeURIComponent(name)}`;
        if (opts.folder) exportUrl += `?folder=${encodeURIComponent(opts.folder)}`;
        const r = await fetch(exportUrl, { method: 'POST', body: blob });
        if (r.ok) savedFiles.push({ name, url: exportUrl });
      } catch (e) {}
      files++;
    }
  } else {
  for (let c = 0; c < activeBuckets.length; c++) {
    const blobs = activeBuckets[c].idx.map(i => wavClips.get(i)).filter(Boolean);"""

content = content.replace(loop_find, loop_replace)

# 5. Fix references inside the else block
content = content.replace(
    "encMsg.textContent = `Merging chapter ${c + 1} / ${buckets.length}…`;",
    "encMsg.textContent = `Merging chapter ${c + 1} / ${activeBuckets.length}…`;"
)
content = content.replace(
    "encMsg.textContent = `Encoding chapter ${c + 1} / ${buckets.length}…`;",
    "encMsg.textContent = `Encoding chapter ${c + 1} / ${activeBuckets.length}…`;"
)
content = content.replace(
    "if (c > 0 && buckets[c].title && (pauses.chapterMs > 0 || _abChapterSfxDataUrl())) {",
    "if (c > 0 && activeBuckets[c].title && (pauses.chapterMs > 0 || _abChapterSfxDataUrl())) {"
)
content = content.replace(
    "const ch = buckets[c].title ? ' ' + readerSafeName(buckets[c].title) : '';\n    const name = (realChapters || buckets.length > 1) ? `${title} - ${String(c + 1).padStart(2, '0')}${ch}.${ext}` : `${title}.${ext}`;",
    "const ch = activeBuckets[c].title ? ' ' + readerSafeName(activeBuckets[c].title) : '';\n    const name = (realChapters || activeBuckets.length > 1) ? `${title} - ${String(c + 1).padStart(2, '0')}${ch}.${ext}` : `${title}.${ext}`;"
)
content = content.replace(
    "if (r.ok) savedFiles.push({ name, url: `/api/audiobook-export/${encodeURIComponent(bookForExport)}/${encodeURIComponent(name)}` });",
    "if (r.ok) savedFiles.push({ name, url: exportUrl });"
)
# add folder query param in else block
content = content.replace(
    "const r = await fetch(`/api/audiobook-export/${encodeURIComponent(bookForExport)}/${encodeURIComponent(name)}`, { method: 'POST', body: blob });",
    "let exportUrl = `/api/audiobook-export/${encodeURIComponent(bookForExport)}/${encodeURIComponent(name)}`;\n    if (opts.folder) exportUrl += `?folder=${encodeURIComponent(opts.folder)}`;\n    const r = await fetch(exportUrl, { method: 'POST', body: blob });"
)


# close the else block
close_else = """    files++;
    await new Promise(r => setTimeout(r, 400)); // breathing room between chunk encode/uploads
  }
  } // end else
"""
content = content.replace(
    "await new Promise(r => setTimeout(r, 400)); // breathing room between chunk encode/uploads\n  }",
    close_else
)

# 6. Add audiobookExportDialog
dialog_code = """
window.audiobookExportDialog = function() {
  if (_audiobook.running) return;
  if (!window.rehState || !(rehState.lines || []).length) { toast('Open a script in the rehearser first', 'error'); return; }
  if (!rehState.backend) { toast('Select a TTS backend in the rehearser first', 'error'); return; }
  
  const speakable = i => {
    const l = rehState.lines[i];
    if (!l || l.ignored || l.hidden) return false;
    if (l.type === 'dialog') { const c = rehState.cast[l.speaker]; return !!(c && c.voice && c.voice !== 'me'); }
    return !!(rehState.narratorVoice && (l.text || '').trim());
  };
  const buckets = [];
  let cur = null;
  rehState.lines.forEach((l, i) => {
    if (audiobookIsChapter(l)) { cur = { title: (typeof stripMarkdown === 'function' ? stripMarkdown(l.text) : l.text).trim().slice(0, 50), idx: [] }; buckets.push(cur); }
    if (speakable(i)) { if (!cur) { cur = { title: '', idx: [] }; buckets.push(cur); } cur.idx.push(i); }
  });

  const listEl = document.getElementById('ab-export-chapter-list');
  if (listEl) {
    listEl.innerHTML = '';
    const countEl = document.getElementById('ab-export-chapter-count');
    if (countEl) countEl.textContent = buckets.length;
    buckets.forEach((b, i) => {
      const lbl = document.createElement('label');
      lbl.style.display = 'block';
      lbl.style.marginBottom = '4px';
      const chk = document.createElement('input');
      chk.type = 'checkbox';
      chk.checked = true;
      chk.value = i;
      chk.className = 'ab-export-chapter-cb';
      lbl.appendChild(chk);
      lbl.appendChild(document.createTextNode(' Chapter ' + (i+1) + (b.title ? ': ' + b.title : '')));
      listEl.appendChild(lbl);
    });
  }

  const warn = document.getElementById('reh-tone-warn');
  const slot = document.getElementById('ab-export-tone-warn-slot');
  if (warn && slot) {
    warn.hidden = false;
    slot.appendChild(warn);
  }
  
  if (typeof _abInitPauseUI === 'function') _abInitPauseUI();

  const modal = document.getElementById('audiobook-export-modal');
  if (modal) modal.hidden = false;
};

document.addEventListener('DOMContentLoaded', () => {
  document.getElementById('ab-export-modal-close')?.addEventListener('click', () => document.getElementById('audiobook-export-modal').hidden = true);
  document.getElementById('ab-export-modal-cancel')?.addEventListener('click', () => document.getElementById('audiobook-export-modal').hidden = true);
  document.getElementById('ab-export-modal-start')?.addEventListener('click', () => {
    document.getElementById('audiobook-export-modal').hidden = true;
    const formatEl = document.querySelector('input[name="ab_export_format"]:checked');
    const format = formatEl ? formatEl.value : 'chapter';
    const folderEl = document.getElementById('ab-export-server-folder');
    const folder = folderEl ? folderEl.value : '';
    const checkboxes = Array.from(document.querySelectorAll('.ab-export-chapter-cb'));
    const selectedChapterIndices = checkboxes.filter(cb => cb.checked).map(cb => parseInt(cb.value, 10));
    audiobookExport({ format, selectedChapterIndices, folder });
  });
});
"""

content += dialog_code

with open('static/js/audiobook.js', 'w') as f:
    f.write(content)