"""ArchiveAI TTS — convert SRT/VTT subtitle files to spoken audio, and review
AI-drafted interpreter turns before they go into the pipeline's video output."""

import math
import os
import re
import tempfile
import threading
import time
from pathlib import Path

import gradio as gr

from tts import (
    synthesize_text,
    apply_glossary,
    load_glossary,
    lookup_term,
    save_term,
    remove_term,
    parse_turns_json,
    parse_vtt_cues,
    build_turn_segments,
    synthesize_turn_audio,
    is_mantra_only,
)

VOICE_ID = "7Oj8ZD5BDEDKYENop6zW"  # David-Test
MAX_SLOTS = 20
TURN_PAGE_SIZE = 6

MODELS = {
    "v3 (best quality)":      "eleven_v3",
    "Multilingual v2":        "eleven_multilingual_v2",
    "Flash v2.5 (fastest)":   "eleven_flash_v2_5",
}
DEFAULT_MODEL_LABEL = "Multilingual v2"  # matches the July Launch pipeline's actual workflow

# ---------------------------------------------------------------------------
# Turn-audio temp files — every "Generate" click in Turn Review writes an
# mp3 to disk for the Audio player to serve. Isolated into its own directory
# (rather than scattered across the shared OS temp dir) so a background
# sweep can safely clean up anything old without touching unrelated files.
# ---------------------------------------------------------------------------

TURN_AUDIO_DIR = Path(tempfile.gettempdir()) / "archiveai_tts_turn_audio"
TURN_AUDIO_MAX_AGE_S = 2 * 3600  # eventual cleanup window
TURN_AUDIO_SWEEP_INTERVAL_S = 30 * 60

TURN_AUDIO_DIR.mkdir(parents=True, exist_ok=True)


def _remove_file_quietly(path: str | None) -> None:
    if not path:
        return
    try:
        os.remove(path)
    except OSError:
        pass  # already gone, or never existed — nothing to clean up


def sweep_stale_turn_audio(max_age_s: float = TURN_AUDIO_MAX_AGE_S) -> int:
    """Delete turn-audio temp files older than max_age_s. Returns count removed."""
    now = time.time()
    removed = 0
    for entry in TURN_AUDIO_DIR.glob("*.mp3"):
        try:
            if now - entry.stat().st_mtime > max_age_s:
                entry.unlink()
                removed += 1
        except OSError:
            pass  # raced with another cleanup or the file vanished — fine
    return removed


def _start_turn_audio_sweeper() -> None:
    def _loop():
        while True:
            sweep_stale_turn_audio()
            time.sleep(TURN_AUDIO_SWEEP_INTERVAL_S)

    threading.Thread(target=_loop, daemon=True).start()


_start_turn_audio_sweeper()

# ---------------------------------------------------------------------------
# Subtitle parsers — return list of {"timestamp": str, "text": str}
# ---------------------------------------------------------------------------

_TIMESTAMP_RE = re.compile(
    r"^\d{2}:\d{2}:\d{2}[,\.]\d{3}\s*-->\s*\d{2}:\d{2}:\d{2}[,\.]\d{3}"
)


def _parse_srt(content: str) -> list[dict]:
    content = content.replace("\r\n", "\n").replace("\r", "\n")
    blocks = re.split(r"\n{2,}", content.strip())
    segments = []
    for block in blocks:
        lines = block.strip().splitlines()
        if len(lines) < 2:
            continue
        idx = 0
        if lines[idx].strip().isdigit():
            idx += 1
        if idx >= len(lines) or not _TIMESTAMP_RE.match(lines[idx].strip()):
            continue
        timestamp = lines[idx].strip()
        text = " ".join(line.strip() for line in lines[idx + 1:] if line.strip())
        if text:
            segments.append({"timestamp": timestamp, "text": text})
    return segments


