"""
Higgs Audio v3 TTS — high-fidelity edition (Hugging Face Spaces build)
• Smart sentence-aware chunking for texts up to ~50 000 chars
• Retry logic per chunk (3 attempts with back-off)
• Asks the API for raw 24 kHz/16-bit PCM instead of MP3, so the model's
  output is never lossily compressed before you've even heard it. Falls
  back to MP3 automatically if the account/region doesn't support PCM.
• Per-chunk loudness matching so volume doesn't jump between API calls
• Edge-silence trimming so chunks don't carry long dead air at the seams
• Glitch-free joins: chunks are stitched as decoded audio via pydub
  (short adaptive crossfade), with real silence spliced in for pauses
• Lossless WAV/FLAC export by default (MP3 320 kbps offered as an
  option for smaller files)
• Progress bar so you can see long jobs moving
Get a free API key at https://boson.ai/workspace — paste it into the
"Boson API key" field in the UI below. The app never generates,
stores, or fetches a key on your behalf; you provide it directly to
the Boson AI API exactly the way the API requires (a single header:
Authorization: Bearer <key>).

── Changes in this revision ─────────────────────────────────────────────
1. Speaker name/audio/transcript boxes are now zipped together as
   (name, audio, transcript) triples at creation time and passed through
   run_generate() as triples, instead of being flattened into three
   parallel lists and re-sliced by position. A future edit to the input
   ordering can no longer silently pair the wrong audio with the wrong
   speaker name.
2. Dropped `stream=True` for PCM requests — the API's PCM response is
   a plain binary body, not an SSE/chunked stream, so treating it as one
   risked handing bytes_to_audiosegment a malformed buffer with no
   error raised. If Boson's docs specify streaming is required for PCM,
   flip STREAM_PCM back to True (see constant below).
3. toggle_speaker_rows() now clears the name/audio/transcript fields of
   any row that becomes hidden, so shrinking "Number of speakers" can no
   longer leave a stale, invisible speaker silently included in the
   roster.
4. Replaced tempfile.mktemp() (documented as insecure — a race window
   between name generation and file creation) with
   tempfile.NamedTemporaryFile(delete=False), which reserves the file
   atomically.
5. gr.Blocks now uses delete_cache=(3600, 3600) so generated audio files
   in the Gradio temp dir are garbage-collected instead of accumulating
   forever on a long-running Space.
6. Status text truncates the per-chunk error list to the first 10
   entries with a "…and N more" summary, so a job with many failures
   doesn't dump an unreadable wall of text into the status box.
7. If cloning mode is used without a reference transcript, the status
   message now includes a visible warning that clone quality may suffer,
   instead of silently sending ref_text=None.

── Long-form support (10+ minute scripts) ───────────────────────────────
8. generate() is no longer wrapped in @spaces.GPU. ZeroGPU's `duration`
   is a hard kill-switch — the previous duration=120 would silently
   terminate any job whose sequential chunk calls took longer than two
   minutes, which a 10-minute script (roughly 60-100 chunks) routinely
   would. Since this app never touches the GPU (it only calls Boson's
   HTTP API), a separate trivial `_gpu_keepalive()` function stays
   decorated with @spaces.GPU purely to satisfy the Space's hardware
   validation, and is called once per job with a short, safe duration —
   the actual generation work runs outside any GPU time limit.
9. Chunks are now fetched CONCURRENTLY (bounded thread pool, default 4
   workers, adjustable in the UI) instead of one HTTP request at a time.
   For a 10-minute script this cuts wall-clock generation time roughly
   in proportion to the worker count, since each request is I/O-bound
   (waiting on the API), not CPU-bound. The first chunk is still fetched
   alone first, both to surface auth/format errors immediately and to
   pin down whether PCM or the MP3 fallback is in use before fanning out
   the rest of the requests. Segment order and pause placement remain
   exactly as specified in the script — concurrency only affects the
   order requests are *sent*, never the order audio is *assembled* in.
10. A live "estimated length" label under the text box gives a rough
    spoken-duration estimate as you type/paste, so you can gauge whether
    a script will land near your 10-minute target before spending API
    calls on it.
"""

import gradio as gr
import requests, os, re, time, tempfile, statistics
from pydub import AudioSegment, effects
from pydub.silence import detect_leading_silence
from concurrent.futures import ThreadPoolExecutor, as_completed
import io

# ZeroGPU (the free-tier hardware on Spaces) requires at least one function
# decorated with @spaces.GPU to pass its startup check — even though this
# app does no local GPU work at all (it only calls the Boson API over
# HTTP). Unlike the previous revision, the actual generation loop is NOT
# decorated anymore (see _gpu_keepalive below) — only a trivial no-op call
# is, so long jobs are never subject to ZeroGPU's hard duration cutoff.
import spaces

# ── constants ─────────────────────────────────────────────────────────────────
BASE_URL    = "https://api.boson.ai/v1/audio/speech"
CHUNK_SIZE  = 900          # safe chars per API call (keeps tokens well under limit)
MAX_RETRIES = 3

# Concurrent chunk requests. Higher = faster wall-clock time for long
# scripts, but more simultaneous load on your Boson account's rate limit.
# 4 is a safe default; the UI exposes a slider from 1 (fully sequential,
# old behavior) to 8. call_api()'s existing 429 back-off still applies
# per-thread, so hitting a rate limit slows a thread down rather than
# failing the job outright.
DEFAULT_CONCURRENCY = 4
MAX_CONCURRENCY = 8

# Rough spoken-word rate used only for the UI's "estimated length" label.
# Actual speed varies by language, punctuation/pauses, and the "Slow
# speech" toggle — this is guidance, not a precise prediction.
WORDS_PER_MINUTE_ESTIMATE = 150

# Whether to request a streamed response for PCM. Boson's PCM output is a
# plain binary body in practice, not chunked/SSE, so this defaults to
# False. Flip to True only if you've confirmed the API requires it.
STREAM_PCM = False

