app.py

"""
Qwen3-TTS Audiobook Generator
Two modes: Single Speaker | Multi-Speaker
Self-hosted on HF Spaces with ZeroGPU
"""

import os
import json
import re
import tempfile
import time
import subprocess
import shutil
import struct

import torch
import spaces
import gradio as gr
import soundfile as sf
import requests as http_requests
from qwen_tts import Qwen3TTSModel

try:
    from openai import OpenAI
    HAS_OPENAI = True
except ImportError:
    HAS_OPENAI = False

# ==========================================
# CONFIG
# ==========================================
OMNI_MODEL = "qwen3.5-omni-plus"
DASHSCOPE_BASE_URL = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
ELEVENLABS_TTS_URL = "https://api.elevenlabs.io/v1/text-to-speech"

OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "qwen3_tts_out")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Languages with engine routing
LANGUAGE_CONFIG = {
    # Qwen3-TTS (local GPU, free)
    "English": {"engine": "qwen"},
    "Chinese": {"engine": "qwen"},
    "Japanese": {"engine": "qwen"},
    "Korean": {"engine": "qwen"},
    "German": {"engine": "qwen"},
    "French": {"engine": "qwen"},
    "Russian": {"engine": "qwen"},
    "Portuguese": {"engine": "qwen"},
    "Spanish": {"engine": "qwen"},
    "Italian": {"engine": "qwen"},
    # ElevenLabs (API)
    "Arabic": {"engine": "elevenlabs"},
    "English (Jamaican)": {"engine": "elevenlabs"},
}

LANGUAGES = ["Auto"] + list(LANGUAGE_CONFIG.keys())

def get_engine(lang):
    return LANGUAGE_CONFIG.get(lang, {}).get("engine", "qwen")

# Qwen speakers
SPEAKERS = {
    "Ryan": {"desc": "Dynamic male, strong rhythmic drive", "gender": "male"},
    "Aiden": {"desc": "Sunny American male, clear midrange", "gender": "male"},
    "Dylan": {"desc": "Youthful Beijing male, clear natural", "gender": "male"},
    "Uncle_Fu": {"desc": "Seasoned male, low mellow timbre", "gender": "male"},
    "Eric": {"desc": "Lively Chengdu male, slightly husky", "gender": "male"},
    "Vivian": {"desc": "Bright, edgy young female", "gender": "female"},
    "Serena": {"desc": "Warm, gentle young female", "gender": "female"},
    "Ono_Anna": {"desc": "Playful Japanese female, light nimble", "gender": "female"},
    "Sohee": {"desc": "Warm Korean female, rich emotion", "gender": "female"},
}

SPEAKER_CHOICES = [f"{n} -- {s['desc']}" for n, s in SPEAKERS.items()]
MALE_SPEAKERS = [n for n, s in SPEAKERS.items() if s["gender"] == "male"]
FEMALE_SPEAKERS = [n for n, s in SPEAKERS.items() if s["gender"] == "female"]

# ElevenLabs voices
ELEVENLABS_VOICES = {
    "Arabic": [
        {"name": "Rachel", "id": "21m00Tcm4TlvDq8ikWAM", "desc": "Calm female", "gender": "female"},
        {"name": "Drew", "id": "29vD33N1CtxCmqQRPOHJ", "desc": "Rounded male", "gender": "male"},
        {"name": "Paul", "id": "5Q0t7uMcjvnagumLfvZi", "desc": "Authoritative male", "gender": "male"},
        {"name": "Matilda", "id": "XrExE9yKIg1WjnnlVkGX", "desc": "Warm female", "gender": "female"},
    ],
    "English (Jamaican)": [
        {"name": "Clyde", "id": "2EiwWnXFnvU5JabPnv8n", "desc": "Deep, masculine male", "gender": "male"},
        {"name": "Freya", "id": "jsCqWAovK2LkecY7zXl4", "desc": "Young expressive female", "gender": "female"},
        {"name": "Antoni", "id": "ErXwobaYiN019PkySvjV", "desc": "Warm rounded male", "gender": "male"},
        {"name": "Rachel", "id": "21m00Tcm4TlvDq8ikWAM", "desc": "Calm warm female", "gender": "female"},
    ],
}

EL_SPEAKER_CHOICES = {
    lang: [f"{v['name']} -- {v['desc']}" for v in voices]
    for lang, voices in ELEVENLABS_VOICES.items()
}

EL_MALE = {lang: [v for v in voices if v["gender"] == "male"] for lang, voices in ELEVENLABS_VOICES.items()}
EL_FEMALE = {lang: [v for v in voices if v["gender"] == "female"] for lang, voices in ELEVENLABS_VOICES.items()}

def get_el_voice_id(lang, label):
    name = label.split("--")[0].strip()
    for v in ELEVENLABS_VOICES.get(lang, ELEVENLABS_VOICES.get("Arabic", [])):
        if v["name"] == name:
            return v["id"]
    # Fallback to first voice
    voices = ELEVENLABS_VOICES.get(lang, ELEVENLABS_VOICES.get("Arabic", []))
    return voices[0]["id"] if voices else "21m00Tcm4TlvDq8ikWAM"

# ==========================================
# MODEL LOADING
# ==========================================
_models = {}

def get_model(model_type):
    if model_type not in _models:
        model_map = {
            "custom": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
            "clone": "Qwen/Qwen3-TTS-12Hz-0.6B-Base",  # Use smaller model for cloning to save VRAM
        }
        print(f"[TTS] Loading {model_map[model_type]}...")
        _models[model_type] = Qwen3TTSModel.from_pretrained(
            model_map[model_type], device_map="cuda:0", dtype=torch.bfloat16,
        )
    return _models[model_type]


def get_llm_client():
    key = os.environ.get("DASHSCOPE_API_KEY", "")
    if key and HAS_OPENAI:
        return OpenAI(api_key=key, base_url=DASHSCOPE_BASE_URL)
    return None