def _parse_vtt(content: str) -> list[dict]:
    content = content.replace("\r\n", "\n").replace("\r", "\n")
    blocks = re.split(r"\n{2,}", content.strip())
    if not blocks or not blocks[0].strip().startswith("WEBVTT"):
        raise ValueError("Not a valid WebVTT file (missing WEBVTT header).")
    segments = []
    for block in blocks[1:]:
        lines = block.strip().splitlines()
        if len(lines) < 2:
            continue
        idx = 0
        if not _TIMESTAMP_RE.match(lines[idx].strip()):
            idx += 1
        if idx >= len(lines) or not _TIMESTAMP_RE.match(lines[idx].strip()):
            continue
        timestamp = lines[idx].strip()
        text = " ".join(line.strip() for line in lines[idx + 1:] if line.strip())
        if text:
            segments.append({"timestamp": timestamp, "text": text})
    return segments


def parse_subtitle_file(file_obj) -> list[dict]:
    path = file_obj if isinstance(file_obj, str) else file_obj.name
    with open(path, encoding="utf-8") as f:
        content = f.read()
    if Path(path).suffix.lower() == ".vtt":
        return _parse_vtt(content)
    return _parse_srt(content)


# ---------------------------------------------------------------------------
# State — Quick Synthesize tab
# ---------------------------------------------------------------------------

def _make_state() -> dict:
    return {"page": 0, "segments": []}


def _save_page_edits(state: dict, slot_texts: list[str]) -> dict:
    segments = state["segments"]
    start = state["page"] * MAX_SLOTS
    for i, text in enumerate(slot_texts):
        idx = start + i
        if idx >= len(segments):
            break
        segments[idx]["text"] = text or ""
    state["segments"] = segments
    return state


def _render_page(state: dict) -> tuple:
    segments = state["segments"]
    total = len(segments)
    page = state["page"]
    total_pages = max(1, math.ceil(total / MAX_SLOTS))
    state["page"] = min(page, total_pages - 1)
    start = state["page"] * MAX_SLOTS
    end = min(start + MAX_SLOTS, total)
    n = end - start

    if total == 0:
        nav = ""
    elif total_pages > 1:
        nav = f"<p style='text-align:center;margin:0'>Page {state['page'] + 1} of {total_pages} &middot; Segments {start + 1}&ndash;{end} of {total}</p>"
    else:
        nav = f"<p style='text-align:center;margin:0'>{total} segment{'s' if total != 1 else ''}</p>"

    group_updates, ts_updates, text_updates = [], [], []
    for i in range(MAX_SLOTS):
        if i < n:
            seg = segments[start + i]
            group_updates.append(gr.update(visible=True))
            ts_updates.append(gr.update(value=f"`{seg['timestamp']}`"))
            text_updates.append(gr.update(value=seg["text"]))
        else:
            group_updates.append(gr.update(visible=False))
            ts_updates.append(gr.update(value=""))
            text_updates.append(gr.update(value=""))

    return (
        state,
        gr.update(value=nav),
        gr.update(interactive=state["page"] > 0),
        gr.update(interactive=state["page"] < total_pages - 1),
        *group_updates,
        *ts_updates,
        *text_updates,
    )


# ---------------------------------------------------------------------------
# Handlers — Quick Synthesize tab
# ---------------------------------------------------------------------------

def _load_path(path: str):
    """Shared logic: given a local file path, return page_outputs + group/textbox updates."""
    ext = Path(path).suffix.lower()
    if ext == ".txt":
        with open(path, encoding="utf-8") as f:
            content = f.read()
        return (*_render_page(_make_state()), gr.update(visible=True), gr.update(value=content))
    try:
        segments = parse_subtitle_file(path)
    except Exception as e:
        raise gr.Error(f"Could not parse subtitle file: {e}")
    state = _make_state()
    state["segments"] = segments
    return (*_render_page(state), gr.update(visible=False), gr.update(value=""))


def handle_file_change(file_obj, state):
    if file_obj is None:
        return (*_render_page(_make_state()), gr.update(visible=False), gr.update(value=""))
    path = file_obj if isinstance(file_obj, str) else file_obj.name
    return _load_path(path)


def handle_drive_url(url: str, state):
    url = (url or "").strip()
    if not url:
        return (*_render_page(_make_state()), gr.update(visible=False), gr.update(value=""))
    try:
        from tts.utils.drive import download_from_drive
        local_path, _ = download_from_drive(url)
    except Exception as e:
        raise gr.Error(str(e))
    return _load_path(local_path)


