app.py

"""
Visual Storybook / Animated Audiobook Generator
Powered by Qwen3.5-Omni-Plus (narration) + HappyHorse 1.0 (video)

Features:
  - English text input (paste, .txt, .pdf, .docx)
  - Translates + narrates into 36 languages (preset or cloned voice)
  - Generates AI video per scene via HappyHorse 1.0
  - Three video modes: text-to-video, image-to-video, auto scene prompts
  - Two HappyHorse API providers: happyhorse.app, DashScope (Wan 2.1 fallback)
  - Combines narrated audio + video scenes into final MP4

Deploy as a Hugging Face Space:
  1. Create a new Space (SDK: Gradio)
  2. Upload app.py and requirements.txt
  3. Add secrets: DASHSCOPE_API_KEY (required for audio)
     Plus one of: FAL_API_KEY, HAPPYHORSE_API_KEY, or use DashScope for video too
"""

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

import gradio as gr
import requests as http_requests
from openai import OpenAI

# Optional document parsers
try:
    import pypdf
    HAS_PYPDF = True
except ImportError:
    HAS_PYPDF = False

try:
    import docx
    HAS_DOCX = True
except ImportError:
    HAS_DOCX = False


# ==========================================
# CONFIGURATION
# ==========================================
OMNI_MODEL = "qwen3.5-omni-plus"
TTS_VC_MODEL = "qwen3-tts-vc-2026-01-22"
VOICE_CLONE_MODEL = "qwen-voice-enrollment"

DASHSCOPE_BASE_URL = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
DASHSCOPE_API_URL = "https://dashscope-intl.aliyuncs.com/api/v1"
VOICE_CLONE_URL = f"{DASHSCOPE_API_URL}/services/audio/tts/customization"
TTS_SYNTHESIS_URL = f"{DASHSCOPE_API_URL}/services/aigc/multimodal-generation/generation"

# HappyHorse API endpoints per provider
HAPPYHORSE_PROVIDERS = {
    "happyhorse.app": {
        "generate": "https://happyhorse.app/api/generate",
        "status": "https://happyhorse.app/api/status",
        "key_env": "HAPPYHORSE_API_KEY",
        "auth_header": "Authorization",
        "auth_prefix": "Bearer ",
    },
    "DashScope (Bailian)": {
        "generate": f"{DASHSCOPE_API_URL}/services/aigc/video-generation/generation",
        "status": f"{DASHSCOPE_API_URL}/tasks",
        "key_env": "DASHSCOPE_API_KEY",
        "auth_header": "Authorization",
        "auth_prefix": "Bearer ",
    },
}

MAX_CHARS_PER_CHUNK = 1200  # Slightly shorter for scene-based splitting

LANGUAGES = {
    "English": {"code": "en", "native": "English", "tier": "core"},
    "Chinese (Mandarin)": {"code": "zh", "native": "Chinese", "tier": "core"},
    "Japanese": {"code": "ja", "native": "Japanese", "tier": "core"},
    "Korean": {"code": "ko", "native": "Korean", "tier": "core"},
    "German": {"code": "de", "native": "Deutsch", "tier": "core"},
    "French": {"code": "fr", "native": "Francais", "tier": "core"},
    "Russian": {"code": "ru", "native": "Russian", "tier": "core"},
    "Portuguese": {"code": "pt", "native": "Portugues", "tier": "core"},
    "Spanish": {"code": "es", "native": "Espanol", "tier": "core"},
    "Italian": {"code": "it", "native": "Italiano", "tier": "core"},
    "Arabic": {"code": "ar", "native": "Arabic", "tier": "extended"},
    "Dutch": {"code": "nl", "native": "Nederlands", "tier": "extended"},
    "Polish": {"code": "pl", "native": "Polski", "tier": "extended"},
    "Turkish": {"code": "tr", "native": "Turkce", "tier": "extended"},
    "Vietnamese": {"code": "vi", "native": "Tieng Viet", "tier": "extended"},
    "Thai": {"code": "th", "native": "Thai", "tier": "extended"},
    "Indonesian": {"code": "id", "native": "Bahasa Indonesia", "tier": "extended"},
    "Hindi": {"code": "hi", "native": "Hindi", "tier": "extended"},
    "Swahili": {"code": "sw", "native": "Kiswahili", "tier": "extended"},
    "Tamil": {"code": "ta", "native": "Tamil", "tier": "extended"},
}

VOICE_CLONE_LANGUAGES = {
    "English", "Chinese (Mandarin)", "Japanese", "Korean", "German",
    "French", "Russian", "Portuguese", "Spanish", "Italian",
}

PRESET_VOICES = [
    "Cherry -- Sunny, friendly",
    "Serena -- Gentle, soft",
    "Jennifer -- Cinematic narrator",
    "Katerina -- Mature, rich rhythm",
    "Ethan -- Warm, energetic",
    "Ryan -- Dramatic, rhythmic",
    "Kai -- Soothing, calm",
    "Neil -- Precise, clear",
    "Lenn -- Rational, steady",
    "Eldric Sage -- Authoritative narrator",
    "Arthur -- Classic, mature",
    "Bella -- Elegant, warm",
]


def get_voice_name(label):
    return label.split("--")[0].strip()