# ==========================================
# AUDIO HELPERS
# ==========================================
def concatenate_wavs(files, out):
    if not files:
        return
    if len(files) == 1:
        shutil.copy2(files[0], out)
        return
    lst = out + ".txt"
    with open(lst, "w") as f:
        for w in files:
            f.write(f"file '{w}'\n")
    subprocess.run(["ffmpeg", "-y", "-f", "concat", "-safe", "0",
                    "-i", lst, "-c", "copy", out], capture_output=True, check=True)
    os.remove(lst)


def make_silence(dur, path):
    subprocess.run(["ffmpeg", "-y", "-f", "lavfi", "-i", "anullsrc=r=24000:cl=mono",
                    "-t", str(dur), "-acodec", "pcm_s16le", path],
                   capture_output=True, check=True)


# ==========================================
# TEXT HELPERS
# ==========================================
def resolve_text(text_input, file_input):
    if file_input is not None:
        ext = os.path.splitext(file_input)[1].lower()
        if ext == ".pdf":
            import pypdf
            reader = pypdf.PdfReader(file_input)
            return "\n\n".join(p.extract_text().strip() for p in reader.pages if p.extract_text())
        elif ext == ".docx":
            import docx
            doc = docx.Document(file_input)
            return "\n\n".join(p.text.strip() for p in doc.paragraphs if p.text.strip())
        else:
            with open(file_input, "r", encoding="utf-8", errors="replace") as f:
                return f.read()
    elif text_input and text_input.strip():
        return text_input.strip()
    raise gr.Error("Please enter text or upload a file.")


def extract_pdf_sections(filepath):
    import pypdf
    reader = pypdf.PdfReader(filepath)
    full = ""
    for p in reader.pages:
        t = p.extract_text()
        if t:
            full += t + "\n\n"

    sections = []
    pattern = r'(?m)^(?:(?:Chapter|CHAPTER|Part|PART|Section|SECTION)\s+\w+[.:)?\s].*|(?:\d+[.)]\s+[A-Z].*))'
    matches = list(re.finditer(pattern, full))

    if matches:
        for i, m in enumerate(matches):
            start = m.start()
            end = matches[i + 1].start() if i + 1 < len(matches) else len(full)
            title = m.group().strip()[:80]
            content = full[start:end].strip()
            if len(content) > 20:
                sections.append({"title": title, "content": content, "chars": len(content)})

    if not sections:
        paragraphs = re.split(r'\n\s*\n', full)
        current, num = "", 1
        for para in paragraphs:
            para = para.strip()
            if not para:
                continue
            if len(current) + len(para) > 3000 and current:
                sections.append({"title": f"Section {num}", "content": current.strip(), "chars": len(current)})
                num += 1
                current = para
            else:
                current += "\n\n" + para
        if current.strip():
            sections.append({"title": f"Section {num}", "content": current.strip(), "chars": len(current)})

    return sections


# ==========================================
# TRANSLATION
# ==========================================
SOURCE_LANGUAGES = ["English", "Chinese", "Japanese", "Korean", "German",
                    "French", "Russian", "Portuguese", "Spanish", "Italian", "Arabic",
                    "Hindi", "Swahili", "Auto-detect"]

def translate_text(client, text, source_lang, target_lang):
    """Translate text between languages. Handles long texts by chunking."""
    if not client:
        raise gr.Error("DASHSCOPE_API_KEY needed for translation.")
    if source_lang == target_lang:
        return text

    chunks = split_for_llm(text, max_chars=4000)
    translated_parts = []

    for ci, chunk in enumerate(chunks):
        source_hint = f" (source language: {source_lang})" if source_lang != "Auto-detect" else ""
        response = client.chat.completions.create(
            model=OMNI_MODEL, modalities=["text"],
            messages=[{
                "role": "system",
                "content": f"Translate the following text into {target_lang}. Output ONLY the translation.",
            }, {
                "role": "user",
                "content": f"Translate this{source_hint}:\n\n{chunk}",
            }],
        )
        translated_parts.append(response.choices[0].message.content.strip())
        print(f"[Translate] Chunk {ci+1}/{len(chunks)} done")

    result = "\n\n".join(translated_parts)
    print(f"[Translate] {source_lang} -> {target_lang}: {len(text)} -> {len(result)} chars")
    return result


# ==========================================
# EMOTION ANALYSIS
# ==========================================
def split_for_llm(text, max_chars=4000):
    """Split long text into chunks at paragraph boundaries for LLM processing."""
    if len(text) <= max_chars:
        return [text]
    chunks, paragraphs, current = [], re.split(r'\n\s*\n', text), ""
    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
        if len(current) + len(para) + 2 > max_chars and current:
            chunks.append(current.strip())
            current = para
        else:
            current = (current + "\n\n" + para).strip()
    if current.strip():
        chunks.append(current.strip())
    return chunks if chunks else [text]


def analyze_emotions(client, text):
    """Split text into segments with emotion instructions. Handles long texts."""
    if not client:
        return [{"text": text, "emotion": ""}]

    chunks = split_for_llm(text, max_chars=4000)
    all_segments = []

    for ci, chunk in enumerate(chunks):
        try:
            response = client.chat.completions.create(
                model=OMNI_MODEL, modalities=["text"],
                messages=[{
                    "role": "system",
                    "content": (
                        "You are an audiobook director. Split text into segments where the emotional "
                        "tone changes. For each segment, provide a specific emotion/delivery instruction.\n\n"
                        "Output ONLY valid JSON:\n"
                        '{"segments": [\n'
                        '  {"text": "The lighthouse stood tall.", "emotion": "Atmospheric, steady narration"},\n'
                        '  {"text": "She ran!", "emotion": "Urgent, breathless, rising tension"}\n'
                        "]}\n\n"
                        "Rules: Include ALL text. 1-4 sentences per segment. Be specific with emotions. "
                        "No markdown. ONLY JSON."
                    ),
                }, {"role": "user", "content": f"Direct this:\n\n{chunk}"}],
            )
            raw = response.choices[0].message.content.strip()
            raw = re.sub(r'^```json\s*', '', raw)
            raw = re.sub(r'\s*```$', '', raw)
            data = json.loads(raw)
            segs = data.get("segments", [])
            if segs:
                all_segments.extend(segs)
            else:
                all_segments.append({"text": chunk, "emotion": ""})
            print(f"[Emotions] Chunk {ci+1}/{len(chunks)}: {len(segs)} segments")
        except Exception as e:
            print(f"[Emotions] Chunk {ci+1} failed: {e}")
            all_segments.append({"text": chunk, "emotion": ""})

    print(f"[Emotions] Total: {len(all_segments)} segments from {len(chunks)} chunks")
    return all_segments if all_segments else [{"text": text, "emotion": ""}]