def handle_prev(state, *slot_texts):
    state = _save_page_edits(state, list(slot_texts))
    state["page"] = max(0, state["page"] - 1)
    return _render_page(state)


def handle_next(state, *slot_texts):
    state = _save_page_edits(state, list(slot_texts))
    state["page"] += 1
    return _render_page(state)


def save_edits(state, *slot_texts):
    return _save_page_edits(state, list(slot_texts))


def handle_synthesize(state, plain_text, prose_speed, mantra_speed, mantra_mode, model_label):
    segments = state["segments"]
    if segments:
        raw_text = " ".join(s["text"] for s in segments if s["text"].strip())
    elif (plain_text or "").strip():
        raw_text = plain_text.strip()
    else:
        raise gr.Error("Upload a file first.")

    api_key = os.environ.get("ELEVENLABS_API_KEY", "")
    if not api_key:
        raise gr.Error("ELEVENLABS_API_KEY secret is not set on this Space.")

    model_id = MODELS.get(model_label, MODELS[DEFAULT_MODEL_LABEL])
    full_text = apply_glossary(raw_text)

    stop = threading.Event()
    result = [None]
    exc = [None]

    def _synth():
        try:
            result[0] = synthesize_text(
                full_text,
                voice_id=VOICE_ID,
                prose_speed=prose_speed,
                mantra_speed=mantra_speed,
                mantra_mode=mantra_mode,
                api_key=api_key,
                model_id=model_id,
            )
        except Exception as e:
            exc[0] = e

    thread = threading.Thread(target=_synth, daemon=True)
    thread.start()

    try:
        while thread.is_alive():
            thread.join(timeout=0.5)
            if thread.is_alive():
                yield gr.update()
    except GeneratorExit:
        stop.set()
        raise

    if exc[0]:
        raise gr.Error(str(exc[0]))
    if result[0] is None:
        raise gr.Error("No audio was produced.")
    yield result[0]


# ---------------------------------------------------------------------------
# State — Turn Review tab
# ---------------------------------------------------------------------------

def _make_turn_state() -> dict:
    return {"page": 0, "segments": []}


def _save_turn_page_edits(state: dict, slot_texts: list[str]) -> dict:
    segments = state["segments"]
    start = state["page"] * TURN_PAGE_SIZE
    for i, text in enumerate(slot_texts):
        idx = start + i
        if idx >= len(segments):
            break
        segments[idx]["text"] = text or ""
    return state


def _fmt_time(seconds: float) -> str:
    m, s = divmod(int(seconds), 60)
    h, m = divmod(m, 60)
    return f"{h:02d}:{m:02d}:{s:02d}"


def _render_turn_page(state: dict) -> tuple:
    segments = state["segments"]
    total = len(segments)
    total_pages = max(1, math.ceil(total / TURN_PAGE_SIZE))
    state["page"] = min(state["page"], total_pages - 1)
    start = state["page"] * TURN_PAGE_SIZE
    end = min(start + TURN_PAGE_SIZE, total)
    n = end - start

    if total == 0:
        nav = ""
    elif total_pages > 1:
        nav = f"<p style='text-align:center;margin:0'>Page {state['page'] + 1} of {total_pages} &middot; Turns {start + 1}&ndash;{end} of {total}</p>"
    else:
        nav = f"<p style='text-align:center;margin:0'>{total} turn{'s' if total != 1 else ''}</p>"

    group_updates, header_updates, text_updates, audio_updates = [], [], [], []
    for i in range(TURN_PAGE_SIZE):
        if i < n:
            seg = segments[start + i]
            group_updates.append(gr.update(visible=True))
            mantra_note = " &middot; **mantra-only, no interpreter audio needed**" if seg["is_mantra_only"] else ""
            header_updates.append(gr.update(
                value=f"**Turn {seg['index'] + 1}** &middot; `{_fmt_time(seg['start'])} – {_fmt_time(seg['end'])}`"
                      f"{mantra_note}\n\n*{seg['rationale']}*"
            ))
            text_updates.append(gr.update(value=seg["text"]))
            audio_updates.append(gr.update(value=None))
        else:
            group_updates.append(gr.update(visible=False))
            header_updates.append(gr.update(value=""))
            text_updates.append(gr.update(value=""))
            audio_updates.append(gr.update(value=None))

    return (
        state,
        gr.update(value=nav),
        gr.update(interactive=state["page"] > 0),
        gr.update(interactive=state["page"] < total_pages - 1),
        *group_updates,
        *header_updates,
        *text_updates,
        *audio_updates,
    )