# ==========================================
# AUDIO HELPERS
# ==========================================
def base64_to_wav(b64_data, output_path):
    audio_bytes = base64.b64decode(b64_data)
    sr, nc, bps = 24000, 1, 16
    br = sr * nc * bps // 8
    ba = nc * bps // 8
    ds = len(audio_bytes)
    with open(output_path, "wb") as f:
        f.write(b"RIFF")
        f.write(struct.pack("<I", 36 + ds))
        f.write(b"WAVE")
        f.write(b"fmt ")
        f.write(struct.pack("<I", 16))
        f.write(struct.pack("<H", 1))
        f.write(struct.pack("<H", nc))
        f.write(struct.pack("<I", sr))
        f.write(struct.pack("<I", br))
        f.write(struct.pack("<H", ba))
        f.write(struct.pack("<H", bps))
        f.write(b"data")
        f.write(struct.pack("<I", ds))
        f.write(audio_bytes)


def get_audio_duration(filepath):
    result = subprocess.run(
        ["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
         "-of", "default=noprint_wrappers=1:nokey=1", filepath],
        capture_output=True, text=True,
    )
    return float(result.stdout.strip())


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


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


# ==========================================
# DOCUMENT EXTRACTION
# ==========================================
def extract_text_from_file(filepath):
    ext = os.path.splitext(filepath)[1].lower()
    if ext == ".pdf":
        if not HAS_PYPDF:
            raise gr.Error("pypdf not installed.")
        reader = pypdf.PdfReader(filepath)
        return "\n\n".join(p.extract_text().strip() for p in reader.pages if p.extract_text())
    elif ext in (".docx", ".doc"):
        if ext == ".doc":
            raise gr.Error("Please save as .docx or .pdf.")
        if not HAS_DOCX:
            raise gr.Error("python-docx not installed.")
        doc = docx.Document(filepath)
        return "\n\n".join(p.text.strip() for p in doc.paragraphs if p.text.strip())
    else:
        with open(filepath, "r", encoding="utf-8", errors="replace") as f:
            return f.read()


# ==========================================
# TEXT SPLITTING INTO SCENES
# ==========================================
def split_into_scenes(text, max_chars=MAX_CHARS_PER_CHUNK):
    text = text.strip()
    if not text:
        return []
    if len(text) <= max_chars:
        return [text]

    chunks = []
    paragraphs = re.split(r"\n\s*\n", text)
    current = ""

    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
        if len(current) + len(para) + 2 <= max_chars:
            current = (current + "\n\n" + para).strip()
        else:
            if current:
                chunks.append(current)
            if len(para) > max_chars:
                sentences = re.split(r"(?<=[.!?])\s+", para)
                current = ""
                for s in sentences:
                    if len(current) + len(s) + 1 <= max_chars:
                        current = (current + " " + s).strip()
                    else:
                        if current:
                            chunks.append(current)
                        current = s
            else:
                current = para

    if current:
        chunks.append(current)
    return chunks


# ==========================================
# SCENE PROMPT GENERATION (via Qwen)
# ==========================================
def generate_scene_prompt(client, scene_text, scene_index):
    """Use Qwen to create a cinematic video prompt from story text."""
    try:
        response = client.chat.completions.create(
            model=OMNI_MODEL,
            modalities=["text"],
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are a cinematic scene director. Given a passage from a story, "
                        "create a single vivid video generation prompt (max 200 words) that "
                        "describes the visual scene. Include: setting, lighting, camera angle, "
                        "character actions, mood. Do NOT include dialogue or text overlays. "
                        "Output ONLY the video prompt, nothing else."
                    ),
                },
                {
                    "role": "user",
                    "content": f"Create a cinematic video prompt for this scene:\n\n{scene_text[:1500]}",
                },
            ],
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        # Fallback: use first 200 chars as prompt
        return scene_text[:200]


# ==========================================
# VOICE CLONING
# ==========================================
def prepare_clone_audio(audio_path):
    result = subprocess.run(
        ["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
         "-of", "default=noprint_wrappers=1:nokey=1", audio_path],
        capture_output=True, text=True,
    )
    duration = float(result.stdout.strip())

    if duration < 10:
        raise ValueError(f"Audio too short ({duration:.1f}s). Need at least 10 seconds.")

    tmp_prepared = audio_path + "_prepared.wav"
    if duration <= 60:
        subprocess.run(
            ["ffmpeg", "-y", "-i", audio_path, "-ar", "24000", "-ac", "1",
             "-acodec", "pcm_s16le", tmp_prepared],
            capture_output=True, check=True,
        )
    else:
        start = min(5, duration - 60)
        subprocess.run(
            ["ffmpeg", "-y", "-ss", str(start), "-t", "60", "-i", audio_path,
             "-ar", "24000", "-ac", "1", "-acodec", "pcm_s16le", tmp_prepared],
            capture_output=True, check=True,
        )
    return tmp_prepared


def clone_voice(audio_path, api_key):
    prepared = prepare_clone_audio(audio_path)
    b64 = base64.b64encode(pathlib.Path(prepared).read_bytes()).decode()
    try:
        os.remove(prepared)
    except OSError:
        pass

    resp = http_requests.post(VOICE_CLONE_URL, json={
        "model": VOICE_CLONE_MODEL,
        "input": {
            "action": "create",
            "target_model": TTS_VC_MODEL,
            "preferred_name": "storybook_voice",
            "audio": {"data": f"data:audio/wav;base64,{b64}"},
        },
    }, headers={
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }, timeout=60)

    if resp.status_code != 200:
        raise RuntimeError(f"Voice clone failed: {resp.text[:300]}")
    return resp.json()["output"]["voice"]