# The hosted model's native output is 24 kHz / 16-bit / mono. Requesting
# response_format "pcm" gets that audio with zero lossy compression in
# between the model and your speakers — MP3 is only used as a fallback if
# PCM isn't available for some reason (e.g. an account/region restriction).
PCM_SAMPLE_RATE = 24000
PCM_SAMPLE_WIDTH = 2   # bytes (16-bit)
PCM_CHANNELS = 1       # mono

VOICES = ["default", "chloe", "eleanor", "jake", "marcus", "nora", "oliver"]
EMOTIONS = [
    "none", "elation", "enthusiasm", "amusement", "sadness", "anger",
    "fear", "awe", "surprise", "disgust", "contemplation", "pride",
    "contentment", "affection", "relief", "confusion", "longing", "bitterness",
]
STYLES = ["none", "whisper", "shout"]

OUTPUT_FORMATS = [
    "WAV — lossless (recommended)",
    "FLAC — lossless, smaller file",
    "MP3 320 kbps — smallest file",
]

MAX_SPEAKERS = 6
MAX_ERRORS_SHOWN = 10  # cap on per-chunk error lines shown in the status box
# Matches a [SpeakerName] tag at the start of a line/segment. Names can
# contain letters (incl. non-Latin scripts), numbers, spaces, underscores,
# and hyphens — basically anything except the closing bracket.
_SPEAKER_TAG_RE = re.compile(r'\[([^\[\]]+)\]')

# Pause tokens (kept ONLY for backward-compat / display in the textbox).
# They are stripped out before being sent to the model — see
# split_text_into_pieces() below. Pauses are instead spliced in as real
# silence with pydub, which is 100% deterministic and doesn't depend on
# the model choosing to honor a prosody token.
PAUSE_SHORT = "<|prosody:pause|>"        # ≈ 400–700 ms   → we splice 550 ms
PAUSE_LONG  = "<|prosody:long_pause|>"   # ≈ 700–1500 ms  → we splice 1100 ms

PAUSE_DURATIONS_MS = {
    PAUSE_SHORT: 550,
    PAUSE_LONG: 1100,
}
# Matches either pause tag, used to split text around them
_PAUSE_SPLIT_RE = re.compile(
    r'(' + re.escape(PAUSE_SHORT) + r'|' + re.escape(PAUSE_LONG) + r')'
)

# ── audio-quality tuning knobs ──────────────────────────────────────────────
SILENCE_THRESH_DBFS  = -45     # what counts as "silence" when trimming edges
MAX_EDGE_SILENCE_MS  = 120     # keep up to this much natural lead-in/out per chunk
MIN_LEN_TO_TRIM_MS   = 2 * MAX_EDGE_SILENCE_MS  # don't bother trimming tiny chunks
MAX_LOUDNESS_GAIN_DB = 12      # clamp how hard we'll push a quiet/loud chunk
FINAL_PEAK_HEADROOM_DB = 1.0   # final mix peaks at -1 dBFS (safety margin, no clipping)
DEFAULT_CROSSFADE_MS = 30      # short, adaptive crossfade between spoken chunks


# ── helpers ───────────────────────────────────────────────────────────────────

def split_text_into_pieces(text: str) -> list[dict]:
    """
    Split text on pause tokens into an ordered list of pieces:
      {"type": "text",   "content": "..."}
      {"type": "pause",  "ms": 550}
    Each "text" piece still needs further sentence/length chunking before
    being sent to the API (see expand_text_piece_to_chunks below) — this
    function's only job is to pull the pauses out as their own atomic
    units so they can never be mangled, dropped, or left to the model's
    discretion.
    """
    pieces = []
    parts = _PAUSE_SPLIT_RE.split(text)
    for part in parts:
        if part in PAUSE_DURATIONS_MS:
            pieces.append({"type": "pause", "ms": PAUSE_DURATIONS_MS[part]})
        elif part.strip():
            pieces.append({"type": "text", "content": part.strip()})
    return pieces


def split_text_by_speaker(text: str, known_speakers: list[str]) -> list[tuple[str, str]]:
    """
    Parse a multi-speaker script tagged like:
        [Maya]
        Hello there, how are you?
        [Omar]
        I'm doing well, thanks for asking.
        [Maya]
        Glad to hear it.
    Returns an ordered list of (speaker_name, text_segment) tuples. Text
    before the first tag (if any) is assigned to the first known speaker
    as a fallback, so a script that forgets the very first tag still
    produces audio instead of silently losing that text.
    known_speakers is the roster's case-sensitive name list — a tag that
    doesn't match any known speaker is treated as plain spoken text
    belonging to whichever speaker is currently active (most likely a
    typo'd tag), so nothing gets silently dropped.
    """
    if not known_speakers:
        return []

    matches = list(_SPEAKER_TAG_RE.finditer(text))
    if not matches:
        stripped = text.strip()
        return [(known_speakers[0], stripped)] if stripped else []

    segments: list[tuple[str, str]] = []
    current_speaker = known_speakers[0]

    # Text before the first tag, if any, belongs to the first speaker.
    lead = text[: matches[0].start()].strip()
    if lead:
        segments.append((current_speaker, lead))

    for i, m in enumerate(matches):
        tag_name = m.group(1).strip()
        seg_start = m.end()
        seg_end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
        seg_text = text[seg_start:seg_end].strip()

        if tag_name in known_speakers:
            current_speaker = tag_name
        else:
            # Unrecognized tag — don't silently drop it. Put the
            # bracketed text back as part of the spoken line for the
            # currently active speaker (likely a typo'd speaker name).
            seg_text = f"[{tag_name}] {seg_text}".strip()

        if seg_text:
            segments.append((current_speaker, seg_text))

    return segments