def detect_characters_and_emotions(client, text):
    """Detect characters + emotions for multi-speaker mode. Handles long texts."""
    chunks = split_for_llm(text, max_chars=5000)
    all_characters = {}
    all_segments = []

    for ci, chunk in enumerate(chunks):
        try:
            response = client.chat.completions.create(
                model=OMNI_MODEL, modalities=["text"],
                messages=[{
                    "role": "system",
                    "content": (
                        "You are an audiobook director. Analyze this story:\n"
                        "1. Identify all characters with genders\n"
                        "2. Split into segments by speaker\n"
                        "3. Add emotion instructions per segment\n\n"
                        "Output ONLY valid JSON:\n"
                        '{"characters": [{"name": "Narrator", "gender": "neutral"}, '
                        '{"name": "Elena", "gender": "female"}],\n'
                        '"segments": [{"speaker": "Narrator", "text": "...", "emotion": "Calm narration"}, '
                        '{"speaker": "Elena", "text": "...", "emotion": "Wistful, dreamy"}]}\n\n'
                        "Rules: Narrator handles non-dialogue. Include ALL text. "
                        "Be specific with emotions. No markdown. ONLY JSON."
                    ),
                }, {"role": "user", "content": f"Direct this story:\n\n{chunk}"}],
            )
            raw = response.choices[0].message.content.strip()
            raw = re.sub(r'^```json\s*', '', raw)
            raw = re.sub(r'\s*```$', '', raw)
            data = json.loads(raw)

            # Merge characters (keep unique by name)
            for c in data.get("characters", []):
                name = c.get("name", "Narrator")
                if name not in all_characters:
                    all_characters[name] = c

            all_segments.extend(data.get("segments", []))
            print(f"[MultiSpeaker] Chunk {ci+1}/{len(chunks)}: {len(data.get('segments', []))} segments")
        except Exception as e:
            print(f"[MultiSpeaker] Chunk {ci+1} failed: {e}")
            all_segments.append({"speaker": "Narrator", "text": chunk, "emotion": ""})

    # Ensure Narrator exists
    if "Narrator" not in all_characters:
        all_characters["Narrator"] = {"name": "Narrator", "gender": "neutral"}

    characters = list(all_characters.values())
    # Put Narrator first
    characters.sort(key=lambda c: 0 if c["name"] == "Narrator" else 1)

    print(f"[MultiSpeaker] Total: {len(characters)} characters, {len(all_segments)} segments")
    return characters, all_segments


# ==========================================
# ELEVENLABS TTS
# ==========================================
def inject_audio_tags(client, text):
    """Add ElevenLabs v3 audio tags for emotional delivery."""
    if not client:
        return text
    try:
        response = client.chat.completions.create(
            model=OMNI_MODEL, modalities=["text"],
            messages=[{
                "role": "system",
                "content": (
                    "Add ElevenLabs v3 audio tags to this text for expressive narration. "
                    "Tags are words in [brackets] like [whispers], [excited], [sighs], [laughs], "
                    "[softly], [firmly], [dramatically]. Place before the phrase they apply to. "
                    "Use sparingly (1 tag per 2-3 sentences). Output ONLY the tagged text."
                ),
            }, {"role": "user", "content": text}],
        )
        return response.choices[0].message.content.strip()
    except Exception:
        return text


def clone_voice_elevenlabs(audio_path, api_key):
    """Clone a voice using ElevenLabs Instant Voice Cloning. Returns voice_id."""
    headers = {"xi-api-key": api_key}
    with open(audio_path, "rb") as f:
        files = [("files", (os.path.basename(audio_path), f, "audio/mpeg"))]
        data = {"name": f"clone_{int(time.time())}", "description": "Cloned voice", "remove_background_noise": "true"}
        resp = http_requests.post("https://api.elevenlabs.io/v1/voices/add",
                                  headers=headers, data=data, files=files, timeout=60)
    if resp.status_code != 200:
        raise RuntimeError(f"Voice clone failed ({resp.status_code}): {resp.text[:300]}")
    voice_id = resp.json().get("voice_id")
    if not voice_id:
        raise RuntimeError(f"No voice_id in response: {resp.text[:300]}")
    print(f"[EL] Voice cloned: {voice_id}")
    return voice_id


def tts_elevenlabs(text, voice_id, api_key, seg_idx, tmp_dir, language=None, client=None):
    """Generate speech via ElevenLabs API. Returns (wav_path, error)."""
    output_mp3 = os.path.join(tmp_dir, f"el_{seg_idx:04d}.mp3")
    output_wav = os.path.join(tmp_dir, f"el_{seg_idx:04d}.wav")

    # Inject audio tags for emotional delivery
    final_text = text
    if client:
        final_text = inject_audio_tags(client, text)

    # Add accent instruction for Jamaican
    model_id = "eleven_v3"
    if language == "English (Jamaican)":
        final_text = f"[Jamaican accent] {final_text}"

    headers = {"xi-api-key": api_key, "Content-Type": "application/json"}
    payload = {
        "text": final_text,
        "model_id": model_id,
        "voice_settings": {"stability": 0.5, "similarity_boost": 0.75, "style": 0.3},
    }

    try:
        resp = http_requests.post(f"{ELEVENLABS_TTS_URL}/{voice_id}", headers=headers, json=payload, timeout=120)
        print(f"[EL] Seg {seg_idx}: status={resp.status_code}, {len(resp.content)} bytes")
        if resp.status_code != 200:
            print(f"[EL] Error: {resp.text[:300]}")
            return None, f"ElevenLabs failed ({resp.status_code}): {resp.text[:100]}"
        with open(output_mp3, "wb") as f:
            f.write(resp.content)
        subprocess.run(["ffmpeg", "-y", "-i", output_mp3, "-ar", "24000", "-ac", "1",
                        "-acodec", "pcm_s16le", output_wav], capture_output=True, check=True)
        return output_wav, None
    except Exception as e:
        return None, str(e)