# ==========================================
# AUDIO NARRATION
# ==========================================
def narrate_scene_preset(client, text, voice, language, lang_config, translate):
    """Narrate with preset voice via Qwen3.5-Omni-Plus. Returns (audio_b64_parts, transcript)."""
    if translate and language != "English":
        sys_prompt = (
            f"You are a professional audiobook narrator. "
            f"Translate the English text into {language} ({lang_config['native']}) "
            f"and narrate it expressively. Respond ONLY with the spoken {language} narration."
        )
        user_text = f"Translate into {language} and narrate:\n\n{text}"
    else:
        sys_prompt = "You are a professional audiobook narrator. Read expressively. Respond ONLY with narration."
        user_text = f"Narrate:\n\n{text}"

    completion = client.chat.completions.create(
        model=OMNI_MODEL,
        messages=[
            {"role": "system", "content": sys_prompt},
            {"role": "user", "content": user_text},
        ],
        modalities=["text", "audio"],
        audio={"voice": voice, "format": "wav"},
        stream=True,
        stream_options={"include_usage": True},
    )

    audio_parts, text_parts = [], []
    for event in completion:
        if not event.choices:
            continue
        delta = event.choices[0].delta
        if hasattr(delta, "content") and delta.content:
            text_parts.append(delta.content)
        if hasattr(delta, "audio") and delta.audio:
            if isinstance(delta.audio, dict) and "data" in delta.audio:
                audio_parts.append(delta.audio["data"])
            elif hasattr(delta.audio, "data") and delta.audio.data:
                audio_parts.append(delta.audio.data)

    return "".join(audio_parts), "".join(text_parts)


def narrate_scene_cloned(client, text, voice_id, language, lang_config, translate, api_key):
    """Translate then synthesize with cloned voice. Returns (audio_url, transcript)."""
    final_text = text
    if translate and language != "English":
        resp = client.chat.completions.create(
            model=OMNI_MODEL, modalities=["text"],
            messages=[
                {"role": "system", "content": f"Translate English to {language}. Output ONLY the translation."},
                {"role": "user", "content": f"Translate:\n\n{text}"},
            ],
        )
        final_text = resp.choices[0].message.content.strip()

    lang_map = {
        "English": "English", "Chinese (Mandarin)": "Chinese", "Japanese": "Japanese",
        "Korean": "Korean", "German": "German", "French": "French",
        "Russian": "Russian", "Portuguese": "Portuguese", "Spanish": "Spanish", "Italian": "Italian",
    }

    resp = http_requests.post(TTS_SYNTHESIS_URL, json={
        "model": TTS_VC_MODEL,
        "input": {"text": final_text, "voice": voice_id, "language_type": lang_map.get(language, "English")},
    }, headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, timeout=120)

    if resp.status_code != 200:
        return None, final_text

    audio_url = resp.json().get("output", {}).get("audio", {}).get("url")
    return audio_url, final_text


# ==========================================
# HAPPYHORSE VIDEO GENERATION
# ==========================================
def sanitize_prompt(prompt, max_length=2000):
    """Clean and truncate prompt for HappyHorse API."""
    # Replace smart quotes and special chars
    prompt = prompt.replace("\u2018", "'").replace("\u2019", "'")
    prompt = prompt.replace("\u201c", '"').replace("\u201d", '"')
    prompt = prompt.replace("\u2014", " -- ").replace("\u2013", " - ")
    # Remove any control characters
    prompt = re.sub(r'[\x00-\x1f\x7f-\x9f]', ' ', prompt)
    # Collapse multiple spaces
    prompt = re.sub(r'\s+', ' ', prompt).strip()
    # Truncate
    if len(prompt) > max_length:
        prompt = prompt[:max_length].rsplit(' ', 1)[0]
    return prompt