# ---------------------------------------------------------------------------
# Handlers — Turn Review tab
# ---------------------------------------------------------------------------

def _cleanup_turn_state_audio(state: dict | None) -> None:
    """Delete every generated audio file tracked in state's segments.

    Called whenever a turn_state (or its segments) is about to be replaced —
    otherwise those temp files would only get caught by the periodic sweep.
    """
    if not state:
        return
    for seg in state.get("segments", []):
        _remove_file_quietly(seg.get("audio_path"))


def _load_turns_and_vtt(turns_path: str, vtt_path: str, prev_state: dict | None = None) -> dict:
    try:
        turns = parse_turns_json(turns_path)
    except Exception as e:
        raise gr.Error(f"Could not parse turns.json: {e}")
    try:
        cues = parse_vtt_cues(vtt_path)
    except Exception as e:
        raise gr.Error(f"Could not parse translated VTT: {e}")
    _cleanup_turn_state_audio(prev_state)
    state = _make_turn_state()
    state["segments"] = build_turn_segments(turns, cues)
    return state


def handle_turns_load(turns_file, vtt_file, prev_state=None):
    if turns_file is None or vtt_file is None:
        _cleanup_turn_state_audio(prev_state)
        return _render_turn_page(_make_turn_state())
    turns_path = turns_file if isinstance(turns_file, str) else turns_file.name
    vtt_path = vtt_file if isinstance(vtt_file, str) else vtt_file.name
    state = _load_turns_and_vtt(turns_path, vtt_path, prev_state)
    return _render_turn_page(state)


def handle_turns_drive_load(turns_url, vtt_url, prev_state=None):
    turns_url = (turns_url or "").strip()
    vtt_url = (vtt_url or "").strip()
    if not turns_url or not vtt_url:
        _cleanup_turn_state_audio(prev_state)
        return _render_turn_page(_make_turn_state())
    from tts.utils.drive import download_from_drive
    try:
        turns_path, _ = download_from_drive(turns_url)
        vtt_path, _ = download_from_drive(vtt_url)
    except Exception as e:
        raise gr.Error(str(e))
    state = _load_turns_and_vtt(turns_path, vtt_path, prev_state)
    return _render_turn_page(state)


def handle_turn_prev(state, *slot_texts):
    state = _save_turn_page_edits(state, list(slot_texts))
    state["page"] = max(0, state["page"] - 1)
    return _render_turn_page(state)


def handle_turn_next(state, *slot_texts):
    state = _save_turn_page_edits(state, list(slot_texts))
    state["page"] += 1
    return _render_turn_page(state)


def save_turn_edits(state, *slot_texts):
    return _save_turn_page_edits(state, list(slot_texts))