def split_into_sentences(text: str) -> list[str]:
    """
    Split on sentence-ending punctuation while keeping the delimiter attached.
    Falls back to hard-cutting at CHUNK_SIZE if a single sentence is huge.
    NOTE: by the time text reaches this function, pause tokens have
    already been extracted out by split_text_into_pieces(), so there's no
    longer any risk of a hard-split landing inside a pause tag.
    """
    # Include Urdu/Arabic sentence-ending punctuation (۔ ؟ and Arabic ؛)
    # alongside Latin .!? so non-Latin scripts get split into real sentences
    # instead of being treated as one giant blob.
    raw = re.split(r'(?<=[.!?۔؟؛])\s+', text.strip())
    sentences = []
    for s in raw:
        s = s.strip()
        if not s:
            continue
        # If a single sentence is still too long, hard-split it at word
        # boundaries.
        while len(s) > CHUNK_SIZE:
            cut = s.rfind(' ', 0, CHUNK_SIZE)
            if cut == -1:
                cut = CHUNK_SIZE
            sentences.append(s[:cut].strip())
            s = s[cut:].strip()
        if s:
            sentences.append(s)
    return sentences


def pack_chunks(sentences: list[str], max_len: int = CHUNK_SIZE) -> list[str]:
    """Greedily merge sentences into chunks that stay under max_len chars."""
    chunks, current = [], ""
    for s in sentences:
        candidate = (current + " " + s).strip() if current else s
        if len(candidate) <= max_len:
            current = candidate
        else:
            if current:
                chunks.append(current)
            current = s
    if current:
        chunks.append(current)
    return chunks


def build_generation_plan(text: str, known_speakers: list[str] | None = None) -> list[dict]:
    """
    Turn raw user text into an ordered, flat list of work items ready for
    generation:
      {"type": "text",  "content": "<chunk under CHUNK_SIZE chars>", "speaker": "Maya" or None}
      {"type": "pause", "ms": 550 or 1100}
    If known_speakers is given (multi-speaker mode), the text is first
    split on [SpeakerName] tags so each chunk knows which cloned voice to
    use. Pauses are pulled out next (atomically, can never be split or
    dropped, and work the same whether or not the text is single- or
    multi-speaker), then each remaining text segment is independently run
    through sentence-aware chunking. This guarantees a pause always sits
    exactly where the user placed it relative to surrounding chunks, and
    that a chunk never accidentally merges lines from two different
    speakers together.
    """
    plan: list[dict] = []

    if known_speakers:
        speaker_segments = split_text_by_speaker(text, known_speakers)
    else:
        speaker_segments = [(None, text)]

    for speaker, seg_text in speaker_segments:
        for piece in split_text_into_pieces(seg_text):
            if piece["type"] == "pause":
                plan.append(piece)
            else:
                sentences = split_into_sentences(piece["content"])
                for chunk in pack_chunks(sentences, max_len=CHUNK_SIZE):
                    chunk = chunk.strip()
                    if chunk:
                        plan.append({
                            "type": "text",
                            "content": chunk,
                            "speaker": speaker,
                        })
    return plan


def build_tagged_text(text: str, emotion: str, style: str, slow: bool) -> str:
    tags = ""
    if emotion != "none":
        tags += f"<|emotion:{emotion}|>"
    if style != "none":
        tags += f"<|style:{style}|>"
    if slow:
        tags += "<|prosody:speed_slow|>"
    return tags + text


def insert_token_at_cursor(text: str, cursor_pos, token: str):
    """
    Insert a token at the current cursor position in the textbox.
    Gradio Textbox doesn't expose live cursor position via a normal event,
    so we fall back to appending at the end with a leading space — this
    keeps the button reliable across Gradio versions while still letting
    users place tokens by editing the text afterwards.
    """
    text = text or ""
    if text and not text.endswith((" ", "\n")):
        text += " "
    text += token + " "
    return text


@spaces.GPU(duration=10)
def _gpu_keepalive():
    """
    Trivial no-op used only to satisfy ZeroGPU's requirement that a Space
    have at least one @spaces.GPU-decorated function it actually calls.
    This app does no local GPU work — it only calls the Boson HTTP API —
    so the real generation loop deliberately does NOT carry this
    decorator (see the header notes). Keeping the duration tiny (10s)
    means this call returns almost instantly and never meaningfully
    limits how long a job can run.
    """
    return True


def estimate_spoken_duration(text: str) -> str:
    """
    Rough "how long will this take to speak" estimate for the UI, based
    on word count at WORDS_PER_MINUTE_ESTIMATE. Ignores pause tokens'
    exact timing and language-specific speech rate differences — it's a
    ballpark so you can see roughly where a script lands (e.g. relative
    to a 10-minute target) before spending API calls on it.
    """
    text = (text or "").strip()
    if not text:
        return "≈ 0:00 estimated"
    # Strip prosody tags before counting words so they don't inflate the count.
    visible = _PAUSE_SPLIT_RE.sub(" ", text)
    visible = re.sub(r'<\|[^|]*\|>', ' ', visible)
    words = len(visible.split())
    total_seconds = int(round((words / WORDS_PER_MINUTE_ESTIMATE) * 60))
    minutes, seconds = divmod(total_seconds, 60)
    return f"≈ {minutes}:{seconds:02d} estimated at ~{WORDS_PER_MINUTE_ESTIMATE} wpm (varies by language & pauses)"


class UnsupportedFormatError(RuntimeError):
    """Raised when the API rejects the requested response_format."""
    pass


def call_api(headers: dict, payload: dict) -> bytes:
    """Call the API with retry + exponential back-off. Returns raw audio bytes."""
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            r = requests.post(BASE_URL, headers=headers, json=payload, timeout=60)
            if r.ok:
                return r.content
            # 429 → wait and retry
            if r.status_code == 429:
                wait = 2 ** attempt
                time.sleep(wait)
                continue
            # Other errors → raise immediately with context
            try:
                detail = r.json()
            except Exception:
                detail = r.text[:400]
            detail_str = str(detail)
            if r.status_code in (400, 422) and "format" in detail_str.lower():
                raise UnsupportedFormatError(f"HTTP {r.status_code}: {detail}")
            raise RuntimeError(f"HTTP {r.status_code}: {detail}")
        except requests.exceptions.Timeout:
            if attempt == MAX_RETRIES:
                raise RuntimeError("Request timed out after 60 s")
            time.sleep(2 ** attempt)
    raise RuntimeError(f"Failed after {MAX_RETRIES} retries")


