"""
Supertonic TTS Gradio Demo — concurrent chunk streaming with playback sync
Splits text into statements, synthesizes in a background thread, streams
audio chunk-by-chunk.  A JS timeupdate listener tracks audio.currentTime
to bold exactly the chunk currently playing in the browser.

Usage:
  python app.py
"""
import json
import logging
import queue
import re
import threading
from pathlib import Path

import gradio as gr
import numpy as np

from supertonic import TTS, AVAILABLE_LANGUAGES

logger = logging.getLogger(__name__)

CLONED_VOICE_JSON = Path(__file__).resolve().parent / "recorded_voice_supertonic-3.json"
CLONED_VOICE_LABEL = "Custom (recorded_voice_supertonic-3.json)"

print("Loading Supertonic TTS model...")
tts = TTS(auto_download=True)

builtin_voices = tts.voice_style_names  # M1-M5, F1-F5
all_voices = ([CLONED_VOICE_LABEL] + builtin_voices) if CLONED_VOICE_JSON.exists() else builtin_voices

_SENTENCE_SPLIT_RE = re.compile(r'(?<=[.!?;])\s+')
_SENTINEL = object()

# JavaScript injected once at page load.
# Watches for changes to the hidden timing component, then attaches a
# timeupdate listener that bolds whichever .chunk-text span is currently
# playing based on audio.currentTime.
_SYNC_JS = """
() => {
    function attachAudioSync(chunks) {
        // Find the streaming audio element — Gradio wraps it inside a
        // data-testid="audio" block; fall back to first <audio> on the page.
        const container = document.querySelector('[data-testid="audio"]') ||
                          document.querySelector('.audio-container');
        const audio = container ? container.querySelector('audio')
                                : document.querySelector('audio');
        if (!audio) {
            // Audio element not yet in DOM — retry shortly.
            setTimeout(() => attachAudioSync(chunks), 150);
            return;
        }

        // Remove any previous listener we registered.
        if (audio._ttsSyncHandler) {
            audio.removeEventListener('timeupdate', audio._ttsSyncHandler);
        }

        audio._ttsSyncHandler = () => {
            const t = audio.currentTime;
            document.querySelectorAll('.chunk-item').forEach((item, i) => {
                const textEl = item.querySelector('.chunk-text');
                if (!textEl || i >= chunks.length) return;
                const active = t >= chunks[i].start && t < chunks[i].end;
                textEl.style.fontWeight = active ? '700' : 'normal';
            });
        };

        audio.addEventListener('timeupdate', audio._ttsSyncHandler);
    }

    // Observe the hidden timing textbox for value changes.
    const observer = new MutationObserver(() => {
        const box = document.querySelector('#tts-timing-box textarea');
        if (!box || !box.value) return;
        let chunks;
        try { chunks = JSON.parse(box.value); } catch(e) { return; }
        if (chunks.length) attachAudioSync(chunks);
    });
    observer.observe(document.body, { childList: true, subtree: true, characterData: true });
}
"""


def split_into_statements(text: str) -> list[str]:
    parts = _SENTENCE_SPLIT_RE.split(text.strip())
    return [p.strip() for p in parts if p.strip()]


def _chunks_html(statements: list[str], playing_idx: int, generating_idx: int) -> str:
    rows = []
    for i, stmt in enumerate(statements):
        if i < playing_idx:
            icon, color = "✅", "#16a34a"
        elif i == playing_idx:
            icon, color = "▶️", "#2563eb"
        elif i == generating_idx:
            icon, color = "⏳", "#d97706"
        else:
            icon, color = "○", "#9ca3af"

        rows.append(
            f'<div class="chunk-item" style="display:flex;gap:8px;padding:3px 0;">'
            f'<span style="color:{color};min-width:18px;font-size:13px;">{icon}</span>'
            f'<span class="chunk-text" style="font-size:13px;color:#374151;line-height:1.4;">{stmt}</span>'
            f'</div>'
        )
    return '<div style="margin-top:8px;">' + "".join(rows) + "</div>"