def generate_video_happyhorse_app(prompt, api_key, duration=5, aspect="16:9", image_url=None):
    """Generate video via happyhorse.app."""
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}

    clean_prompt = sanitize_prompt(prompt, max_length=2000)

    payload = {
        "model": "happyhorse-1.0/video",
        "prompt": clean_prompt,
        "mode": "pro",
        "duration": int(duration),
        "aspect_ratio": aspect,
        "sound": False,
    }
    if image_url and not image_url.startswith("data:"):
        payload["image_urls"] = [image_url]

    generate_url = HAPPYHORSE_PROVIDERS["happyhorse.app"]["generate"]
    print(f"[HappyHorse] Submitting to {generate_url}")
    print(f"[HappyHorse] Payload: {json.dumps(payload)[:300]}")

    resp = http_requests.post(generate_url, json=payload, headers=headers, timeout=60)
    print(f"[HappyHorse] Submit response: {resp.status_code}")

    if resp.status_code != 200:
        resp_text = resp.text[:500]
        print(f"[HappyHorse] Error: {resp_text}")
        raise RuntimeError(f"happyhorse.app submit failed ({resp.status_code}): {resp_text}")

    resp_data = resp.json()
    print(f"[HappyHorse] Response data: {json.dumps(resp_data)[:300]}")

    # Handle different response structures
    task_id = (
        resp_data.get("data", {}).get("task_id")
        or resp_data.get("task_id")
        or resp_data.get("id")
    )
    if not task_id:
        raise RuntimeError(f"No task_id in response: {json.dumps(resp_data)[:300]}")

    print(f"[HappyHorse] Task ID: {task_id}")

    # Poll for completion
    status_url = HAPPYHORSE_PROVIDERS["happyhorse.app"]["status"]
    for attempt in range(90):  # ~15 minutes max
        time.sleep(10)
        try:
            s = http_requests.get(
                f"{status_url}?task_id={task_id}",
                headers=headers, timeout=30,
            ).json()
        except Exception as e:
            print(f"[HappyHorse] Poll error (attempt {attempt}): {e}")
            continue

        # Handle different status response structures
        status_data = s.get("data", s)
        status = (
            status_data.get("status", "")
            or s.get("status", "")
        )
        print(f"[HappyHorse] Poll {attempt}: status={status}")

        if status in ("SUCCESS", "COMPLETED", "completed", "succeeded"):
            # Try different result URL paths
            urls = (
                status_data.get("response", {}).get("resultUrls", [])
                or status_data.get("resultUrls", [])
                or [status_data.get("video_url")]
                or [status_data.get("output", {}).get("video_url")]
            )
            urls = [u for u in urls if u]  # Filter None
            if urls:
                print(f"[HappyHorse] Video URL: {urls[0][:100]}")
                return urls[0]
            raise RuntimeError(f"Status SUCCESS but no video URL found in: {json.dumps(s)[:300]}")

        elif status in ("FAILED", "failed", "error"):
            error_msg = (
                status_data.get("error_message", "")
                or status_data.get("error", "")
                or "Unknown error"
            )
            raise RuntimeError(f"happyhorse.app failed: {error_msg}")

    raise RuntimeError("happyhorse.app timeout after 15 minutes")


def generate_video_dashscope(prompt, api_key, duration=5, aspect="16:9", image_url=None):
    """Generate video via DashScope Bailian (async task API).
    Uses wan2.1-t2v-plus since happyhorse-1.0 is not yet available on DashScope.
    """
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "X-DashScope-Async": "enable",
    }

    clean_prompt = sanitize_prompt(prompt, max_length=2000)

    # Use Wan 2.1 T2V Plus (available on DashScope) as HappyHorse is not live yet
    if image_url and not image_url.startswith("data:"):
        model = "wan2.1-i2v-plus"
        payload = {
            "model": model,
            "input": {"prompt": clean_prompt, "img_url": image_url},
            "parameters": {"resolution": "1280*720"},
        }
    else:
        model = "wan2.1-t2v-plus"
        payload = {
            "model": model,
            "input": {"prompt": clean_prompt},
            "parameters": {"size": "1280*720"},
        }

    generate_url = HAPPYHORSE_PROVIDERS["DashScope (Bailian)"]["generate"]
    print(f"[DashScope] Submitting to {generate_url}")
    print(f"[DashScope] Prompt: {prompt[:100]}...")

    resp = http_requests.post(generate_url, json=payload, headers=headers, timeout=60)
    print(f"[DashScope] Submit response: {resp.status_code}")

    if resp.status_code != 200:
        resp_text = resp.text[:500]
        print(f"[DashScope] Error: {resp_text}")
        raise RuntimeError(f"DashScope submit failed ({resp.status_code}): {resp_text}")

    resp_data = resp.json()
    print(f"[DashScope] Response: {json.dumps(resp_data)[:300]}")

    task_id = resp_data.get("output", {}).get("task_id")
    if not task_id:
        raise RuntimeError(f"No task_id from DashScope: {json.dumps(resp_data)[:300]}")

    print(f"[DashScope] Task ID: {task_id}")

    for attempt in range(90):
        time.sleep(10)
        try:
            s = http_requests.get(
                f"{HAPPYHORSE_PROVIDERS['DashScope (Bailian)']['status']}/{task_id}",
                headers={"Authorization": f"Bearer {api_key}"}, timeout=30,
            ).json()
        except Exception as e:
            print(f"[DashScope] Poll error (attempt {attempt}): {e}")
            continue

        status = s.get("output", {}).get("task_status", "")
        print(f"[DashScope] Poll {attempt}: status={status}")

        if status == "SUCCEEDED":
            results = s.get("output", {}).get("results", [])
            if results:
                video_url = results[0].get("url")
                print(f"[DashScope] Video URL: {video_url[:100] if video_url else 'None'}")
                return video_url
            raise RuntimeError(f"SUCCEEDED but no results: {json.dumps(s)[:300]}")
        elif status == "FAILED":
            msg = s.get("output", {}).get("message", "Unknown error")
            raise RuntimeError(f"DashScope failed: {msg}")

    raise RuntimeError("DashScope timeout after 15 minutes")