def bytes_to_audiosegment(raw: bytes, response_format: str) -> AudioSegment:
    """Decode API audio bytes to an AudioSegment, honoring the format used."""
    if response_format == "pcm":
        # Raw 16-bit/24kHz/mono PCM — construct directly, no ffmpeg/lossy
        # decode step involved at all.
        return AudioSegment(
            data=raw,
            sample_width=PCM_SAMPLE_WIDTH,
            frame_rate=PCM_SAMPLE_RATE,
            channels=PCM_CHANNELS,
        )
    return AudioSegment.from_file(io.BytesIO(raw), format="mp3")


def trim_edge_silence(seg: AudioSegment,
                       silence_thresh: int = SILENCE_THRESH_DBFS,
                       max_keep_ms: int = MAX_EDGE_SILENCE_MS) -> AudioSegment:
    """
    Cap the amount of near-silence at the start/end of a generated chunk
    to at most max_keep_ms, instead of removing it entirely — this keeps
    a natural breath/lead-in but stops occasional long dead-air the model
    sometimes appends from creating awkward gaps once chunks are joined.
    Skipped for very short segments to avoid edge cases.
    """
    if len(seg) < MIN_LEN_TO_TRIM_MS or seg.rms == 0:
        return seg

    lead_sil = detect_leading_silence(seg, silence_threshold=silence_thresh)
    trim_front = max(0, lead_sil - max_keep_ms)
    if trim_front:
        seg = seg[trim_front:]

    trail_sil = detect_leading_silence(seg.reverse(), silence_threshold=silence_thresh)
    trim_back = max(0, trail_sil - max_keep_ms)
    if trim_back and trim_back < len(seg):
        seg = seg[:-trim_back]

    return seg


def match_loudness(segments: list[AudioSegment]) -> list[AudioSegment]:
    """
    Bring all spoken (non-pause) chunks to a common loudness so volume
    doesn't visibly jump between separate API calls. Uses the median dBFS
    across chunks as the target (robust to one unusually loud/quiet
    chunk), and clamps the applied gain so a near-silent chunk doesn't get
    pushed into audible noise. Pause/silence segments are left untouched.
    """
    voiced_dbfs = [
        s.dBFS for s in segments
        if s.rms > 0 and s.dBFS != float("-inf")
    ]
    if not voiced_dbfs:
        return segments

    target = statistics.median(voiced_dbfs)
    matched = []
    for s in segments:
        if s.rms == 0 or s.dBFS == float("-inf"):
            matched.append(s)
            continue
        gain = target - s.dBFS
        gain = max(-MAX_LOUDNESS_GAIN_DB, min(MAX_LOUDNESS_GAIN_DB, gain))
        matched.append(s.apply_gain(gain))
    return matched


def concatenate_segments(segments: list[AudioSegment],
                          crossfade_ms: int = DEFAULT_CROSSFADE_MS) -> AudioSegment:
    """
    Join segments with a short, adaptive crossfade to avoid click/pop
    glitches at the seams — except around deliberately-inserted silence
    (pause) segments, where crossfading would eat into the intended pause
    duration and shorten/blur it. Plain concatenation (+) is used on those
    joins instead, which is silence-safe (no glitch risk crossfading
    into/out of pure silence anyway).
    """
    if not segments:
        raise RuntimeError("No audio segments to join.")

    def is_silence(seg: AudioSegment) -> bool:
        return seg.rms == 0

    result = segments[0]
    for seg in segments[1:]:
        if is_silence(result) or is_silence(seg):
            result = result + seg
        else:
            # Never crossfade longer than either side actually has, or
            # pydub will raise/clip oddly on very short chunks.
            safe_fade = max(0, min(crossfade_ms, len(result) - 1, len(seg) - 1))
            if safe_fade <= 0:
                result = result + seg
            else:
                result = result.append(seg, crossfade=safe_fade)
    return result


# ── main generation function ──────────────────────────────────────────────────