def generate_audio(text: str, voice_name: str, language: str, speed: float, steps: int):
    text = (text or "").strip()
    if not text:
        yield gr.update(), "Please enter some text.", "", "[]"
        return

    try:
        style = (
            tts.get_voice_style_from_path(CLONED_VOICE_JSON)
            if voice_name == CLONED_VOICE_LABEL
            else tts.get_voice_style(voice_name)
        )

        statements = split_into_statements(text) or [text]
        n = len(statements)

        chunk_q: queue.Queue = queue.Queue(maxsize=2)

        def _worker():
            for i, stmt in enumerate(statements):
                try:
                    wav, _ = tts.synthesize(
                        stmt,
                        voice_style=style,
                        lang=language,
                        speed=speed,
                        total_steps=int(steps),
                    )
                    chunk_q.put((i, wav.squeeze()))
                except Exception as exc:
                    chunk_q.put((i, exc))
                    return
            chunk_q.put(_SENTINEL)

        worker = threading.Thread(target=_worker, daemon=True)
        worker.start()

        yield gr.update(), f"Generating chunk 1/{n}…", _chunks_html(statements, -1, 0), "[]"

        timing: list[dict] = []   # [{start, end}, …] — grows as chunks arrive
        cumulative = 0.0

        while True:
            item = chunk_q.get()
            if item is _SENTINEL:
                break

            i, payload = item
            if isinstance(payload, Exception):
                yield gr.update(), f"Generation failed on chunk {i + 1}: {payload}", "", json.dumps(timing)
                return

            wav = payload
            dur = len(wav) / tts.sample_rate
            timing.append({"start": round(cumulative, 4), "end": round(cumulative + dur, 4)})
            cumulative += dur

            next_gen = i + 1 if i + 1 < n else -1
            status = (
                f"Playing {i + 1}/{n}, generating {next_gen + 1}/{n}…"
                if next_gen >= 0
                else f"Playing {i + 1}/{n}"
            )

            yield (
                (tts.sample_rate, wav),
                status,
                _chunks_html(statements, i, next_gen),
                json.dumps(timing),   # JS picks this up via MutationObserver
            )

        worker.join()
        total_dur = cumulative
        yield (
            gr.update(),
            f"Done — {n} chunks, {total_dur:.1f}s total",
            _chunks_html(statements, n, -1),
            json.dumps(timing),
        )

    except Exception as e:
        logger.exception("Audio generation failed")
        yield gr.update(), f"Generation failed: {type(e).__name__}: {e}", "", "[]"


with gr.Blocks(title="Supertonic TTS Demo", js=_SYNC_JS) as demo:
    gr.Markdown(f"# Supertonic TTS Demo\n**Model:** `supertonic-3` · **Voices:** {len(all_voices)}")

    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="Text",
                placeholder="Enter text to synthesize…",
                lines=5,
            )

            with gr.Row():
                lang_input = gr.Dropdown(
                    label="Language",
                    choices=AVAILABLE_LANGUAGES,
                    value="en",
                    interactive=True,
                )
                voice_input = gr.Dropdown(
                    label="Voice Style",
                    choices=all_voices,
                    value=all_voices[0],
                    interactive=True,
                )

            with gr.Row():
                speed_input = gr.Slider(label="Speed", minimum=0.7, maximum=2.0, value=1.05, step=0.05)
                steps_input = gr.Slider(label="Quality Steps", minimum=1, maximum=20, value=8, step=1)

            play_btn = gr.Button("Play", variant="primary", size="lg")

        with gr.Column(scale=3):
            # streaming=True: each yielded chunk is appended to the playback buffer
            audio_out = gr.Audio(label="Generated Speech", type="numpy", autoplay=True, streaming=True)
            status_out = gr.Textbox(label="Status", lines=1, interactive=False)
            chunks_out = gr.HTML()
            # Hidden component — carries [{start, end}] timing per chunk to JS
            timing_out = gr.Textbox(visible=False, elem_id="tts-timing-box")

    play_btn.click(
        generate_audio,
        inputs=[text_input, voice_input, lang_input, speed_input, steps_input],
        outputs=[audio_out, status_out, chunks_out, timing_out],
    )

if __name__ == "__main__":
    demo.queue().launch()