def make_handle_generate_turn(slot_idx: int):
    def _handle(state, text, model_label, speed, comma_pause, period_pause,
                stability, similarity_boost, style):
        global_idx = state["page"] * TURN_PAGE_SIZE + slot_idx
        segments = state["segments"]
        if global_idx >= len(segments):
            raise gr.Error("No turn loaded in this slot.")

        text = (text or "").strip()
        if not text:
            raise gr.Error("This turn has no text to synthesize.")

        segments[global_idx]["text"] = text
        segments[global_idx]["is_mantra_only"] = is_mantra_only(text)
        if segments[global_idx]["is_mantra_only"]:
            raise gr.Error("This turn is mantra-only — no interpreter audio needed, skipping.")

        api_key = os.environ.get("ELEVENLABS_API_KEY", "")
        if not api_key:
            raise gr.Error("ELEVENLABS_API_KEY secret is not set on this Space.")

        model_id = MODELS.get(model_label, MODELS[DEFAULT_MODEL_LABEL])
        full_text = apply_glossary(text)

        result = [None]
        exc = [None]

        def _synth():
            try:
                audio = synthesize_turn_audio(
                    full_text,
                    voice_id=VOICE_ID,
                    model_id=model_id,
                    api_key=api_key,
                    speed=speed,
                    comma_pause_s=comma_pause,
                    period_pause_s=period_pause,
                    stability=stability,
                    similarity_boost=similarity_boost,
                    style=style,
                )
                fd, path = tempfile.mkstemp(
                    suffix=f"_turn_{global_idx + 1:02d}.mp3", dir=TURN_AUDIO_DIR,
                )
                os.close(fd)
                audio.export(path, format="mp3", bitrate="128k")
                result[0] = path
            except Exception as e:
                exc[0] = e

        thread = threading.Thread(target=_synth, daemon=True)
        thread.start()
        thread.join()

        if exc[0]:
            raise gr.Error(str(exc[0]))

        # Regenerating this slot orphans its previous file — clean it up now
        # rather than waiting on the sweep, then remember the new one.
        _remove_file_quietly(segments[global_idx].get("audio_path"))
        segments[global_idx]["audio_path"] = result[0]
        return gr.update(value=result[0])

    return _handle


# ---------------------------------------------------------------------------
# Glossary handlers
# ---------------------------------------------------------------------------

def handle_lookup(word):
    return lookup_term(word)

def handle_save(word, pronunciation):
    return save_term(word, pronunciation)

def handle_remove(word):
    return remove_term(word)