def generate_video(prompt, provider, api_key, duration=5, aspect="16:9", image_url=None):
    """Route to the correct provider."""
    if provider == "happyhorse.app":
        return generate_video_happyhorse_app(prompt, api_key, duration, aspect, image_url)
    elif provider == "DashScope (Bailian)":
        return generate_video_dashscope(prompt, api_key, duration, aspect, image_url)
    else:
        raise ValueError(f"Unknown provider: {provider}")


def download_video(url, output_path):
    """Download a video from URL."""
    resp = http_requests.get(url, timeout=120, stream=True)
    if resp.status_code != 200:
        raise RuntimeError(f"Failed to download video: {resp.status_code}")
    with open(output_path, "wb") as f:
        for chunk in resp.iter_content(chunk_size=8192):
            f.write(chunk)
    return output_path


# ==========================================
# VIDEO + AUDIO COMPOSITING
# ==========================================
def create_scene_video(video_path, audio_path, output_path):
    """Combine a video clip with audio. Loop/trim video to match audio duration."""
    audio_dur = get_audio_duration(audio_path)

    # Create a video that matches audio length (loop if needed, trim if longer)
    subprocess.run([
        "ffmpeg", "-y",
        "-stream_loop", "-1", "-i", video_path,
        "-i", audio_path,
        "-map", "0:v:0", "-map", "1:a:0",
        "-c:v", "libx264", "-preset", "fast", "-crf", "23",
        "-c:a", "aac", "-b:a", "128k",
        "-t", str(audio_dur),
        "-shortest",
        output_path,
    ], capture_output=True, check=True)
    return output_path


def concatenate_videos(video_files, output_path):
    """Concatenate multiple MP4 files."""
    if len(video_files) == 1:
        shutil.copy2(video_files[0], output_path)
        return

    # Re-encode to ensure compatible streams
    normalized = []
    for i, vf in enumerate(video_files):
        norm = vf.replace(".mp4", f"_norm_{i}.mp4")
        subprocess.run([
            "ffmpeg", "-y", "-i", vf,
            "-c:v", "libx264", "-preset", "fast", "-crf", "23",
            "-c:a", "aac", "-b:a", "128k",
            "-ar", "24000", "-ac", "1",
            "-r", "24",
            "-s", "1280x720",
            norm,
        ], capture_output=True, check=True)
        normalized.append(norm)

    list_file = output_path + ".txt"
    with open(list_file, "w") as f:
        for vf in normalized:
            f.write(f"file '{vf}'\n")

    subprocess.run([
        "ffmpeg", "-y", "-f", "concat", "-safe", "0",
        "-i", list_file, "-c", "copy", output_path,
    ], capture_output=True, check=True)
    os.remove(list_file)