# ==========================================
# SINGLE SPEAKER: Generate with emotions
# ==========================================
@spaces.GPU(duration=600)
def generate_single_speaker(text_input, file_input, source_lang, target_lang, speaker_label,
                             use_clone, clone_audio, clone_transcript,
                             progress=gr.Progress()):
    resolved = resolve_text(text_input, file_input)
    if len(resolved) < 5:
        raise gr.Error("Text too short.")

    speaker = speaker_label.split("--")[0].strip()
    lang = target_lang if target_lang != "Auto" else "Auto"
    is_clone = use_clone and clone_audio is not None
    needs_translation = source_lang != target_lang and source_lang != "Auto-detect" and target_lang != "Auto"

    # Translate if needed
    client = get_llm_client()
    print(f"[Single] source={source_lang}, target={target_lang}, needs_translation={needs_translation}, is_clone={is_clone}")
    if needs_translation:
        progress(0.03, desc=f"Translating {source_lang} to {target_lang}...")
        if not client:
            raise gr.Error("DASHSCOPE_API_KEY needed for translation.")
        resolved = translate_text(client, resolved, source_lang, target_lang)
        print(f"[Single] Translation complete: {len(resolved)} chars in {target_lang}")

    # Determine engine
    engine = get_engine(target_lang)
    el_key = os.environ.get("ELEVENLABS_API_KEY", "")
    is_elevenlabs = engine == "elevenlabs"

    if is_elevenlabs and not el_key:
        raise gr.Error("ELEVENLABS_API_KEY needed for this language. Add it in Settings > Secrets.")

    # Only load Qwen model for Qwen languages
    model = None
    if not is_elevenlabs:
        model_type = "clone" if is_clone else "custom"
        model = get_model(model_type)

    # Analyze emotions
    progress(0.05, desc="Analyzing emotions...")
    segments = analyze_emotions(client, resolved) if client else [{"text": resolved, "emotion": ""}]

    # Prepare clone prompt once (for Qwen clone only)
    clone_prompt = None
    if is_clone and not is_elevenlabs:
        progress(0.08, desc="Preparing voice clone...")

        # Trim reference audio to max 10 seconds to avoid OOM
        trimmed_ref = os.path.join(OUTPUT_DIR, f"ref_trimmed_{int(time.time())}.wav")
        subprocess.run([
            "ffmpeg", "-y", "-i", clone_audio,
            "-t", "10", "-ar", "24000", "-ac", "1", "-acodec", "pcm_s16le",
            trimmed_ref,
        ], capture_output=True, check=True)

        # Clear GPU cache before clone processing
        torch.cuda.empty_cache()

        clone_kwargs = {"ref_audio": trimmed_ref}
        if clone_transcript and clone_transcript.strip():
            clone_kwargs["ref_text"] = clone_transcript.strip()
            clone_kwargs["x_vector_only_mode"] = False
        else:
            clone_kwargs["x_vector_only_mode"] = True

        try:
            clone_prompt = model.create_voice_clone_prompt(**clone_kwargs)
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            raise gr.Error("GPU out of memory preparing voice clone. Try a shorter audio sample (3-5 seconds).")

    # ElevenLabs voice cloning
    el_cloned_voice_id = None
    if is_clone and is_elevenlabs:
        progress(0.08, desc="Cloning voice...")
        try:
            el_cloned_voice_id = clone_voice_elevenlabs(clone_audio, el_key)
        except Exception as e:
            raise gr.Error(f"Voice cloning failed: {e}")

    # Generate each segment
    tmp_dir = os.path.join(OUTPUT_DIR, f"single_{int(time.time())}")
    os.makedirs(tmp_dir, exist_ok=True)
    audio_files = []
    transcripts = []

    pause_path = os.path.join(tmp_dir, "pause.wav")
    make_silence(0.6, pause_path)

    total = len(segments)
    for i, seg in enumerate(segments):
        frac = 0.10 + 0.80 * (i / max(total, 1))
        seg_text = seg.get("text", "").strip()
        emotion = seg.get("emotion", "")
        if not seg_text:
            continue

        progress(frac, desc=f"Generating segment {i+1}/{total}...")
        path = os.path.join(tmp_dir, f"seg_{i:04d}.wav")

        try:
            if is_elevenlabs:
                # Use cloned voice_id if available, otherwise preset voice
                if el_cloned_voice_id:
                    voice_id = el_cloned_voice_id
                else:
                    voice_id = get_el_voice_id(target_lang, speaker_label)
                wav_path, error = tts_elevenlabs(
                    seg_text, voice_id, el_key, i, tmp_dir,
                    language=target_lang, client=client,
                )
                if wav_path:
                    audio_files.append(wav_path)
                else:
                    raise Exception(error or "ElevenLabs returned no audio")
            else:
                # Qwen local TTS
                # Clear cache periodically to prevent OOM buildup
                if i > 0 and i % 10 == 0:
                    torch.cuda.empty_cache()

                if is_clone:
                    wavs, sr = model.generate_voice_clone(
                        text=seg_text, language=lang, voice_clone_prompt=clone_prompt,
                    )
                else:
                    kwargs = {"text": seg_text, "language": lang, "speaker": speaker}
                    if emotion:
                        kwargs["instruct"] = emotion
                    wavs, sr = model.generate_custom_voice(**kwargs)

                sf.write(path, wavs[0], sr)
                audio_files.append(path)
            emotion_tag = f" *({emotion})*" if emotion else ""
            transcripts.append(f"{emotion_tag} {seg_text[:200]}{'...' if len(seg_text) > 200 else ''}")
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            print(f"[Single] OOM at segment {i}/{total} — saving partial result")
            transcripts.append(f"**GPU memory full at segment {i+1}/{total}. Partial audiobook saved.**")
            break  # Save what we have so far
        except Exception as e:
            print(f"[Single] Seg {i} failed: {e}")
            fail = os.path.join(tmp_dir, f"fail_{i}.wav")
            make_silence(1.0, fail)
            audio_files.append(fail)
            transcripts.append(f"**FAILED:** {str(e)[:100]}")

        if i < total - 1:
            audio_files.append(pause_path)

    completed = len([f for f in audio_files if "pause" not in f and "fail" not in f])
    is_partial = completed < total

    if not audio_files:
        raise gr.Error("No audio generated.")

    # Assemble
    progress(0.92, desc="Assembling...")
    final_wav = os.path.join(tmp_dir, "output.wav")
    print(f"[Single] Concatenating {len(audio_files)} files...")
    concatenate_wavs(audio_files, final_wav)

    if not os.path.exists(final_wav):
        raise gr.Error("Assembly failed — no output WAV created.")

    wav_size = os.path.getsize(final_wav) / (1024 * 1024)
    print(f"[Single] WAV assembled: {wav_size:.1f} MB")

    progress(0.96, desc="Converting to MP3...")
    final_mp3 = os.path.join(OUTPUT_DIR, f"single_{int(time.time())}.mp3")
    result = subprocess.run(["ffmpeg", "-y", "-i", final_wav, "-codec:a", "libmp3lame",
                    "-b:a", "128k", "-ar", "24000", "-ac", "1", final_mp3],
                   capture_output=True, text=True)
    if result.returncode != 0:
        print(f"[Single] FFmpeg error: {result.stderr[-500:]}")
        raise gr.Error(f"MP3 conversion failed: {result.stderr[-200:]}")

    if not os.path.exists(final_mp3) or os.path.getsize(final_mp3) < 100:
        raise gr.Error("MP3 conversion produced empty file.")

    print(f"[Single] MP3 ready: {final_mp3}, {os.path.getsize(final_mp3) / (1024*1024):.1f} MB")

    progress(1.0, desc="Done!")
    size = os.path.getsize(final_mp3) / (1024 * 1024)
    voice_info = "Cloned voice" if is_clone else speaker_label
    lang_info = f"{source_lang} → {target_lang}" if needs_translation else target_lang
    partial_note = f"\n\n> ⚠️ **Partial result:** {completed}/{total} segments generated. Text was too long for available GPU memory. Try shorter text or use the PDF section selector." if is_partial else ""
    stats = (
        f"**Audiobook Generated{'  (Partial)' if is_partial else ''}!**\n\n"
        f"- **Segments:** {completed}/{total} (with auto-emotions)\n"
        f"- **Voice:** {voice_info}\n"
        f"- **Language:** {lang_info}\n"
        f"- **File size:** {size:.1f} MB\n"
        f"{partial_note}"
    )
    transcript = "\n\n".join(transcripts)
    return final_mp3, stats, transcript