def generate(
    api_key, mode, text, emotion, style, slow,
    preset_voice, ref_audio_path, ref_transcript,
    speaker_triples,  # list of (name, audio_path, transcript) — see FIX #1
    output_format, enhance_audio, concurrency,
    progress=gr.Progress(track_tqdm=False),
):
    # Satisfy ZeroGPU's "must actually call a @spaces.GPU function" check
    # without subjecting the (potentially many-minute) generation loop
    # below to that decorator's hard duration cutoff. See header notes.
    _gpu_keepalive()

    if not api_key.strip():
        return None, "⚠️  Please enter your Boson API key."
    if not text.strip():
        return None, "⚠️  Please enter some text."
    if mode == "Clone my voice" and ref_audio_path is None:
        return None, "⚠️  Please upload a reference audio file."

    headers = {
        "Authorization": f"Bearer {api_key.strip()}",
        "Content-Type": "application/json",
    }

    clone_quality_warning = None

    # Pre-load reference audio once (for single-speaker cloning mode).
    # ref_audio is sent as a data URI, which the API accepts directly
    # alongside base64/URL — this avoids any ambiguity about encoding.
    ref_audio_data_uri = None
    if mode == "Clone my voice":
        ref_audio_data_uri = _file_to_data_uri(ref_audio_path)
        if not (ref_transcript or "").strip():
            # FIX #7: don't silently send ref_text=None — flag it, since
            # missing transcripts can noticeably hurt clone accuracy.
            clone_quality_warning = (
                "ℹ️  No reference transcript was provided for voice cloning — "
                "the clone may be less accurate. Adding the exact words spoken "
                "in the reference audio is recommended."
            )

    # Build the speaker roster for multi-speaker mode: name -> {ref_audio, ref_text}
    # FIX #1: speaker_triples is already a clean list of (name, audio_path,
    # transcript) tuples built at the UI-wiring boundary, so there's no
    # positional re-slicing here that could misalign name/audio/transcript.
    roster: dict[str, dict] = {}
    known_speakers: list[str] = []
    missing_transcripts: list[str] = []
    if mode == "Multi-speaker (cloned voices)":
        for name, audio_path, transcript in speaker_triples:
            name = (name or "").strip()
            if not name or audio_path is None:
                continue
            if name in roster:
                return None, f"⚠️  Duplicate speaker name '{name}' — names must be unique."
            transcript = (transcript or "").strip()
            if not transcript:
                missing_transcripts.append(name)
            roster[name] = {
                "ref_audio": _file_to_data_uri(audio_path),
                "ref_text": transcript or None,
            }
            known_speakers.append(name)

        if not known_speakers:
            return None, "⚠️  Add at least one speaker with a name and reference audio."

        if missing_transcripts:
            clone_quality_warning = (
                "ℹ️  No reference transcript given for: "
                f"{', '.join(missing_transcripts)} — clone accuracy may suffer for them."
            )

        tags_in_text = _SPEAKER_TAG_RE.findall(text)
        unmatched = sorted(set(tags_in_text) - set(known_speakers))
        if unmatched and not any(s in known_speakers for s in tags_in_text):
            # No tag in the text matches any configured speaker at all —
            # almost certainly a naming mismatch, worth surfacing clearly
            # rather than silently dumping everything on speaker 1.
            return None, (
                f"⚠️  None of the [tags] in your text match a configured speaker name.\n"
                f"   Found in text: {', '.join(unmatched)}\n"
                f"   Configured speakers: {', '.join(known_speakers)}"
            )

    # Build the ordered plan: text chunks ready for the API, interleaved
    # with pause items that we'll splice in as real silence (not sent to
    # the model at all — guarantees the pause actually happens).
    plan  = build_generation_plan(text, known_speakers=known_speakers or None)
    total = sum(1 for item in plan if item["type"] == "text")
    if total == 0:
        return None, "⚠️  No speakable text found (only pauses?)."

    progress(0, desc=f"Generating 0 / {total} chunks…")

    # Prefer raw PCM (no lossy compression at all). If the API rejects the
    # format for this account/region, fall back to MP3 for the rest of the
    # job rather than failing the whole generation.
    response_format = "pcm"
    format_fallback_notice = None
    errors: list[str] = []

    def build_payload_for_item(item):
        """Build the API payload for a single 'text' plan item, or return
        (None, error_message) if no voice is configured for it (e.g. a
        speaker mismatch that slipped past the earlier validation)."""
        chunk  = item["content"]
        tagged = build_tagged_text(chunk, emotion, style, slow)
        payload = {
            "model": "higgs-audio-v3-tts",
            "input": tagged,
        }
        if mode == "Preset voice":
            payload["voice"] = preset_voice
        elif mode == "Clone my voice":
            payload["ref_audio"] = ref_audio_data_uri
            payload["ref_text"]  = ref_transcript.strip() or None
        else:  # Multi-speaker (cloned voices)
            speaker = item.get("speaker")
            voice_info = roster.get(speaker)
            if voice_info is None:
                return None, f"no voice configured for speaker '{speaker}'"
            payload["ref_audio"] = voice_info["ref_audio"]
            payload["ref_text"]  = voice_info["ref_text"]
        return payload, None

    # Index every 'text' item in the plan so results can be matched back
    # to their original position after concurrent fetching — order in
    # `plan` is what ultimately determines playback order, not the order
    # requests complete in.
    text_indices = [i for i, item in enumerate(plan) if item["type"] == "text"]
    raw_by_index: dict[int, bytes] = {}
    done = 0

    # Fetch the first chunk alone, sequentially. This does double duty:
    # (1) surfaces an auth/payload error immediately instead of after
    # firing off a whole batch of doomed concurrent requests, and
    # (2) pins down whether PCM or the MP3 fallback is actually in use
    # for this account, so every concurrent request below can go out
    # with the correct format from the start rather than each thread
    # discovering (and retrying around) the fallback independently.
    first_idx = text_indices[0]
    payload, build_err = build_payload_for_item(plan[first_idx])
    if build_err:
        errors.append(f"Chunk 1: {build_err}")
    else:
        try:
            raw, used_format = _call_with_format_fallback(headers, payload, response_format)
            if used_format != response_format:
                response_format = used_format
                format_fallback_notice = (
                    "ℹ️  This account/region didn't accept PCM output; "
                    "fell back to MP3 (a small quality trade-off vs. raw audio)."
                )
            raw_by_index[first_idx] = raw
        except Exception as e:
            errors.append(f"Chunk 1: {e}")
    done += 1
    progress(done / total, desc=f"Generated {done} / {total} chunks…")

    # Remaining chunks fan out across a bounded thread pool. Each request
    # is I/O-bound (waiting on Boson's API over HTTP), so concurrency here
    # cuts real wall-clock time roughly in proportion to worker count —
    # this is what makes a 10-minute script (dozens of chunks) practical
    # instead of a long, fully-sequential wait. call_api()'s existing 429
    # back-off still applies independently within each thread.
    remaining_indices = text_indices[1:]
    if remaining_indices:
        workers = max(1, min(int(concurrency or DEFAULT_CONCURRENCY), MAX_CONCURRENCY, len(remaining_indices)))
        with ThreadPoolExecutor(max_workers=workers) as executor:
            future_to_idx = {}
            for idx in remaining_indices:
                payload, build_err = build_payload_for_item(plan[idx])
                if build_err:
                    errors.append(f"Chunk {text_indices.index(idx) + 1}: {build_err}")
                    continue
                future = executor.submit(_call_with_format_fallback, headers, payload, response_format)
                future_to_idx[future] = idx

            for future in as_completed(future_to_idx):
                idx = future_to_idx[future]
                chunk_num = text_indices.index(idx) + 1
                try:
                    raw, used_format = future.result()
                    if used_format != response_format:
                        # A later chunk hit the PCM/MP3 fallback independently
                        # (rare once the probe chunk above has already run,
                        # but possible under a flaky/region-inconsistent
                        # account) — decode it with the format it actually
                        # came back as rather than assuming the probe's format.
                        format_fallback_notice = format_fallback_notice or (
                            "ℹ️  Some chunks fell back from PCM to MP3 output "
                            "(a small quality trade-off vs. raw audio)."
                        )
                        raw_by_index[idx] = ("mp3", raw)
                    else:
                        raw_by_index[idx] = raw
                except Exception as e:
                    speaker = plan[idx].get("speaker")
                    label = f"Chunk {chunk_num}" + (f" ({speaker})" if speaker else "")
                    errors.append(f"{label}: {e}")
                done += 1
                progress(done / total, desc=f"Generated {done} / {total} chunks…")

    # Assembly happens strictly in plan order (fast, in-memory, no network
    # calls) — this is what guarantees pauses land exactly where the
    # script placed them and speech never gets reordered, regardless of
    # which order the concurrent requests above actually completed in.
    segments: list[AudioSegment] = []
    for idx, item in enumerate(plan):
        if item["type"] == "pause":
            # Real, deterministic silence — does not depend on the model
            # honoring a prosody token. Match frame rate/channels to the
            # most recent real speech segment (falls back to pydub's
            # default if no speech has been generated yet) to avoid any
            # sample-rate mismatch artifacts when segments are joined.
            ref = next((s for s in reversed(segments) if s.rms > 0), None)
            silence = AudioSegment.silent(duration=item["ms"])
            if ref is not None:
                silence = silence.set_frame_rate(ref.frame_rate) \
                                  .set_channels(ref.channels)
            segments.append(silence)
            continue

        result = raw_by_index.get(idx)
        if result is None:
            # This chunk failed (already recorded in `errors` above) —
            # skip it but keep assembling the rest, so one bad chunk in a
            # long script doesn't discard everything around it.
            continue
        if isinstance(result, tuple):
            # A per-chunk MP3 fallback (see the concurrent loop above) —
            # decode with its own format rather than the job-wide one.
            chunk_format, raw = result
        else:
            chunk_format, raw = response_format, result
        segments.append(bytes_to_audiosegment(raw, chunk_format))

    if not segments or all(isinstance(s, AudioSegment) and s.duration_seconds == 0 for s in segments):
        err_detail = "\n".join(errors)
        return None, f"❌ All chunks failed:\n{err_detail}"

    if enhance_audio:
        # Cap stray long silences the model sometimes appends at chunk
        # edges, then even out per-chunk loudness so volume doesn't jump
        # between separate API calls.
        segments = [
            s if s.rms == 0 else trim_edge_silence(s)
            for s in segments
        ]
        segments = match_loudness(segments)

    combined = concatenate_segments(segments, crossfade_ms=DEFAULT_CROSSFADE_MS)

    if enhance_audio:
        # Remove any DC offset/sub-audible rumble, then bring the final
        # mix up to a safe, consistent peak level without clipping.
        combined = combined.high_pass_filter(60)
        combined = effects.normalize(combined, headroom=FINAL_PEAK_HEADROOM_DB)

    out_path = _export_combined(combined, output_format)

    status_parts = [f"✅ Done! {len(segments)} / {total} chunks rendered."]
    status_parts.append(f"🎚️  Source audio: {'raw PCM (lossless)' if response_format == 'pcm' else 'MP3'} → exported as {output_format}.")
    if format_fallback_notice:
        status_parts.append(format_fallback_notice)
    if clone_quality_warning:
        status_parts.append(clone_quality_warning)
    if errors:
        # FIX #6: cap how many individual error lines get dumped into the
        # status box so a job with dozens of failures stays readable.
        shown = errors[:MAX_ERRORS_SHOWN]
        remaining = len(errors) - len(shown)
        status_parts.append(f"⚠️  {len(errors)} chunk(s) had errors (skipped):")
        status_parts.extend(f"  • {e}" for e in shown)
        if remaining > 0:
            status_parts.append(f"  …and {remaining} more.")

    return out_path, "\n".join(status_parts)