# ==========================================
# MAIN PIPELINE
# ==========================================
def generate_storybook(
    text_input, file_input, target_language, voice_mode, preset_voice_label,
    clone_audio, video_provider, video_mode, scene_images, aspect_ratio,
    video_duration, progress=gr.Progress(),
):
    # -- Resolve text --
    if file_input is not None:
        progress(0.01, desc="Extracting text...")
        text = extract_text_from_file(file_input)
    elif text_input and text_input.strip():
        text = text_input.strip()
    else:
        raise gr.Error("Please provide text or upload a file.")

    if len(text) < 20:
        raise gr.Error("Text is too short for a storybook.")

    # -- API keys --
    ds_key = os.environ.get("DASHSCOPE_API_KEY", "")
    if not ds_key:
        raise gr.Error("DASHSCOPE_API_KEY not set (needed for audio narration).")

    provider_config = HAPPYHORSE_PROVIDERS[video_provider]
    video_key = os.environ.get(provider_config["key_env"], "")
    if not video_key:
        raise gr.Error(f"{provider_config['key_env']} not set. Add it in Settings > Secrets.")

    lang_config = LANGUAGES[target_language]
    use_clone = voice_mode == "Clone a Voice"
    translate = target_language != "English"
    client = OpenAI(api_key=ds_key, base_url=DASHSCOPE_BASE_URL)
    tmp_dir = tempfile.mkdtemp(prefix="storybook_")

    # -- Voice cloning --
    cloned_voice_id = None
    if use_clone:
        if clone_audio is None:
            raise gr.Error("Upload a voice sample for cloning.")
        if target_language not in VOICE_CLONE_LANGUAGES:
            raise gr.Error(f"Voice cloning only supports: {', '.join(sorted(VOICE_CLONE_LANGUAGES))}")
        progress(0.03, desc="Cloning voice...")
        cloned_voice_id = clone_voice(clone_audio, ds_key)

    # -- Parse scene images if provided --
    image_list = []
    if scene_images:
        image_list = scene_images if isinstance(scene_images, list) else [scene_images]

    try:
        # -- Split into scenes --
        progress(0.05, desc="Splitting text into scenes...")
        scenes = split_into_scenes(text)
        total = len(scenes)
        scene_videos = []
        all_transcripts = []

        for i, scene_text in enumerate(scenes):
            base_frac = 0.08 + 0.85 * (i / total)

            # --- Step 1: Generate video prompt ---
            progress(base_frac, desc=f"Scene {i+1}/{total}: Creating video prompt...")
            if video_mode == "Text-to-Video (auto scene prompts)":
                video_prompt = generate_scene_prompt(client, scene_text, i)
                image_url = None
            elif video_mode == "Image-to-Video (animate uploaded images)":
                video_prompt = generate_scene_prompt(client, scene_text, i)
                # Use matching image if available, otherwise cycle through
                if image_list:
                    img_index = i % len(image_list)
                    img_path = image_list[img_index]
                    # Upload image and get URL -- for now, encode as data URI
                    # Note: some providers need a public URL; we base64 encode
                    img_b64 = base64.b64encode(pathlib.Path(img_path).read_bytes()).decode()
                    ext = pathlib.Path(img_path).suffix.lower()
                    mime = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".webp": "image/webp"}.get(ext, "image/jpeg")
                    image_url = f"data:{mime};base64,{img_b64}"
                else:
                    image_url = None
            else:  # Both
                video_prompt = generate_scene_prompt(client, scene_text, i)
                if image_list and i < len(image_list):
                    img_path = image_list[i]
                    img_b64 = base64.b64encode(pathlib.Path(img_path).read_bytes()).decode()
                    ext = pathlib.Path(img_path).suffix.lower()
                    mime = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png"}.get(ext, "image/jpeg")
                    image_url = f"data:{mime};base64,{img_b64}"
                else:
                    image_url = None

            # --- Step 2: Generate video ---
            progress(base_frac + 0.02, desc=f"Scene {i+1}/{total}: Generating video...")
            video_error = None
            try:
                video_url = generate_video(
                    video_prompt, video_provider, video_key,
                    duration=video_duration, aspect=aspect_ratio, image_url=image_url,
                )
                scene_video_path = os.path.join(tmp_dir, f"scene_{i:03d}_video.mp4")
                if video_url:
                    download_video(video_url, scene_video_path)
                    # Verify the downloaded file is a valid video
                    file_size = os.path.getsize(scene_video_path)
                    if file_size < 1000:
                        video_error = f"Downloaded file too small ({file_size} bytes) - likely not a video"
                        scene_video_path = None
                else:
                    video_error = "No video URL returned from API"
                    scene_video_path = None
            except Exception as e:
                video_error = str(e)
                scene_video_path = None

            if scene_video_path is None:
                # Create a placeholder with text overlay showing the error
                scene_video_path = os.path.join(tmp_dir, f"scene_{i:03d}_placeholder.mp4")
                # Escape special chars for ffmpeg drawtext
                safe_prompt = video_prompt[:80].replace("'", "").replace('"', '').replace(':', ' ')
                subprocess.run([
                    "ffmpeg", "-y", "-f", "lavfi", "-i",
                    f"color=c=0x1a1a2e:s=1280x720:d={video_duration}:r=24",
                    "-vf", f"drawtext=text='Scene {i+1}':fontsize=48:fontcolor=white:x=(w-text_w)/2:y=(h-text_h)/2-40,"
                           f"drawtext=text='{safe_prompt[:60]}':fontsize=20:fontcolor=0xaaaaaa:x=(w-text_w)/2:y=(h-text_h)/2+30",
                    "-c:v", "libx264", "-preset", "fast",
                    scene_video_path,
                ], capture_output=True, check=True)
                all_transcripts.append(f"**Scene {i+1} VIDEO ERROR:** {video_error}")

            # --- Step 3: Generate narration audio ---
            progress(base_frac + 0.04, desc=f"Scene {i+1}/{total}: Narrating...")
            scene_audio_path = os.path.join(tmp_dir, f"scene_{i:03d}_audio.wav")

            try:
                if use_clone:
                    audio_url, transcript = narrate_scene_cloned(
                        client, scene_text, cloned_voice_id, target_language, lang_config, translate, ds_key,
                    )
                    if audio_url:
                        download_video(audio_url, scene_audio_path)
                    else:
                        generate_silence(5, scene_audio_path)
                else:
                    voice = get_voice_name(preset_voice_label)
                    audio_b64, transcript = narrate_scene_preset(
                        client, scene_text, voice, target_language, lang_config, translate,
                    )
                    if audio_b64:
                        base64_to_wav(audio_b64, scene_audio_path)
                    else:
                        generate_silence(5, scene_audio_path)
                        transcript = "No audio generated"

                if transcript:
                    all_transcripts.append(f"**Scene {i+1}:** {transcript[:200]}...")
            except Exception as e:
                generate_silence(5, scene_audio_path)
                all_transcripts.append(f"Scene {i+1} narration failed: {str(e)[:100]}")

            # --- Step 4: Combine video + audio for this scene ---
            progress(base_frac + 0.06, desc=f"Scene {i+1}/{total}: Compositing...")
            scene_final = os.path.join(tmp_dir, f"scene_{i:03d}_final.mp4")
            try:
                create_scene_video(scene_video_path, scene_audio_path, scene_final)
                scene_videos.append(scene_final)
            except Exception as e:
                all_transcripts.append(f"Scene {i+1} compositing failed: {str(e)[:100]}")

        if not scene_videos:
            raise gr.Error("No scenes were successfully generated.")

        # -- Final assembly --
        progress(0.95, desc="Assembling final storybook video...")
        final_video = os.path.join(tmp_dir, "storybook.mp4")
        concatenate_videos(scene_videos, final_video)

        progress(1.0, desc="Done!")

        # Stats
        video_size = os.path.getsize(final_video) / (1024 * 1024)
        stats = (
            f"**Visual Storybook Generated!**\n\n"
            f"- **Scenes:** {len(scene_videos)} / {total}\n"
            f"- **Language:** {target_language} ({lang_config['native']})\n"
            f"- **Voice:** {'Cloned' if use_clone else preset_voice_label}\n"
            f"- **Video Provider:** {video_provider}\n"
            f"- **File size:** {video_size:.1f} MB\n"
        )

        transcript_text = "\n\n---\n\n".join(all_transcripts) if all_transcripts else ""

        return final_video, stats, transcript_text

    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Pipeline error: {str(e)}")


