"""JoyAI-Echo x LTX-2.3 surgical merge — multi-shot narrated video (video + audio).

This Space runs the official JoyAI-Echo inference graph
(`ltx-core` / `ltx-pipelines` / `ltx-distillation`) with the community
"surgical merge" checkpoint
`joeygambino/joyai-echo-ltx23-echoVid-ltxAud-surgical` dropped in as the single
`model_file`. The merge keeps JoyAI-Echo's video / conditioning branch (its
slot-paired cross-modal memory bank, which holds a character's identity across
shots) and LTX-2.3-distilled-1.1's audio branch (natural voice + lip-sync).

The fp8 checkpoint (~23.4 GB) is used so weights fit the ZeroGPU budget; the
JoyAI-Echo pipeline's own 3-phase module hot-swap keeps peak VRAM manageable.
Gemma-3-12B-it is the (gated) text encoder — requires an `HF_TOKEN` secret with
access to `google/gemma-3-12b-it`.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import json
import shutil
import subprocess
import sys
import time
from pathlib import Path

import gradio as gr
import spaces
from huggingface_hub import snapshot_download, hf_hub_download

# --------------------------------------------------------------------------- #
# Constants / paths
# --------------------------------------------------------------------------- #
ROOT = Path("/tmp/joyai_echo")
REPO = ROOT / "JoyAI-Echo"
VENV = REPO / ".venv"
VENV_PY = VENV / "bin" / "python"
CKPT_REPO = "joeygambino/joyai-echo-ltx23-echoVid-ltxAud-surgical"
CKPT_FILE = "ltx23_echoVid-ltxAud_surgical_fp8.safetensors"
GEMMA_REPO = "google/gemma-3-12b-it"
JOYAI_GIT = "https://github.com/jd-opensource/JoyAI-Echo.git"

TOKEN = os.environ.get("HF_TOKEN")

# One recurring-character example script (three coherent shots).
EXAMPLE_SCRIPT = 'ID_A is Maya Chen, a woman in her early thirties with shoulder-length black hair, warm brown eyes and a calm, grounded presence, wearing a charcoal wool coat over a cream sweater. ID_A\'s voice is a clear, warm young female voice in a casual American accent. At normal speed, ID_A stands still on a quiet city rooftop at dusk, facing the camera, and says, "I used to think the view got smaller the higher you climbed. It doesn\'t - you just start seeing how everything connects." Realistic cinematic footage, soft blue-hour light, gentle handheld framing. A locked medium close-up keeps her face large in frame and her mouth clearly readable while she speaks. Background: a hazy city skyline, string lights on the railing, warm windows below. Quiet diegetic sound only: soft wind and a distant traffic hum under her voice.\n---\nID_A is Maya Chen, a woman in her early thirties with shoulder-length black hair, warm brown eyes and a calm, grounded presence, wearing a charcoal wool coat over a cream sweater. ID_A\'s voice is a clear, warm young female voice in a casual American accent. At normal speed, ID_A stands still in a narrow indoor gallery in front of a wall of framed photographs, facing the camera, and says, "Every one of these was somebody\'s whole evening. You forget that when you\'re only counting the numbers." Realistic cinematic footage, warm gallery lighting. A locked medium close-up keeps her face large in frame and her mouth clearly readable while she speaks. Background: a softly lit gallery wall of framed photographs, polished wooden floor. Quiet diegetic sound only: faint room tone and the muffled sounds of the building under her voice.'


def _run(cmd, cwd=None, env=None):
    """Run a shell command, streaming combined output; raise with the log on failure."""
    proc = subprocess.run(
        cmd,
        cwd=str(cwd) if cwd else None,
        shell=True,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        env=env,
        check=False,
    )
    print(proc.stdout, flush=True)
    if proc.returncode != 0:
        raise RuntimeError(proc.stdout[-4000:])
    return proc.stdout


# --------------------------------------------------------------------------- #
# One-time setup at module scope (runs in the main process at startup).
#   - clone the JoyAI-Echo inference repo (pure python; deps are baked into the
#     Space image via requirements.txt so no runtime venv build is needed)
#   - download the merged fp8 checkpoint + gemma-3-12b-it text encoder
#   - point the repo's config at our checkpoint
# --------------------------------------------------------------------------- #
def setup():
    ROOT.mkdir(parents=True, exist_ok=True)

    if not (REPO / "inference.py").exists():
        _run(f"git clone --depth 1 {JOYAI_GIT} JoyAI-Echo", ROOT)

    # Build the isolated inference venv (JoyAI-Echo's own pinned stack). This is
    # separate from the main Gradio process so transformers 4.57.6 (needs
    # huggingface-hub <1.0) doesn't clash with the runtime's gradio/hub 1.x.
    if not VENV_PY.exists():
        print("[setup] building inference venv (uv)...", flush=True)
        _run(f"{sys.executable} -m uv venv --python 3.11 {VENV}", REPO)
        _run(
            f"{sys.executable} -m uv pip install --python {VENV_PY} "
            f"--extra-index-url https://download.pytorch.org/whl/cu128 "
            f"-r requirements.txt pyyaml huggingface_hub soundfile",
            REPO,
        )
        print("[setup] venv ready", flush=True)

    # UPSTREAM FPS HOTFIX (2026-07-23): the release code hardcodes
    # VIDEO_FPS = 24.0 in the video RoPE position math (ltx_wrapper.py) while
    # this Space renders at 25 fps. That 4% clock skew accumulates ~40ms of
    # mouth-ahead-of-audio drift per second - visibly off ~10s into a shot.
    # Patch every cloned/installed copy to the render fps.
    import pathlib
    for root in (REPO, VENV):
        for f in pathlib.Path(root).rglob("ltx_wrapper.py"):
            try:
                txt = f.read_text(encoding="utf-8")
                if "VIDEO_FPS = 24.0" in txt:
                    f.write_text(txt.replace("VIDEO_FPS = 24.0", "VIDEO_FPS = 25.0"),
                                 encoding="utf-8")
                    print(f"[setup] fps hotfix applied: {f}", flush=True)
            except Exception as e:
                print(f"[setup] fps hotfix skipped for {f}: {e}", flush=True)

    # Ensure the audio-write backend exists even on a pre-existing venv.
    # torchaudio 2.8 has no built-in WAV encoder; it dispatches to soundfile
    # (bundles libsndfile) for `.wav` output. Without it the run crashes at
    # the very end when muxing the combined audio track.
    _run(f"{sys.executable} -m uv pip install --python {VENV_PY} soundfile", REPO)

    checkpoints = REPO / "checkpoints"
    prompts = REPO / "prompts"
    checkpoints.mkdir(exist_ok=True)
    prompts.mkdir(exist_ok=True)

    if not TOKEN:
        raise RuntimeError(
            "Missing HF_TOKEN Space secret. Add a Hugging Face token with "
            "access to the gated google/gemma-3-12b-it model."
        )

    # Merged surgical fp8 checkpoint -> checkpoints/model.safetensors
    ckpt_target = checkpoints / "model.safetensors"
    if not ckpt_target.exists():
        print(f"[setup] downloading {CKPT_REPO}/{CKPT_FILE} ...", flush=True)
        local = hf_hub_download(CKPT_REPO, CKPT_FILE, token=TOKEN)
        try:
            os.symlink(local, ckpt_target)
        except OSError:
            shutil.copy2(local, ckpt_target)

    # Gemma-3-12B-it text encoder
    gemma_dir = checkpoints / "gemma-3-12b"
    if not (gemma_dir / "config.json").exists():
        print(f"[setup] downloading {GEMMA_REPO} ...", flush=True)
        snapshot_download(
            GEMMA_REPO,
            local_dir=gemma_dir,
            token=TOKEN,
            ignore_patterns=["*.gguf", "original/*"],
        )

    # Point the pipeline config at our checkpoint + gemma (inference.py has no
    # --checkpoint / --gemma-path CLI flags, so we edit the YAML in place).
    cfg_path = REPO / "configs" / "inference.yaml"
    import yaml as _yaml
    cfg = _yaml.safe_load(cfg_path.read_text()) or {}
    cfg.setdefault("paths", {})
    cfg["paths"]["checkpoint"] = str(ckpt_target)
    cfg["paths"]["gemma_path"] = str(gemma_dir)
    cfg_path.write_text(_yaml.safe_dump(cfg, sort_keys=False), encoding="utf-8")
    print(f"[setup] config checkpoint -> {ckpt_target}", flush=True)

    print("[setup] complete", flush=True)


try:
    setup()
    SETUP_OK = True
    SETUP_ERR = ""
except Exception as exc:  # surface at request time instead of crashing boot
    SETUP_OK = False
    SETUP_ERR = str(exc)
    print(f"[setup] FAILED: {SETUP_ERR}", flush=True)


def _seconds_for(num_shots: int, frames: int) -> int:
    """Duration estimate: cold GPU worker + model load, then per-shot work.

    Measured on zero-a10g: a full cold call (worker spawn + 25GB fp8 + 24GB
    Gemma load + Stage1 encode + Stage2 gen + decode + mux) for 1 shot / 49
    frames ran ~80s wall. We budget a cold-load base plus per-shot cost that
    scales with frame count, capped at the ZeroGPU ceiling.
    """
    per_shot = 30 + frames * 0.6
    return int(min(1500, 130 + num_shots * per_shot))


def _estimate(script, seed=0, frames=49, height=480, width=832,
              memory_max_size=7, *a, **k):
    shots = max(1, len([s for s in str(script).split("---") if s.strip()]))
    f = max(9, 1 + 8 * round((int(frames) - 1) / 8))
    return _seconds_for(shots, f)


@spaces.GPU(duration=_estimate)
def generate(
    script: str,
    seed: int = 12345,
    frames: int = 49,
    height: int = 480,
    width: int = 832,
    memory_max_size: int = 7,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a multi-shot narrated video (with synchronized audio).

    Args:
        script: The story. Separate shots with a line containing only `---`.
            Each shot is one self-contained prompt describing the character
            (use a stable label like `ID_A`), the spoken line, style, camera,
            background and sound. Reusing the same character description across
            shots lets the paired memory bank hold identity and voice.
        seed: Base RNG seed (each shot uses seed + shot index).
        frames: Frames per shot (25 fps). Lower is faster / less VRAM.
        height: Video height in pixels (multiple of 32).
        width: Video width in pixels (multiple of 32).
        memory_max_size: Cross-shot memory bank size (paired audio-video slots).

    Returns:
        (video_path, status_text) — an MP4 with muxed audio, plus a status log.
    """
    if not SETUP_OK:
        return None, f"Setup failed at startup:\n{SETUP_ERR}"

    started = time.time()

    shots = [s.strip() for s in str(script).split("---") if s.strip()]
    if not shots:
        return None, "Please enter at least one shot."

    # The pipeline requires num_frames == 1 + 8*k, and H/W multiples of 32.
    frames = int(frames)
    frames = max(9, 1 + 8 * round((frames - 1) / 8))
    height = max(256, int(round(int(height) / 32) * 32))
    width = max(256, int(round(int(width) / 32) * 32))

    # Write the prompt JSON the pipeline consumes.
    prompt_path = REPO / "prompts" / "space_prompt.json"
    prompt_path.write_text(json.dumps({"prompts": shots}, indent=2), encoding="utf-8")

    # Isolated output dir per request (no fixed shared path).
    out_root = REPO / "inference_result" / f"run_{int(time.time()*1000)}"

    env = dict(os.environ)
    env["HF_TOKEN"] = TOKEN or ""

    cmd = (
        f"{VENV_PY} inference.py "
        f"--prompts-dir prompts --prompts-glob space_prompt.json "
        f"--output-root {out_root} "
        f"--seed {int(seed)} --num-frames {int(frames)} "
        f"--video-height {int(height)} --video-width {int(width)} "
        f"--memory-max-size {int(memory_max_size)}"
    )
    print(f"[generate] {len(shots)} shot(s): {cmd}", flush=True)
    log = _run(cmd, REPO, env=env)

    # Prefer the combined (all shots concatenated) MP4; fall back to any MP4.
    combined = sorted(out_root.glob("**/combined_shots.mp4"),
                      key=lambda p: p.stat().st_mtime, reverse=True)
    if not combined:
        combined = sorted(out_root.glob("**/*.mp4"),
                          key=lambda p: p.stat().st_mtime, reverse=True)
    if not combined:
        raise RuntimeError("Inference finished but produced no MP4.\n" + log[-2000:])

    elapsed = time.time() - started
    status = (
        f"Done in {elapsed:.1f}s — {len(shots)} shot(s), "
        f"{int(frames)} frames each @ {int(width)}x{int(height)}."
    )
    return str(combined[0]), status


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🎬 JoyAI-Echo × LTX-2.3 — Multi-Shot Narrated Video\n"
            "A **surgical merge** of JoyAI-Echo (cross-shot character memory) and "
            "LTX-2.3-distilled (natural voice + lip-sync) for joint **audio-video** "
            "generation. Write a story as one or more shots (separate shots with a "
            "line containing only `---`); the paired memory bank keeps the same "
            "character's face and voice consistent across shots.\n\n"
            "Model: [`joeygambino/joyai-echo-ltx23-echoVid-ltxAud-surgical`]"
            "(https://huggingface.co/joeygambino/joyai-echo-ltx23-echoVid-ltxAud-surgical) "
            "· fp8 checkpoint · text encoder: Gemma-3-12B · **non-commercial** "
            "(LTX-2 Community License). Generated content is machine-generated."
        )

        if not SETUP_OK:
            gr.Markdown(f"⚠️ **Startup setup failed:** `{SETUP_ERR}`")

        with gr.Row():
            with gr.Column(scale=3):
                script = gr.Textbox(
                    label="Story script (separate shots with a line of `---`)",
                    lines=14,
                    value=EXAMPLE_SCRIPT,
                    placeholder=(
                        "ID_A is <character description>. ID_A's voice is <voice>. "
                        "At normal speed, ID_A <action>, then says, \"<line>\". "
                        "<style> <camera> <background> <sound>.\n"
                        "---\n"
                        "ID_A is <same character>. ... next shot ..."
                    ),
                )
                run_btn = gr.Button("Generate video", variant="primary")
            with gr.Column(scale=2):
                video = gr.Video(label="Result (video + audio)", autoplay=True)
                status = gr.Textbox(label="Status", lines=3)

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                seed = gr.Number(value=12345, precision=0, label="Seed")
                frames = gr.Slider(9, 121, value=49, step=8,
                                   label="Frames per shot (must be 1 + 8k)")
            with gr.Row():
                height = gr.Slider(256, 736, value=480, step=32, label="Height")
                width = gr.Slider(256, 1280, value=832, step=32, label="Width")
            memory_max_size = gr.Slider(
                1, 7, value=7, step=1,
                label="Memory bank size (cross-shot identity slots)")
            gr.Markdown(
                "Tips from the model card: state the accent affirmatively "
                "(e.g. *\"in a casual American accent\"*) or voices drift British; "
                "keep the same character description across shots for identity to "
                "hold. Larger frame counts and resolutions are slower and use more "
                "VRAM."
            )

        gr.Examples(
            examples=[
                [EXAMPLE_SCRIPT],
                [
                    "ID_A is Leo, a man in his late twenties with short curly hair, "
                    "a trimmed beard and an easy grin, wearing a mustard hoodie. "
                    "ID_A's voice is a bright, upbeat young male voice in a casual "
                    "American accent. At normal speed, ID_A leans toward the camera "
                    "in a small sunlit kitchen and says, \"Okay, batch one is in — "
                    "if this works, breakfast is officially solved forever.\" "
                    "Realistic warm handheld footage, cheerful morning light. A "
                    "close medium shot keeps his face and mouth readable. "
                    "Background: a bright kitchen counter, a steaming pan, plants on "
                    "the windowsill. Sizzling pan, light clatter, an upbeat soft "
                    "music bed under the dialogue."
                ],
                [
                    "ID_A is Dr. Amara, a woman in her forties with locs tied back "
                    "and steady dark eyes, wearing a white lab coat. ID_A's voice "
                    "is a calm, authoritative female voice in a neutral American "
                    "accent. At normal speed, ID_A stands in a clean lab and looks "
                    "into the camera, then says, \"The result surprised us too — but "
                    "the data is clear, and clear data is how everything changes.\" "
                    "Realistic cinematic documentary look, cool even lighting. A "
                    "stable medium shot keeps her face and mouth readable. "
                    "Background: a modern laboratory with glass equipment and soft "
                    "monitor glow. Faint equipment hum, a quiet neutral music bed "
                    "under the dialogue."
                ],
            ],
            inputs=[script],
            outputs=[video, status],
            fn=generate,
            cache_examples=False,
            run_on_click=True,
        )

        run_btn.click(
            generate,
            inputs=[script, seed, frames, height, width, memory_max_size],
            outputs=[video, status],
            api_name="generate",
        )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(mcp_server=True)