def _call_with_format_fallback(headers: dict, base_payload: dict, preferred_format: str):
    """
    Try the API with preferred_format first (PCM by default, for zero
    compression loss). If the server rejects that response_format
    specifically, retry the same chunk once with MP3 instead of failing
    the whole job. Returns (raw_bytes, format_actually_used).
    """
    payload = dict(base_payload)
    payload["response_format"] = preferred_format
    # FIX #2: PCM responses from this API are a plain binary body, not a
    # chunked/SSE stream — requesting stream=True for a non-streamed
    # response risked call_api() reading a malformed buffer with no error
    # surfaced. Controlled by STREAM_PCM in case the API's behavior differs.
    payload["stream"] = (preferred_format == "pcm") and STREAM_PCM
    try:
        return call_api(headers, payload), preferred_format
    except UnsupportedFormatError:
        if preferred_format == "mp3":
            raise
        fallback_payload = dict(base_payload)
        fallback_payload["response_format"] = "mp3"
        fallback_payload["stream"] = False
        return call_api(headers, fallback_payload), "mp3"


def _file_to_data_uri(path: str) -> str:
    import base64, mimetypes
    mime, _ = mimetypes.guess_type(path)
    mime = mime or "audio/wav"
    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    return f"data:{mime};base64,{b64}"