# ==========================================
# GRADIO UI
# ==========================================
SAMPLE_TEXT = """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 like a great golden coin.

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

Her grandfather was a man of few words but many stories. He kept them locked away like treasures, only bringing them out on winter nights when storms howled outside.

Chapter 2: The Storm

The storm came without warning. Dark clouds swallowed the horizon and the sea rose in great heaving walls of grey-green water. Elena pressed her face to the glass of the lamp room and watched lightning split the sky.

Through the rain, she saw something impossible - a ship, its sails torn to ribbons, riding the crest of a monstrous wave. And on its deck, a figure stood perfectly still, as if the fury of the ocean meant nothing at all.

Chapter 3: The Stranger

He arrived at dawn, walking up the rocky path as if he had always known the way. His clothes were dry despite the storm, and his eyes held the same changing colors as the sea itself.

"I've been looking for this lighthouse," he said simply. "I've been looking for a very long time."

Elena's grandfather stood in the doorway, his weathered face unreadable. Then slowly, he stepped aside and let the stranger in."""

DESCRIPTION = """
# Visual Storybook / Animated Audiobook
### English Text to Multi-Language Narrated Video
**Powered by Qwen3.5-Omni-Plus (narration) + HappyHorse 1.0 (video)**

Upload English text and generate a **narrated video storybook**:
- AI generates cinematic video scenes from your story
- Professional narration in 36 languages (preset or cloned voice)
- Each scene gets its own video + synchronized audio
- All scenes assembled into one final MP4

| Component | Model | What it does |
|-----------|-------|-------------|
| **Narration** | Qwen3.5-Omni-Plus | Translates + speaks the story |
| **Video** | HappyHorse 1.0 | Generates cinematic scene videos |
| **Scene Prompts** | Qwen3.5-Omni-Plus | Auto-creates video prompts from text |
"""

lang_choices = []
for name, cfg in LANGUAGES.items():
    if cfg["tier"] == "core":
        lang_choices.append(f"* {name}")
for name, cfg in LANGUAGES.items():
    if cfg["tier"] == "extended":
        lang_choices.append(name)


def clean_lang(choice):
    return choice.replace("* ", "").strip()


def on_voice_mode_change(mode):
    if mode == "Clone a Voice":
        return gr.update(visible=False), gr.update(visible=True)
    return gr.update(visible=True), gr.update(visible=False)


def on_video_mode_change(mode):
    show_images = "Image" in mode or "Both" in mode
    return gr.update(visible=show_images)


def generate_wrapper(text_input, file_input, lang, voice_mode, preset_voice,
                     clone_audio, video_provider, video_mode, scene_images,
                     aspect, duration, progress=gr.Progress()):
    language = clean_lang(lang)
    imgs = None
    if scene_images:
        imgs = [f.name if hasattr(f, 'name') else f for f in scene_images] if isinstance(scene_images, list) else scene_images
    return generate_storybook(
        text_input, file_input, language, voice_mode, preset_voice,
        clone_audio, video_provider, video_mode, imgs, aspect, duration, progress,
    )