# ==========================================
# MULTI-SPEAKER: Generate with character voices
# ==========================================
@spaces.GPU(duration=600)
def generate_multi_speaker(text_input, file_input, source_lang, target_lang,
                           v0, v1, v2, v3, v4, v5, v6, v7,
                           progress=gr.Progress()):
    resolved = resolve_text(text_input, file_input)
    if len(resolved) < 30:
        raise gr.Error("Text too short for multi-speaker.")

    client = get_llm_client()
    if not client:
        raise gr.Error("DASHSCOPE_API_KEY needed. Add it in Settings > Secrets.")

    lang = target_lang if target_lang != "Auto" else "Auto"
    engine = get_engine(target_lang)
    el_key = os.environ.get("ELEVENLABS_API_KEY", "")
    is_elevenlabs = engine == "elevenlabs"
    needs_translation = source_lang != target_lang and source_lang != "Auto-detect" and target_lang != "Auto"

    if is_elevenlabs and not el_key:
        raise gr.Error("ELEVENLABS_API_KEY needed for this language.")

    # Translate first if needed
    if needs_translation:
        progress(0.03, desc=f"Translating {source_lang} to {target_lang}...")
        resolved = translate_text(client, resolved, source_lang, target_lang)

    # Only load Qwen model for Qwen languages
    model = None
    if not is_elevenlabs:
        model = get_model("custom")

    tmp_dir = os.path.join(OUTPUT_DIR, f"multi_{int(time.time())}")
    os.makedirs(tmp_dir, exist_ok=True)

    # Detect characters
    progress(0.05, desc="Detecting characters and emotions...")
    characters, segments = detect_characters_and_emotions(client, resolved)
    char_names = [c["name"] for c in characters]
    print(f"[Multi] {len(characters)} characters: {char_names}, {len(segments)} segments")

    # Assign voices — use custom assignments from dropdowns if provided
    voice_assignments = [v0, v1, v2, v3, v4, v5, v6, v7]
    voice_map = {}
    mi, fi = 0, 0
    for ci, c in enumerate(characters):
        name, gender = c["name"], c.get("gender", "neutral")

        # Use custom assignment if provided
        if ci < len(voice_assignments) and voice_assignments[ci]:
            custom_label = voice_assignments[ci]
            if is_elevenlabs:
                voice_name = custom_label.split("--")[0].strip()
                voice_id = get_el_voice_id(target_lang, custom_label)
                voice_map[name] = {"id": voice_id, "name": voice_name}
            else:
                voice_map[name] = {"speaker": custom_label.split("--")[0].strip()}
        else:
            # Auto-assign by gender
            if is_elevenlabs:
                el_m = EL_MALE.get(target_lang, EL_MALE.get("Arabic", []))
                el_f = EL_FEMALE.get(target_lang, EL_FEMALE.get("Arabic", []))
                if gender == "male" and el_m:
                    voice_map[name] = {"id": el_m[mi % len(el_m)]["id"], "name": el_m[mi % len(el_m)]["name"]}
                    mi += 1
                elif gender == "female" and el_f:
                    voice_map[name] = {"id": el_f[fi % len(el_f)]["id"], "name": el_f[fi % len(el_f)]["name"]}
                    fi += 1
                else:
                    voice_map[name] = {"id": "21m00Tcm4TlvDq8ikWAM", "name": "Rachel"}
            else:
                if name == "Narrator":
                    voice_map[name] = {"speaker": "Ryan"}
                elif gender == "male":
                    voice_map[name] = {"speaker": MALE_SPEAKERS[mi % len(MALE_SPEAKERS)]}
                    mi += 1
                elif gender == "female":
                    voice_map[name] = {"speaker": FEMALE_SPEAKERS[fi % len(FEMALE_SPEAKERS)]}
                    fi += 1
                else:
                    voice_map[name] = {"speaker": "Ryan"}
    print(f"[Multi] Voice map: {voice_map}")

    # Generate segments
    audio_files, transcripts = [], []
    speaker_pause = os.path.join(tmp_dir, "sp.wav")
    section_pause = os.path.join(tmp_dir, "sec.wav")
    make_silence(0.4, speaker_pause)
    make_silence(1.0, section_pause)

    total = len(segments)
    prev_speaker = None

    for i, seg in enumerate(segments):
        frac = 0.10 + 0.80 * (i / max(total, 1))
        speaker_name = seg.get("speaker", "Narrator")
        seg_text = seg.get("text", "").strip()
        emotion = seg.get("emotion", "")
        if not seg_text:
            continue

        voice_info = voice_map.get(speaker_name, voice_map.get("Narrator", {}))
        progress(frac, desc=f"[{speaker_name}] Segment {i+1}/{total}...")

        if prev_speaker and prev_speaker != speaker_name:
            audio_files.append(speaker_pause)

        path = os.path.join(tmp_dir, f"seg_{i:04d}.wav")
        try:
            if is_elevenlabs:
                voice_id = voice_info.get("id", "21m00Tcm4TlvDq8ikWAM")
                wav_path, error = tts_elevenlabs(
                    seg_text, voice_id, el_key, i, tmp_dir,
                    language=target_lang, client=client,
                )
                if wav_path:
                    audio_files.append(wav_path)
                else:
                    raise Exception(error or "ElevenLabs returned no audio")
            else:
                # Clear cache periodically
                if i > 0 and i % 10 == 0:
                    torch.cuda.empty_cache()

                voice = voice_info.get("speaker", "Ryan")
                kwargs = {"text": seg_text, "language": lang, "speaker": voice}
                if emotion:
                    kwargs["instruct"] = emotion
                wavs, sr = model.generate_custom_voice(**kwargs)
                sf.write(path, wavs[0], sr)
                audio_files.append(path)
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            print(f"[Multi] OOM at segment {i}/{total} — saving partial result")
            transcripts.append(f"**GPU memory full at segment {i+1}/{total}. Partial audiobook saved.**")
            break
        except Exception as e:
            print(f"[Multi] Seg {i} failed: {e}")
            fail = os.path.join(tmp_dir, f"fail_{i}.wav")
            make_silence(1.5, fail)
            audio_files.append(fail)

        emotion_tag = f" *({emotion})*" if emotion else ""
        transcripts.append(f"**[{speaker_name}]**{emotion_tag} {seg_text[:200]}{'...' if len(seg_text) > 200 else ''}")

        if i < total - 1:
            audio_files.append(section_pause)
        prev_speaker = speaker_name

    if not audio_files:
        raise gr.Error("No audio generated.")

    # Assemble
    progress(0.92, desc="Assembling audiobook...")
    final_wav = os.path.join(tmp_dir, "output.wav")
    print(f"[Multi] Concatenating {len(audio_files)} files...")
    concatenate_wavs(audio_files, final_wav)

    if not os.path.exists(final_wav):
        raise gr.Error("Assembly failed — no WAV created.")

    wav_size = os.path.getsize(final_wav) / (1024 * 1024)
    print(f"[Multi] WAV assembled: {wav_size:.1f} MB")

    progress(0.96, desc="Converting to MP3...")
    final_mp3 = os.path.join(OUTPUT_DIR, f"multi_{int(time.time())}.mp3")
    result = subprocess.run(["ffmpeg", "-y", "-i", final_wav, "-codec:a", "libmp3lame",
                    "-b:a", "128k", "-ar", "24000", "-ac", "1", final_mp3],
                   capture_output=True, text=True)
    if result.returncode != 0:
        print(f"[Multi] FFmpeg error: {result.stderr[-500:]}")
        raise gr.Error(f"MP3 conversion failed.")

    if not os.path.exists(final_mp3) or os.path.getsize(final_mp3) < 100:
        raise gr.Error("MP3 conversion produced empty file.")

    print(f"[Multi] MP3 ready: {final_mp3}, {os.path.getsize(final_mp3) / (1024*1024):.1f} MB")

    progress(1.0, desc="Done!")
    size = os.path.getsize(final_mp3) / (1024 * 1024)
    cast_lines = []
    for c in characters:
        v = voice_map.get(c['name'], {})
        voice_label = v.get("name", v.get("speaker", "?"))
        cast_lines.append(f"  - **{c['name']}** ({c.get('gender', '?')}) → {voice_label}")
    cast = "\n".join(cast_lines)
    stats = (
        f"**Multi-Speaker Audiobook Generated!**\n\n"
        f"- **Language:** {source_lang + ' → ' + target_lang if needs_translation else target_lang}\n"
        f"- **Segments:** {total}\n"
        f"- **Characters:** {len(characters)}\n"
        f"- **File size:** {size:.1f} MB\n\n"
        f"**Cast:**\n{cast}\n"
    )
    transcript = "\n\n".join(transcripts)
    return final_mp3, stats, transcript