def _export_combined(combined: AudioSegment, output_format: str) -> str:
    # FIX #4: tempfile.mktemp() is documented as insecure — it returns a
    # name with no file actually created, leaving a race window between
    # name generation and the file being written. NamedTemporaryFile with
    # delete=False reserves the file atomically instead.
    if output_format.startswith("FLAC"):
        suffix, fmt, kwargs = ".flac", "flac", {}
    elif output_format.startswith("MP3"):
        suffix, fmt, kwargs = ".mp3", "mp3", {"bitrate": "320k"}
    else:  # WAV — lossless, default
        suffix, fmt, kwargs = ".wav", "wav", {}

    tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    tmp.close()
    combined.export(tmp.name, format=fmt, **kwargs)
    return tmp.name


# ── UI ────────────────────────────────────────────────────────────────────────

CSS = """
body, .gradio-container {
    background: #0f1117 !important;
    color: #e2e8f0 !important;
    font-family: 'Inter', system-ui, sans-serif !important;
}
.gr-box, .gr-form, .gr-panel {
    background: #1a1d27 !important;
    border: 1px solid #2d3148 !important;
    border-radius: 12px !important;
}
.gr-button-primary {
    background: linear-gradient(135deg, #6366f1, #8b5cf6) !important;
    border: none !important;
    border-radius: 8px !important;
    font-weight: 600 !important;
    letter-spacing: 0.02em !important;
    padding: 10px 28px !important;
    color: #fff !important;
    transition: opacity 0.15s !important;
}
.gr-button-primary:hover { opacity: 0.88 !important; }
.pause-btn {
    background: #1f2335 !important;
    border: 1px solid #2d3148 !important;
    border-radius: 8px !important;
    color: #c7d2fe !important;
    font-size: 0.78rem !important;
    padding: 6px 12px !important;
}
.pause-btn:hover { background: #262b42 !important; }
label, .gr-label { color: #94a3b8 !important; font-size: 0.8rem !important; }
textarea, input[type="text"], input[type="password"] {
    background: #0f1117 !important;
    border: 1px solid #2d3148 !important;
    color: #e2e8f0 !important;
    border-radius: 8px !important;
}
.gr-dropdown select {
    background: #0f1117 !important;
    color: #e2e8f0 !important;
    border: 1px solid #2d3148 !important;
}
footer { display: none !important; }
"""

TITLE_MD = """
# 🎙 Higgs Audio v3 — High-Fidelity TTS
Native 24 kHz PCM pipeline (no double lossy compression) • Loudness-matched,
glitch-free joins • Lossless WAV/FLAC export • Long-text auto-chunking •
concurrent generation for long-form scripts (10+ minutes) •
[Get a free key at boson.ai](https://boson.ai/workspace)
"""