with gr.Blocks(
    title="Visual Storybook Generator",
) as demo:

    gr.Markdown(DESCRIPTION)

    with gr.Row():
        # -- LEFT: Inputs --
        with gr.Column(scale=1):
            with gr.Tab("Story"):
                text_input = gr.Textbox(label="English Text", placeholder="Paste your story...", lines=8, max_lines=20)
                file_input = gr.File(
                    label="Or Upload (.txt, .md, .pdf, .docx)",
                    file_types=[".txt", ".md", ".pdf", ".docx"],
                    type="filepath",
                )
                sample_btn = gr.Button("Load Sample Story", variant="secondary", size="sm")

            with gr.Tab("Voice"):
                target_lang = gr.Dropdown(choices=lang_choices, value="* English", label="Target Language")
                voice_mode = gr.Radio(choices=["Preset Voice", "Clone a Voice"], value="Preset Voice", label="Voice Mode")
                preset_voice = gr.Dropdown(choices=PRESET_VOICES, value="Jennifer -- Cinematic narrator", label="Narrator Voice", visible=True)
                clone_audio = gr.Audio(label="Voice Sample (10s-3min)", type="filepath", visible=False)

            with gr.Tab("Video"):
                video_provider = gr.Dropdown(
                    choices=list(HAPPYHORSE_PROVIDERS.keys()),
                    value="happyhorse.app",
                    label="HappyHorse API Provider",
                    info="Each provider needs its own API key in Secrets",
                )
                video_mode = gr.Dropdown(
                    choices=[
                        "Text-to-Video (auto scene prompts)",
                        "Image-to-Video (animate uploaded images)",
                        "Both (images where available, text for rest)",
                    ],
                    value="Text-to-Video (auto scene prompts)",
                    label="Video Generation Mode",
                )
                scene_images = gr.File(
                    label="Scene Images (one per scene, optional)",
                    file_types=["image"],
                    file_count="multiple",
                    visible=False,
                )
                with gr.Row():
                    aspect_ratio = gr.Dropdown(choices=["16:9", "9:16", "1:1"], value="16:9", label="Aspect Ratio")
                    video_duration = gr.Slider(minimum=3, maximum=10, value=5, step=1, label="Video Duration (sec/scene)")

            generate_btn = gr.Button("Generate Visual Storybook", variant="primary", size="lg")
            test_btn = gr.Button("Test Video API Connection", variant="secondary", size="sm")

        # -- RIGHT: Output --
        with gr.Column(scale=1):
            video_output = gr.Video(label="Generated Storybook Video")
            stats_output = gr.Markdown(label="Stats")
            with gr.Accordion("Scene Transcripts", open=False):
                transcript_output = gr.Markdown()
            test_output = gr.Markdown(label="API Test Result", visible=True)

    # Events
    sample_btn.click(fn=lambda: SAMPLE_TEXT, outputs=text_input)
    voice_mode.change(fn=on_voice_mode_change, inputs=voice_mode, outputs=[preset_voice, clone_audio])
    video_mode.change(fn=on_video_mode_change, inputs=video_mode, outputs=[scene_images])

    def test_video_api(provider):
        """Test the video API connection with a simple request."""
        config = HAPPYHORSE_PROVIDERS[provider]
        api_key = os.environ.get(config["key_env"], "")

        if not api_key:
            return f"**FAILED:** `{config['key_env']}` is not set in Secrets."

        results = []
        results.append(f"**Provider:** {provider}")
        results.append(f"**API Key:** `{api_key[:8]}...{api_key[-4:]}`")

        # Step 1: Submit a short test generation
        if provider == "happyhorse.app":
            headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
            payload = {
                "model": "happyhorse-1.0/video",
                "prompt": "A golden sunset over calm ocean waves",
                "mode": "std",
                "duration": 3,
                "aspect_ratio": "16:9",
                "sound": False,
            }
            url = config["generate"]
        else:
            headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json", "X-DashScope-Async": "enable"}
            payload = {
                "model": "wan2.1-t2v-plus",
                "input": {"prompt": "A golden sunset over calm ocean waves"},
                "parameters": {"size": "1280*720"},
            }
            url = config["generate"]

        results.append(f"**URL:** `{url}`")
        results.append(f"**Payload:** `{json.dumps(payload)}`")

        try:
            resp = http_requests.post(url, json=payload, headers=headers, timeout=30)
            results.append(f"**HTTP Status:** {resp.status_code}")
            results.append(f"**Response:** ```{resp.text[:500]}```")

            if resp.status_code == 200:
                data = resp.json()
                task_id = (
                    data.get("data", {}).get("task_id")
                    or data.get("task_id")
                    or data.get("output", {}).get("task_id")
                )
                if task_id:
                    results.append(f"**Task ID:** `{task_id}` -- API is working! Task submitted successfully.")

                    # Quick poll once
                    time.sleep(5)
                    if provider == "happyhorse.app":
                        s_resp = http_requests.get(
                            f"{config['status']}?task_id={task_id}",
                            headers=headers, timeout=15,
                        )
                    else:
                        s_resp = http_requests.get(
                            f"{config['status']}/{task_id}",
                            headers={"Authorization": f"Bearer {api_key}"}, timeout=15,
                        )
                    results.append(f"**Status poll:** {s_resp.status_code}")
                    results.append(f"**Status body:** ```{s_resp.text[:500]}```")
                else:
                    results.append(f"**WARNING:** 200 OK but no task_id found in response")
            elif resp.status_code == 401:
                results.append("**ERROR:** Invalid API key (401 Unauthorized)")
            elif resp.status_code == 402:
                results.append("**ERROR:** Insufficient credits (402 Payment Required)")
            else:
                results.append(f"**ERROR:** Unexpected status {resp.status_code}")

        except Exception as e:
            results.append(f"**EXCEPTION:** {str(e)}")

        return "\n\n".join(results)

    test_btn.click(fn=test_video_api, inputs=[video_provider], outputs=[test_output])

    generate_btn.click(
        fn=generate_wrapper,
        inputs=[text_input, file_input, target_lang, voice_mode, preset_voice,
                clone_audio, video_provider, video_mode, scene_images,
                aspect_ratio, video_duration],
        outputs=[video_output, stats_output, transcript_output],
    )

    gr.Markdown(
        "---\n"
        "**Pipeline:** Text > Split into scenes > Qwen generates video prompts > "
        "HappyHorse generates video per scene > Qwen narrates audio per scene > "
        "FFmpeg composites video+audio > Concatenates all scenes into final MP4\n\n"
        "**API Keys needed:** DASHSCOPE_API_KEY (audio) + HAPPYHORSE_API_KEY (if using happyhorse.app) "
        "or just DASHSCOPE_API_KEY for both audio and video (DashScope provider)\n\n"
        "Built with Gradio | Narration by Qwen | Video by HappyHorse 1.0"
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="amber", secondary_hue="orange", neutral_hue="stone"))