def handle_show_glossary():
    d = load_glossary()
    if not d:
        return "_No entries saved._"
    return "\n".join(f"- **{w}** → {p}" for w, p in sorted(d.items()))


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="ArchiveAI TTS") as demo:
    gr.Markdown(
        "# ArchiveAI TTS\nConvert subtitle files to spoken audio, and review AI-drafted interpreter turns.\n\n"
        "**Pause syntax:** When using the Multilingual v2 model, pauses may be added by "
        "inserting `<break time=\"1.5s\" />` directly in the text (up to 3 seconds per tag). "
        "The v3 model does not support `<break>` tags — use bracketed audio tags instead, "
        "e.g. `[pause]`, `[short pause]`, `[long pause]`, or delivery tags like `[whispers]`, `[laughs]`."
    )

    app_state = gr.State(_make_state())
    turn_state = gr.State(_make_turn_state())

    with gr.Tabs():

        # ── Turn Review tab ───────────────────────────────────────────────
        with gr.Tab("Turn Review"):
            gr.Markdown(
                "Load the speaker-turn plan (`turns.json`) and the translated captions "
                "(`*_en_pe.vtt`) for a chunk. Each turn becomes an editable card — adjust "
                "the text, generate the interpreter audio, listen, and download once approved."
            )

            with gr.Row():
                turns_file_input = gr.File(label="Turns JSON", file_types=[".json"])
                vtt_file_input = gr.File(label="Translated VTT", file_types=[".vtt"])
            with gr.Row():
                turns_drive_input = gr.Textbox(
                    label="Or paste a Google Drive link to turns.json",
                    placeholder="https://drive.google.com/file/d/…",
                )
                vtt_drive_input = gr.Textbox(
                    label="Or paste a Google Drive link to the translated VTT",
                    placeholder="https://drive.google.com/file/d/…",
                )
            load_drive_btn = gr.Button("Load from Drive links")

            with gr.Row():
                turn_model_dropdown = gr.Dropdown(
                    choices=list(MODELS.keys()), value=DEFAULT_MODEL_LABEL, label="Model",
                )
                turn_speed = gr.Slider(0.5, 2.0, value=0.87, step=0.01, label="Speed")
                comma_pause_slider = gr.Slider(
                    0.05, 0.6, value=0.2, step=0.01,
                    label="Comma Pause (s)",
                    info="Medium-phrase tier — short phrases get half this, long clauses get 1.5x.",
                )
                period_pause_slider = gr.Slider(
                    0.1, 1.0, value=0.5, step=0.05, label="Period Pause (s)",
                )

            with gr.Accordion("Advanced voice settings (not recommended)", open=False):
                gr.Markdown(
                    "These control the underlying ElevenLabs voice model directly. "
                    "The defaults below are the values validated in production — "
                    "changing them is **not recommended** unless you're deliberately "
                    "A/B testing voice quality, since it can make the voice sound "
                    "less consistent with previously generated turns."
                )
                with gr.Row():
                    stability_slider = gr.Slider(
                        0.0, 1.0, value=0.75, step=0.01, label="Stability",
                    )
                    similarity_boost_slider = gr.Slider(
                        0.0, 1.0, value=0.75, step=0.01, label="Similarity Boost",
                    )
                    style_slider = gr.Slider(
                        0.0, 1.0, value=0.05, step=0.01, label="Style",
                    )

            with gr.Row():
                turn_prev_btn = gr.Button("◀ Previous", interactive=False, scale=1)
                turn_nav_label = gr.HTML("", scale=3)
                turn_next_btn = gr.Button("Next ▶", interactive=False, scale=1)

            turn_groups, turn_headers, turn_texts, turn_audios, turn_gen_btns = [], [], [], [], []
            for _ in range(TURN_PAGE_SIZE):
                with gr.Group(visible=False) as grp:
                    hdr = gr.Markdown("")
                    txt = gr.Textbox(label="", lines=3, show_label=False)
                    with gr.Row():
                        gen_btn = gr.Button("Generate")
                        aud = gr.Audio(label="Interpreter Audio", type="filepath")
                turn_groups.append(grp)
                turn_headers.append(hdr)
                turn_texts.append(txt)
                turn_audios.append(aud)
                turn_gen_btns.append(gen_btn)

            turn_page_outputs = [turn_state, turn_nav_label, turn_prev_btn, turn_next_btn,
                                 *turn_groups, *turn_headers, *turn_texts, *turn_audios]

            turns_file_input.change(
                handle_turns_load, inputs=[turns_file_input, vtt_file_input, turn_state], outputs=turn_page_outputs,
            )
            vtt_file_input.change(
                handle_turns_load, inputs=[turns_file_input, vtt_file_input, turn_state], outputs=turn_page_outputs,
            )
            load_drive_btn.click(
                handle_turns_drive_load,
                inputs=[turns_drive_input, vtt_drive_input, turn_state],
                outputs=turn_page_outputs,
            )
            turn_prev_btn.click(handle_turn_prev, inputs=[turn_state, *turn_texts], outputs=turn_page_outputs)
            turn_next_btn.click(handle_turn_next, inputs=[turn_state, *turn_texts], outputs=turn_page_outputs)

            for i, gen_btn in enumerate(turn_gen_btns):
                gen_btn.click(
                    save_turn_edits, inputs=[turn_state, *turn_texts], outputs=[turn_state], queue=False,
                ).then(
                    make_handle_generate_turn(i),
                    inputs=[turn_state, turn_texts[i], turn_model_dropdown, turn_speed,
                            comma_pause_slider, period_pause_slider,
                            stability_slider, similarity_boost_slider, style_slider],
                    outputs=[turn_audios[i]],
                )

        # ── Quick Synthesize tab ──────────────────────────────────────────
        with gr.Tab("Quick Synthesize"):
            gr.Markdown(
                "Ad-hoc synthesis of a full subtitle file or block of text — no turn "
                "review, no pause insertion. Use **Turn Review** for the production workflow."
            )
            file_input = gr.File(
                label="Subtitle File or Plain Text (SRT, VTT, TXT)",
                file_types=[".srt", ".vtt", ".txt"],
            )
            drive_url_input = gr.Textbox(
                label="Or paste a Google Drive link (Google Doc, SRT, VTT, TXT — must be publicly shared)",
                placeholder="https://drive.google.com/file/d/… or docs.google.com/document/…",
            )

            with gr.Row():
                with gr.Column(scale=1):
                    model_dropdown = gr.Dropdown(
                        choices=list(MODELS.keys()),
                        value=DEFAULT_MODEL_LABEL,
                        label="Model",
                    )
                    prose_speed = gr.Slider(
                        0.5, 2.0, value=1.0, step=0.05, label="Prose Speed"
                    )
                    mantra_mode_toggle = gr.Checkbox(
                        label="Mantra mode (separate speed for ALL-CAPS runs)",
                        value=True,
                    )
                    mantra_speed = gr.Slider(
                        0.5, 1.5, value=0.75, step=0.05, label="Mantra Speed",
                    )
                with gr.Column(scale=1):
                    synthesize_btn = gr.Button("Synthesize", variant="primary")
                    audio_output = gr.Audio(
                        label="Output Audio",
                        type="numpy",
                        format="wav",
                    )

            with gr.Group(visible=False) as plain_text_group:
                plain_text_box = gr.Textbox(
                    label="Text",
                    lines=15,
                    show_label=False,
                )

            with gr.Column() as editor_col:
                with gr.Row():
                    prev_btn  = gr.Button("◀ Previous", interactive=False, scale=1)
                    nav_label = gr.HTML("", scale=3)
                    next_btn  = gr.Button("Next ▶",     interactive=False, scale=1)

                slot_groups, slot_timestamps, slot_texts = [], [], []
                for _ in range(MAX_SLOTS):
                    with gr.Group(visible=False) as grp:
                        ts  = gr.Markdown("", elem_classes=["timestamp-chip"])
                        txt = gr.Textbox(label="", lines=2, show_label=False)
                    slot_groups.append(grp)
                    slot_timestamps.append(ts)
                    slot_texts.append(txt)

            mantra_mode_toggle.change(
                lambda enabled: gr.update(visible=enabled),
                inputs=[mantra_mode_toggle],
                outputs=[mantra_speed],
            )

            page_outputs = [app_state, nav_label, prev_btn, next_btn,
                            *slot_groups, *slot_timestamps, *slot_texts]

            file_input.change(
                handle_file_change,
                inputs=[file_input, app_state],
                outputs=[*page_outputs, plain_text_group, plain_text_box],
            )
            drive_url_input.submit(
                handle_drive_url,
                inputs=[drive_url_input, app_state],
                outputs=[*page_outputs, plain_text_group, plain_text_box],
            )
            prev_btn.click(handle_prev, inputs=[app_state, *slot_texts], outputs=page_outputs)
            next_btn.click(handle_next, inputs=[app_state, *slot_texts], outputs=page_outputs)

            synthesize_btn.click(
                save_edits,
                inputs=[app_state, *slot_texts],
                outputs=[app_state],
                queue=False,
            ).then(
                handle_synthesize,
                inputs=[app_state, plain_text_box, prose_speed, mantra_speed, mantra_mode_toggle, model_dropdown],
                outputs=[audio_output],
            )

        # ── Glossary tab ───────────────────────────────────────────────────
        with gr.Tab("Glossary"):
            gr.Markdown(
                "Teach the engine how to pronounce words it gets wrong. "
                "Use plain English phonetic spelling — no IPA needed.\n\n"
                "**Example:** `Garchen` → `gar-chen`, `Rinpoche` → `rin-po-cheh`\n\n"
                "Entries are stored in the `tts-glossary` HuggingFace dataset, so they "
                "persist across Space restarts and apply to all future synthesis runs."
            )

            with gr.Row():
                word_input = gr.Textbox(label="Word", placeholder="Garchen")
                pronunciation_input = gr.Textbox(
                    label="Phonetic Spelling", placeholder="gar-chen"
                )

            with gr.Row():
                lookup_btn = gr.Button("Look Up")
                save_btn   = gr.Button("Save",   variant="primary")
                remove_btn = gr.Button("Remove", variant="stop")

            status_output = gr.Textbox(label="Status", interactive=False)

            with gr.Accordion("Show All Entries", open=False) as dict_accordion:
                dict_display = gr.Markdown()

            lookup_btn.click(handle_lookup, inputs=[word_input], outputs=[pronunciation_input])
            save_btn.click(handle_save, inputs=[word_input, pronunciation_input], outputs=[status_output])
            remove_btn.click(handle_remove, inputs=[word_input], outputs=[pronunciation_input, status_output])
            dict_accordion.expand(handle_show_glossary, inputs=[], outputs=[dict_display])


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)