# ==========================================
# UI HELPERS
# ==========================================
def detect_sections_ui(file_input):
    if file_input is None:
        raise gr.Error("Upload a PDF first.")
    if not file_input.lower().endswith(".pdf"):
        raise gr.Error("Section detection works with PDF files.")
    sections = extract_pdf_sections(file_input)
    if not sections:
        return "No sections found.", gr.update(visible=False), gr.update(visible=False), []
    choices = [f"{s['title']} ({s['chars']:,} chars)" for s in sections]
    info = f"**Found {len(sections)} sections:**\n\n"
    for i, s in enumerate(sections):
        preview = s["content"][:100].replace("\n", " ")
        info += f"{i+1}. **{s['title']}** ({s['chars']:,} chars) — {preview}...\n"
    return info, gr.update(visible=True, choices=choices, value=choices[0]), gr.update(visible=True), sections


def load_section_ui(choice, sections):
    if not sections or not choice:
        return ""
    idx = next((i for i, s in enumerate(sections) if f"{s['title']} ({s['chars']:,} chars)" == choice), 0)
    return sections[idx]["content"]


def toggle_clone(use_clone):
    return gr.update(visible=use_clone), gr.update(visible=use_clone)


# ==========================================
# SAMPLE TEXT
# ==========================================
SAMPLE = """Chapter 1: The Lighthouse

The old lighthouse stood at the edge of the world. Each morning, Elena climbed one hundred and forty-seven iron steps to the lamp room and watched the sun rise from the sea.

"One day," she whispered to the seagulls, "I'll follow that sun to wherever it goes."

The gulls said nothing. They merely tilted their heads and launched themselves into the wind.

Her grandfather was a man of few words but many stories.

"Tell me about the ships," Elena would say, curling up in the worn armchair by the fire.

And he would smile that slow, careful smile and begin: "There was a ship once, long ago, that sailed beyond the edge of every map. Its captain was a woman with eyes like starlight and a voice that could calm any storm."

"What happened to her?" Elena asked, leaning forward.

"She found what she was looking for," her grandfather said quietly. "But the price was higher than she imagined."

Elena stared into the fire. "Would you pay it? The price, I mean."

The old man was silent for a long time. "I already did," he finally whispered. "I already did."
"""