# FIX #5: delete_cache periodically GCs files this app writes into
# Gradio's temp dir (e.g. our exported WAV/FLAC/MP3 outputs), instead of
# letting them accumulate forever on a long-running Space.
# (cache_lifetime_seconds, cleanup_interval_seconds) — both set to 1 hour.
with gr.Blocks(title="Higgs Audio v3 TTS", delete_cache=(3600, 3600)) as demo:

    gr.Markdown(TITLE_MD)

    with gr.Row():
        api_key = gr.Textbox(
            label="Boson API key",
            placeholder="bai-xxxx…",
            type="password",
            scale=4,
        )

    with gr.Row():
        # ── Left column: controls ──────────────────────────────────────────
        with gr.Column(scale=3):
            text = gr.Textbox(
                label="Text  (up to ~50 000 chars — long-form scripts, e.g. 10+ minutes, are fully supported)",
                lines=12,
                placeholder="Paste or type your script here. Long texts are automatically split into smooth chunks…",
                max_lines=40,
            )
            duration_estimate = gr.Markdown("≈ 0:00 estimated")
            text.change(estimate_spoken_duration, inputs=text, outputs=duration_estimate)

            with gr.Row():
                add_short_pause_btn = gr.Button(
                    "⏸  Insert short pause (~0.5s)", elem_classes="pause-btn"
                )
                add_long_pause_btn = gr.Button(
                    "⏸⏸  Insert long pause (~1s)", elem_classes="pause-btn"
                )
            gr.Markdown(
                "Pauses are inserted at the end of the text above — drag/cut-paste "
                "the `<|prosody:...|>` tag to move it wherever you want the pause to occur. "
                "These are spliced in as real silence, so the pause is guaranteed "
                "regardless of how the model would otherwise interpret the tag.",
            )

            with gr.Row():
                emotion = gr.Dropdown(EMOTIONS, value="none", label="Emotion")
                style   = gr.Dropdown(STYLES,   value="none", label="Style")
                slow    = gr.Checkbox(label="Slow speech", value=False)

            mode = gr.Radio(
                ["Preset voice", "Clone my voice", "Multi-speaker (cloned voices)"],
                value="Preset voice",
                label="Voice mode",
            )

            with gr.Group() as preset_group:
                preset_voice = gr.Dropdown(
                    VOICES, value="default", label="Voice",
                    allow_custom_value=True,
                    info="Pick a preset, or type any voice name from the Boson Voices page.",
                )

            with gr.Group(visible=False) as clone_group:
                ref_audio = gr.Audio(
                    label="Reference audio (5–30 s recommended)",
                    type="filepath",
                )
                ref_transcript = gr.Textbox(
                    label="Reference transcript (recommended for accuracy)",
                    placeholder="Exact words spoken in the audio above…",
                    lines=2,
                )

            with gr.Group(visible=False) as multispeaker_group:
                gr.Markdown(
                    f"Add up to {MAX_SPEAKERS} speakers, then tag your script with "
                    f"`[Name]` before each speaker's lines, e.g.:\n\n"
                    f"```\n[Maya]\nHello there!\n[Omar]\nHi Maya, good to see you.\n```\n\n"
                    f"Speaker names are case-sensitive and must exactly match the "
                    f"`[tags]` used in the text box above.",
                )
                num_speakers = gr.Slider(
                    minimum=2, maximum=MAX_SPEAKERS, value=2, step=1,
                    label="Number of speakers",
                )

                speaker_rows = []
                speaker_name_boxes = []
                speaker_audio_boxes = []
                speaker_transcript_boxes = []
                for i in range(MAX_SPEAKERS):
                    with gr.Group(visible=(i < 2)) as row:
                        gr.Markdown(f"**Speaker {i + 1}**")
                        name_box = gr.Textbox(
                            label="Name (used in [tags])",
                            placeholder=f"e.g. Speaker{i + 1}",
                        )
                        with gr.Row():
                            audio_box = gr.Audio(
                                label="Reference audio",
                                type="filepath",
                            )
                            transcript_box = gr.Textbox(
                                label="Reference transcript",
                                placeholder="Exact words spoken in the audio…",
                                lines=2,
                            )
                    speaker_rows.append(row)
                    speaker_name_boxes.append(name_box)
                    speaker_audio_boxes.append(audio_box)
                    speaker_transcript_boxes.append(transcript_box)

                def toggle_speaker_rows(n):
                    """
                    FIX #3: previously this only toggled visibility, so a
                    speaker filled in and then hidden by lowering the
                    slider stayed in speaker_name_boxes/audio_boxes/
                    transcript_boxes state and was still picked up by
                    run_generate() → generate(), silently including a
                    "removed" speaker in the roster. Now hidden rows are
                    also cleared, so hiding a row actually drops it.
                    """
                    row_updates = [gr.update(visible=(i < n)) for i in range(MAX_SPEAKERS)]
                    name_updates = [
                        gr.update() if i < n else gr.update(value="")
                        for i in range(MAX_SPEAKERS)
                    ]
                    audio_updates = [
                        gr.update() if i < n else gr.update(value=None)
                        for i in range(MAX_SPEAKERS)
                    ]
                    transcript_updates = [
                        gr.update() if i < n else gr.update(value="")
                        for i in range(MAX_SPEAKERS)
                    ]
                    return row_updates + name_updates + audio_updates + transcript_updates

                num_speakers.change(
                    toggle_speaker_rows,
                    inputs=num_speakers,
                    outputs=(
                        speaker_rows
                        + speaker_name_boxes
                        + speaker_audio_boxes
                        + speaker_transcript_boxes
                    ),
                )

            def toggle_mode(m):
                return (
                    gr.update(visible=(m == "Preset voice")),
                    gr.update(visible=(m == "Clone my voice")),
                    gr.update(visible=(m == "Multi-speaker (cloned voices)")),
                )

            mode.change(
                toggle_mode,
                inputs=mode,
                outputs=[preset_group, clone_group, multispeaker_group],
            )

            with gr.Row():
                output_format = gr.Dropdown(
                    OUTPUT_FORMATS, value=OUTPUT_FORMATS[0], label="Output format",
                )
                enhance_audio = gr.Checkbox(
                    label="Auto-enhance (loudness matching + cleanup)",
                    value=True,
                )

            concurrency = gr.Slider(
                minimum=1, maximum=MAX_CONCURRENCY, value=DEFAULT_CONCURRENCY, step=1,
                label="Parallel chunk requests",
                info=(
                    "Higher = faster generation for long scripts (10+ minutes), but more "
                    "simultaneous load on your Boson account. Set to 1 for the old, fully "
                    "sequential behavior."
                ),
            )

            generate_btn = gr.Button("▶  Generate audio", variant="primary")

        # ── Right column: output ───────────────────────────────────────────
        with gr.Column(scale=2):
            audio_out = gr.Audio(label="Output audio", type="filepath")
            status    = gr.Textbox(
                label="Status",
                interactive=False,
                lines=4,
                placeholder="Status will appear here…",
            )

    # ── wiring ──────────────────────────────────────────────────────────────
    add_short_pause_btn.click(
        fn=lambda t: insert_token_at_cursor(t, None, PAUSE_SHORT),
        inputs=text,
        outputs=text,
    )
    add_long_pause_btn.click(
        fn=lambda t: insert_token_at_cursor(t, None, PAUSE_LONG),
        inputs=text,
        outputs=text,
    )

    def run_generate(api_key, mode, text, emotion, style, slow,
                      preset_voice, ref_audio, ref_transcript,
                      output_format, enhance_audio, concurrency, *speaker_args):
        # FIX #1: zip the three flat argument lists into (name, audio,
        # transcript) triples right here, at the single point where their
        # order is guaranteed to still match (they were laid out in the
        # same per-speaker loop above). generate() now receives already-
        # paired triples instead of re-deriving the pairing from position,
        # so a future change to the input ordering in generate_btn.click
        # can't silently cross-wire one speaker's name with another's audio.
        n = MAX_SPEAKERS
        names        = speaker_args[0:n]
        audios       = speaker_args[n:2 * n]
        transcripts  = speaker_args[2 * n:3 * n]
        speaker_triples = list(zip(names, audios, transcripts))
        return generate(
            api_key, mode, text, emotion, style, slow,
            preset_voice, ref_audio, ref_transcript,
            speaker_triples,
            output_format, enhance_audio, concurrency,
        )

    generate_btn.click(
        fn=run_generate,
        inputs=[
            api_key, mode, text, emotion, style, slow,
            preset_voice, ref_audio, ref_transcript,
            output_format, enhance_audio, concurrency,
            *speaker_name_boxes, *speaker_audio_boxes, *speaker_transcript_boxes,
        ],
        outputs=[audio_out, status],
    )

# ── entry point ─────────────────────────────────────────────────────────────
# Hugging Face Spaces runs this file directly and manages the public URL
# itself — no share=True, no ngrok/frpc tunnel, no Colab proxy involved.
# `demo.queue()` still matters here: it switches Gradio to a persistent
# websocket connection instead of one blocking HTTP request, which is what
# avoids 504s on long multi-chunk generations.
demo.queue()

if __name__ == "__main__":
    demo.launch(css=CSS)