# ==========================================
# GRADIO UI
# ==========================================
DESCRIPTION = """
# Audiobook Generator
### Self-Hosted TTS with Automatic Emotions

| Mode | What it does |
|------|-------------|
| **Single Speaker** | One voice reads your text with auto-detected emotions. Optional voice cloning. |
| **Multi-Speaker** | AI detects characters, assigns unique voices, adds emotions per line. |

12 languages supported: 10 local (free) + Arabic & Jamaican English (API). Upload PDF/DOCX/TXT or paste text.
"""

with gr.Blocks(title="Audiobook Generator") as demo:

    gr.Markdown(DESCRIPTION)

    # ═══════════════════════════════════════
    # TAB 1: SINGLE SPEAKER
    # ═══════════════════════════════════════
    with gr.Tab("Single Speaker"):
        with gr.Row():
            with gr.Column(scale=1):
                ss_text = gr.Textbox(label="Text", lines=8,
                                     placeholder="Paste text or upload a document below...")
                ss_file = gr.File(label="Upload Document (.txt, .pdf, .docx)",
                                  file_types=[".txt", ".md", ".pdf", ".docx"], type="filepath")

                # PDF section selection
                with gr.Accordion("Select PDF Section (optional)", open=False):
                    ss_detect_btn = gr.Button("Detect Sections", variant="secondary", size="sm")
                    ss_sections_info = gr.Markdown("Upload a PDF then click Detect Sections.")
                    ss_section_choice = gr.Dropdown(choices=[], label="Section", visible=False)
                    ss_load_btn = gr.Button("Load Section", variant="secondary", size="sm", visible=False)

                ss_source_lang = gr.Dropdown(choices=SOURCE_LANGUAGES, value="English",
                                             label="Source Language",
                                             info="Language of your input text")
                ss_target_lang = gr.Dropdown(choices=LANGUAGES, value="Auto",
                                             label="Output Language",
                                             info="Language for the generated speech")
                ss_speaker = gr.Dropdown(choices=SPEAKER_CHOICES,
                                         value="Ryan -- Dynamic male, strong rhythmic drive",
                                         label="Voice",
                                         allow_custom_value=True)

                # Voice cloning (Qwen only — ElevenLabs languages use API voices)
                ss_clone = gr.Checkbox(value=False, label="Use Voice Cloning",
                                       info="Clone a voice from an audio sample (3+ seconds). Qwen languages only.")
                ss_clone_audio = gr.Audio(label="Voice Sample (upload or record, 3+ seconds)",
                                          type="filepath", visible=False,
                                          sources=["upload", "microphone"])
                ss_clone_text = gr.Textbox(label="Transcript of sample (optional, improves quality)",
                                           visible=False, placeholder="What the person says in the audio...")

                ss_sample_btn = gr.Button("Load Sample Text", variant="secondary", size="sm")
                ss_btn = gr.Button("Generate Audiobook", variant="primary", size="lg")

            with gr.Column(scale=1):
                ss_audio = gr.Audio(label="Generated Audiobook", type="filepath")
                ss_stats = gr.Markdown()
                with gr.Accordion("Transcript (with emotions)", open=False):
                    ss_transcript = gr.Markdown()

        # Section detection state
        _ss_sections = gr.State(value=[])

        def on_target_lang_change(lang):
            engine = get_engine(lang)
            if engine == "elevenlabs":
                voices = EL_SPEAKER_CHOICES.get(lang, EL_SPEAKER_CHOICES.get("Arabic", []))
                default = voices[0] if voices else "Rachel -- Calm female"
                return gr.update(choices=voices, value=default)
            else:
                return gr.update(choices=SPEAKER_CHOICES, value="Ryan -- Dynamic male, strong rhythmic drive")

        ss_target_lang.change(fn=on_target_lang_change, inputs=[ss_target_lang],
                              outputs=[ss_speaker])
        ss_detect_btn.click(fn=detect_sections_ui, inputs=[ss_file],
                            outputs=[ss_sections_info, ss_section_choice, ss_load_btn, _ss_sections])
        ss_load_btn.click(fn=load_section_ui, inputs=[ss_section_choice, _ss_sections], outputs=[ss_text])
        ss_clone.change(fn=toggle_clone, inputs=[ss_clone], outputs=[ss_clone_audio, ss_clone_text])
        ss_sample_btn.click(fn=lambda: SAMPLE, outputs=ss_text)
        ss_btn.click(fn=generate_single_speaker,
                     inputs=[ss_text, ss_file, ss_source_lang, ss_target_lang, ss_speaker,
                             ss_clone, ss_clone_audio, ss_clone_text],
                     outputs=[ss_audio, ss_stats, ss_transcript])

    # ═══════════════════════════════════════
    # TAB 2: MULTI-SPEAKER
    # ═══════════════════════════════════════
    with gr.Tab("Multi-Speaker"):
        with gr.Row():
            with gr.Column(scale=1):
                ms_text = gr.Textbox(label="Story Text", lines=8,
                                     placeholder="Paste a story with dialogue...")
                ms_file = gr.File(label="Upload Document (.txt, .pdf, .docx)",
                                  file_types=[".txt", ".md", ".pdf", ".docx"], type="filepath")

                # PDF section selection
                with gr.Accordion("Select PDF Section (optional)", open=False):
                    ms_detect_btn = gr.Button("Detect Sections", variant="secondary", size="sm")
                    ms_sections_info = gr.Markdown("Upload a PDF then click Detect Sections.")
                    ms_section_choice = gr.Dropdown(choices=[], label="Section", visible=False)
                    ms_load_btn = gr.Button("Load Section", variant="secondary", size="sm", visible=False)

                ms_source_lang = gr.Dropdown(choices=SOURCE_LANGUAGES, value="English",
                                             label="Source Language",
                                             info="Language of your input text")
                ms_target_lang = gr.Dropdown(choices=LANGUAGES, value="English",
                                             label="Output Language",
                                             info="Language for the generated speech")

                # Character detection + voice assignment
                with gr.Accordion("Cast — Detect & Assign Voices", open=True):
                    ms_char_btn = gr.Button("Detect Characters", variant="secondary")
                    ms_char_info = gr.Markdown("Enter text then click 'Detect Characters' to see who's in the story.")
                    ms_voices = []
                    for idx in range(8):
                        dd = gr.Dropdown(choices=SPEAKER_CHOICES, label=f"Character {idx+1}",
                                         visible=False, allow_custom_value=True)
                        ms_voices.append(dd)

                ms_sample_btn = gr.Button("Load Sample Story", variant="secondary", size="sm")
                ms_btn = gr.Button("Generate Multi-Speaker Audiobook", variant="primary", size="lg")

            with gr.Column(scale=1):
                ms_audio = gr.Audio(label="Multi-Speaker Audiobook", type="filepath")
                ms_stats = gr.Markdown()
                with gr.Accordion("Transcript (with cast & emotions)", open=False):
                    ms_transcript = gr.Markdown()

        # States
        _ms_sections = gr.State(value=[])
        _ms_characters = gr.State(value=[])

        def detect_chars_ui(text_input, file_input, target_lang):
            text = resolve_text(text_input, file_input)
            client = get_llm_client()
            if not client:
                raise gr.Error("DASHSCOPE_API_KEY needed.")
            characters, _ = detect_characters_and_emotions(client, text)
            engine = get_engine(target_lang)

            # Build separate male/female voice lists
            if engine == "elevenlabs":
                all_voices = EL_SPEAKER_CHOICES.get(target_lang, EL_SPEAKER_CHOICES.get("Arabic", []))
                male_voices = [f"{v['name']} -- {v['desc']}" for v in ELEVENLABS_VOICES.get(target_lang, ELEVENLABS_VOICES.get("Arabic", [])) if v["gender"] == "male"]
                female_voices = [f"{v['name']} -- {v['desc']}" for v in ELEVENLABS_VOICES.get(target_lang, ELEVENLABS_VOICES.get("Arabic", [])) if v["gender"] == "female"]
            else:
                all_voices = SPEAKER_CHOICES
                male_voices = [f"{n} -- {s['desc']}" for n, s in SPEAKERS.items() if s["gender"] == "male"]
                female_voices = [f"{n} -- {s['desc']}" for n, s in SPEAKERS.items() if s["gender"] == "female"]

            info = f"**Detected {len(characters)} characters:**\n\n"
            updates = []
            used_male, used_female = 0, 0

            for i in range(8):
                if i < len(characters):
                    c = characters[i]
                    name, gender = c["name"], c.get("gender", "neutral")
                    info += f"{i+1}. **{name}** ({gender})\n"

                    # Assign unique voice per character — no duplicates
                    if gender == "female" and female_voices:
                        default = female_voices[used_female % len(female_voices)]
                        used_female += 1
                    elif gender == "male" and male_voices:
                        default = male_voices[used_male % len(male_voices)]
                        used_male += 1
                    else:
                        # Neutral/narrator — use a male voice by default
                        default = male_voices[used_male % len(male_voices)] if male_voices else all_voices[0]
                        used_male += 1

                    updates.append(gr.update(visible=True, label=f"{name} ({gender})",
                                             choices=all_voices, value=default))
                else:
                    updates.append(gr.update(visible=False))

            info += "\n*Change any voice below, then click Generate.*"
            return [info, characters] + updates

        ms_char_btn.click(
            fn=detect_chars_ui,
            inputs=[ms_text, ms_file, ms_target_lang],
            outputs=[ms_char_info, _ms_characters] + ms_voices,
        )

        ms_detect_btn.click(fn=detect_sections_ui, inputs=[ms_file],
                            outputs=[ms_sections_info, ms_section_choice, ms_load_btn, _ms_sections])
        ms_load_btn.click(fn=load_section_ui, inputs=[ms_section_choice, _ms_sections], outputs=[ms_text])
        ms_sample_btn.click(fn=lambda: SAMPLE, outputs=ms_text)
        ms_btn.click(fn=generate_multi_speaker,
                     inputs=[ms_text, ms_file, ms_source_lang, ms_target_lang] + ms_voices,
                     outputs=[ms_audio, ms_stats, ms_transcript])

    gr.Markdown(
        "---\n"
        "**How it works:** AI splits your text into emotional segments, generates speech per segment "
        "with matching tone, and assembles into one audiobook. Multi-speaker mode also detects "
        "characters and assigns unique voices by gender.\n\n"
        "**Local TTS (free):** English, Chinese, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian\n\n"
        "**API TTS:** Arabic, English (Jamaican) — requires ELEVENLABS_API_KEY\n\n"
        "**Requires:** DASHSCOPE_API_KEY for emotion analysis, character detection, and translation."
    )

if __name__ == "__main__":
    demo.launch(allowed_paths=[OUTPUT_DIR, tempfile.gettempdir()], ssr_mode=False)